go-judge/cmd/executorshell/shell.go
2020-06-02 19:43:18 -04:00

213 lines
4.2 KiB
Go

package main
import (
"context"
"flag"
"fmt"
"io"
"log"
"os"
"os/signal"
"syscall"
"time"
"github.com/creack/pty"
"github.com/criyle/go-judge/pb"
"golang.org/x/crypto/ssh/terminal"
"google.golang.org/grpc"
)
var (
srvAddr = flag.String("srvaddr", "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",
}
func main() {
flag.Parse()
args := flag.Args()
if len(args) == 0 {
args = []string{"/bin/bash"}
}
conn, err := grpc.Dial(*srvAddr, grpc.WithInsecure())
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",
Tty: true,
},
},
},
{
File: &pb.Request_File_StreamOut{
StreamOut: &pb.Request_StreamOutput{
Name: "stdout",
},
},
},
{
File: &pb.Request_File_StreamOut{
StreamOut: &pb.Request_StreamOutput{
Name: "stderr",
},
},
},
},
CPULimit: uint64(cpuLimit),
RealCPULimit: uint64(sessionLimit),
MemoryLimit: memoryLimit,
ProcLimit: procLimit,
}},
}
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 := terminal.MakeRaw(int(os.Stdin.Fd()))
if err != nil {
panic(err)
}
defer func() { _ = terminal.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
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 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
ch := make(chan os.Signal, 1)
signal.Notify(ch, syscall.SIGWINCH)
go func() {
for range ch {
winSize, err := pty.GetsizeFull(os.Stdin)
if err != nil {
log.Println("get win size", err)
return
}
sendCh <- &pb.StreamRequest{
Request: &pb.StreamRequest_ExecResize{
ExecResize: &pb.StreamRequest_Resize{
Name: "stdin",
Rows: uint32(winSize.Rows),
Cols: uint32(winSize.Cols),
X: uint32(winSize.X),
Y: uint32(winSize.Y),
},
},
}
}
}()
ch <- syscall.SIGWINCH // Initial resize.
// 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
}
}
}