mirror of
https://github.com/criyle/go-judge.git
synced 2025-09-26 22:39:12 +08:00
229 lines
4.5 KiB
Go
229 lines
4.5 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
"flag"
|
|
"fmt"
|
|
"io"
|
|
"log"
|
|
"os"
|
|
"os/signal"
|
|
"syscall"
|
|
"time"
|
|
|
|
"github.com/criyle/go-judge/pb"
|
|
"golang.org/x/term"
|
|
"google.golang.org/grpc"
|
|
"google.golang.org/grpc/credentials"
|
|
)
|
|
|
|
var (
|
|
srvAddr = flag.String("srv-addr", "localhost:5051", "GRPC server addr")
|
|
)
|
|
|
|
const (
|
|
cpuLimit = 20 * time.Second
|
|
sessionLimit = 30 * time.Minute
|
|
procLimit = 50
|
|
memoryLimit = 256 << 20 // 256m
|
|
pathEnv = "PATH=/usr/local/bin:/usr/bin:/bin"
|
|
)
|
|
|
|
var env = []string{
|
|
pathEnv,
|
|
"HOME=/tmp",
|
|
"TERM=" + os.Getenv("TERM"),
|
|
}
|
|
|
|
func main() {
|
|
flag.Parse()
|
|
args := flag.Args()
|
|
if len(args) == 0 {
|
|
args = []string{"/bin/bash"}
|
|
}
|
|
|
|
token := os.Getenv("TOKEN")
|
|
opts := []grpc.DialOption{grpc.WithInsecure()}
|
|
if token != "" {
|
|
opts = append(opts, grpc.WithPerRPCCredentials(newTokenAuth(token)))
|
|
}
|
|
conn, err := grpc.Dial(*srvAddr, opts...)
|
|
|
|
if err != nil {
|
|
log.Fatalln("client", err)
|
|
}
|
|
client := pb.NewExecutorClient(conn)
|
|
sc, err := client.ExecStream(context.TODO())
|
|
if err != nil {
|
|
log.Fatalln("ExecStream", err)
|
|
}
|
|
log.Println("Starts", args)
|
|
r, err := run(sc, args)
|
|
log.Println("ExecStream Finished", r, err)
|
|
}
|
|
|
|
func run(sc pb.Executor_ExecStreamClient, args []string) (*pb.Response, error) {
|
|
req := &pb.Request{
|
|
Cmd: []*pb.Request_CmdType{{
|
|
Args: args,
|
|
Env: env,
|
|
Files: []*pb.Request_File{
|
|
{
|
|
File: &pb.Request_File_StreamIn{
|
|
StreamIn: &pb.Request_StreamInput{
|
|
Name: "stdin",
|
|
},
|
|
},
|
|
},
|
|
{
|
|
File: &pb.Request_File_StreamOut{
|
|
StreamOut: &pb.Request_StreamOutput{
|
|
Name: "stdout",
|
|
},
|
|
},
|
|
},
|
|
{
|
|
File: &pb.Request_File_StreamOut{
|
|
StreamOut: &pb.Request_StreamOutput{
|
|
Name: "stderr",
|
|
},
|
|
},
|
|
},
|
|
},
|
|
CpuTimeLimit: uint64(cpuLimit),
|
|
ClockTimeLimit: uint64(sessionLimit),
|
|
MemoryLimit: memoryLimit,
|
|
ProcLimit: procLimit,
|
|
Tty: true,
|
|
}},
|
|
}
|
|
err := sc.Send(&pb.StreamRequest{
|
|
Request: &pb.StreamRequest_ExecRequest{
|
|
ExecRequest: req,
|
|
},
|
|
})
|
|
if err != nil {
|
|
return nil, fmt.Errorf("ExecStream Send request %v", err)
|
|
}
|
|
|
|
// Set stdin in raw mode.
|
|
oldState, err := term.MakeRaw(int(os.Stdin.Fd()))
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
defer func() { _ = term.Restore(int(os.Stdin.Fd()), oldState) }() // Best effort.
|
|
|
|
// pump msg
|
|
sendCh := make(chan *pb.StreamRequest, 64)
|
|
defer close(sendCh)
|
|
go func() {
|
|
for r := range sendCh {
|
|
err := sc.Send(r)
|
|
if err != nil {
|
|
log.Println("input", err)
|
|
return
|
|
}
|
|
}
|
|
}()
|
|
|
|
// pump stdin
|
|
forceQuit := false
|
|
go func() {
|
|
buf := make([]byte, 4096)
|
|
for {
|
|
n, err := os.Stdin.Read(buf)
|
|
if err == io.EOF {
|
|
sendCh <- &pb.StreamRequest{
|
|
Request: &pb.StreamRequest_ExecInput{
|
|
ExecInput: &pb.StreamRequest_Input{
|
|
Name: "stdin",
|
|
Content: []byte("\004"),
|
|
},
|
|
},
|
|
}
|
|
continue
|
|
}
|
|
if n == 1 && buf[0] == 3 {
|
|
if forceQuit {
|
|
sendCh <- &pb.StreamRequest{
|
|
Request: &pb.StreamRequest_ExecCancel{},
|
|
}
|
|
}
|
|
forceQuit = true
|
|
} else {
|
|
forceQuit = false
|
|
}
|
|
|
|
if err != nil {
|
|
log.Println("stdin", err)
|
|
return
|
|
}
|
|
sendCh <- &pb.StreamRequest{
|
|
Request: &pb.StreamRequest_ExecInput{
|
|
ExecInput: &pb.StreamRequest_Input{
|
|
Name: "stdin",
|
|
Content: buf[:n],
|
|
},
|
|
},
|
|
}
|
|
}
|
|
}()
|
|
|
|
sigCh := make(chan os.Signal, 1)
|
|
signal.Notify(sigCh, syscall.SIGINT)
|
|
// pump ^C
|
|
go func() {
|
|
for range sigCh {
|
|
sendCh <- &pb.StreamRequest{
|
|
Request: &pb.StreamRequest_ExecInput{
|
|
ExecInput: &pb.StreamRequest_Input{
|
|
Name: "stdin",
|
|
Content: []byte("\003"),
|
|
},
|
|
},
|
|
}
|
|
}
|
|
}()
|
|
|
|
// pump resize
|
|
handleSizeChange(sendCh)
|
|
|
|
// pump stdout
|
|
for {
|
|
sr, err := sc.Recv()
|
|
if err != nil {
|
|
return nil, fmt.Errorf("ExecStream recv %v", err)
|
|
}
|
|
switch sr := sr.Response.(type) {
|
|
case *pb.StreamResponse_ExecOutput:
|
|
switch sr.ExecOutput.Name {
|
|
case "stdout":
|
|
os.Stdout.Write(sr.ExecOutput.Content)
|
|
case "stderr":
|
|
os.Stderr.Write(sr.ExecOutput.Content)
|
|
}
|
|
case *pb.StreamResponse_ExecResponse:
|
|
return sr.ExecResponse, nil
|
|
}
|
|
}
|
|
}
|
|
|
|
type tokenAuth struct {
|
|
token string
|
|
}
|
|
|
|
func newTokenAuth(token string) credentials.PerRPCCredentials {
|
|
return &tokenAuth{token: token}
|
|
}
|
|
|
|
// Return value is mapped to request headers.
|
|
func (t *tokenAuth) GetRequestMetadata(ctx context.Context, in ...string) (map[string]string, error) {
|
|
return map[string]string{
|
|
"authorization": "Bearer " + t.token,
|
|
}, nil
|
|
}
|
|
|
|
func (*tokenAuth) RequireTransportSecurity() bool {
|
|
return false
|
|
}
|