diff --git a/README.cn.md b/README.cn.md index 0123927..fd2f0a9 100644 --- a/README.cn.md +++ b/README.cn.md @@ -528,6 +528,29 @@ reboot - `-force-gc-target` 默认 `20m`, 堆内存使用超过该值是强制垃圾收集和归还内存 - `-force-gc-interval` 默认 `5s`, 为后台线程检查的频繁程度 +### WebSocket 流接口 + +WebSocket 流接口是用于运行一个程序,同时和它的输入输出进行交互。所有的消息都应该使用 WebSocket 的 binary 格式来发送来避免兼容性问题。 + +```text ++--------+--------+---... +| 类型 | 载荷 ... ++--------|--------+---... +请求: +请求类型 = + 1 - 运行请求 (载荷 = JSON 编码的请求体) + 2 - 设置终端窗口大小 (载荷 = JSON 编码的请求体) + 3 - 输入 (载荷 = 1 字节 (4 位的 命令下标 + 4 位的 文件描述符) + 输入内容) + 4 - 取消 (没有载荷) + +响应: +响应类型 = + 1 - 运行结果 (载荷 = JSON 编码的运行结果) + 2 - 输出 (载荷 = 1 字节 (4 位的 命令下标 + 4 位的 文件描述符) + 输入内容) +``` + +任何的不完整,或者不合法的消息会被认为是错误,并终止运行。 + ### 压力测试 使用 `wrk` 和 `t.lua`: `wrk -s t.lua -c 1 -t 1 -d 30s --latency http://localhost:5050/run`. diff --git a/cmd/go-judge-shell/grpc.go b/cmd/go-judge-shell/grpc.go index f81b5c3..95e580a 100644 --- a/cmd/go-judge-shell/grpc.go +++ b/cmd/go-judge-shell/grpc.go @@ -21,13 +21,13 @@ type grpcWrapper struct { sc pb.Executor_ExecStreamClient } -func newGrpc(args []string, srvAddr *string) Stream { +func newGrpc(args []string, srvAddr string) Stream { token := os.Getenv("TOKEN") opts := []grpc.DialOption{grpc.WithTransportCredentials(insecure.NewCredentials())} if token != "" { opts = append(opts, grpc.WithPerRPCCredentials(newTokenAuth(token))) } - conn, err := grpc.Dial(*srvAddr, opts...) + conn, err := grpc.Dial(srvAddr, opts...) if err != nil { log.Fatalln("client", err) } diff --git a/cmd/go-judge-shell/shell.go b/cmd/go-judge-shell/shell.go index 593cc3d..5628212 100644 --- a/cmd/go-judge-shell/shell.go +++ b/cmd/go-judge-shell/shell.go @@ -16,7 +16,9 @@ import ( ) var ( - srvAddr = flag.String("srv-addr", "localhost:5051", "GRPC server addr") + transport = flag.String("transport", "websocket", "defines transport layer (websocket / grpc)") + wsURL = flag.String("ws-url", "ws://localhost:5050/stream", "HTTP server url") + grpcAddr = flag.String("grpc-addr", "localhost:5051", "GRPC server addr") ) const ( @@ -33,6 +35,7 @@ var env = []string{ "TERM=" + os.Getenv("TERM"), } +// Stream defines the transport layer for stream execution type Stream interface { Send(*stream.Request) error Recv() (*stream.Response, error) @@ -44,8 +47,16 @@ func main() { if len(args) == 0 { args = []string{"/bin/bash"} } - w := newGrpc(args, srvAddr) - r, err := run(w, args) + var s Stream + switch *transport { + case "websocket": + s = newWebsocket(args, *wsURL) + case "grpc": + s = newGrpc(args, *grpcAddr) + default: + log.Fatalln("invalid transport: ", *transport) + } + r, err := run(s, args) log.Printf("finished: %+v %v", r, err) } diff --git a/cmd/go-judge-shell/websocket.go b/cmd/go-judge-shell/websocket.go new file mode 100644 index 0000000..8324b7c --- /dev/null +++ b/cmd/go-judge-shell/websocket.go @@ -0,0 +1,104 @@ +package main + +import ( + "encoding/json" + "fmt" + "io" + "log" + "net/http" + "os" + + "github.com/criyle/go-judge/cmd/go-judge/model" + "github.com/criyle/go-judge/cmd/go-judge/stream" + "github.com/gorilla/websocket" +) + +var _ Stream = &websocketStream{} + +type websocketStream struct { + conn *websocket.Conn +} + +func newWebsocket(args []string, wsURL string) Stream { + header := make(http.Header) + token := os.Getenv("TOKEN") + if token != "" { + header.Add("Authorization", "Bearer "+token) + } + conn, _, err := websocket.DefaultDialer.Dial(wsURL, header) + if err != nil { + log.Fatalln("ws connect: ", err) + } + log.Println("start", args) + return &websocketStream{conn: conn} +} + +// Recv implements Stream. +func (s *websocketStream) Recv() (*stream.Response, error) { + _, r, err := s.conn.ReadMessage() + if err != nil { + return nil, err + } + if len(r) == 0 { + return nil, io.ErrUnexpectedEOF + } + resp := new(stream.Response) + switch r[0] { + case 1: + resp.Response = new(model.Response) + if err := json.Unmarshal(r[1:], resp.Response); err != nil { + return nil, err + } + case 2: + if len(r) < 2 { + return nil, io.ErrUnexpectedEOF + } + resp.Output = new(stream.OutputResponse) + resp.Output.Index = int(r[1]>>4) & 0xf + resp.Output.Fd = int(r[1]) & 0xf + resp.Output.Content = r[2:] + default: + return nil, fmt.Errorf("invalid type code: %d", r[0]) + } + return resp, nil +} + +// Send implements Stream. +func (s *websocketStream) Send(req *stream.Request) error { + w, err := s.conn.NextWriter(websocket.BinaryMessage) + if err != nil { + return err + } + defer w.Close() + + switch { + case req.Request != nil: + if _, err := w.Write([]byte{1}); err != nil { + return err + } + if err := json.NewEncoder(w).Encode(req.Request); err != nil { + return err + } + case req.Resize != nil: + if _, err := w.Write([]byte{2}); err != nil { + return err + } + if err := json.NewEncoder(w).Encode(req.Resize); err != nil { + return err + } + case req.Input != nil: + if _, err := w.Write([]byte{3, byte(req.Input.Index<<4 | req.Input.Fd)}); err != nil { + return err + } + if _, err := w.Write(req.Input.Content); err != nil { + return err + } + case req.Cancel != nil: + if _, err := w.Write([]byte{4}); err != nil { + return err + } + default: + return fmt.Errorf("invalid request") + } + return nil +} diff --git a/cmd/go-judge/main.go b/cmd/go-judge/main.go index f45691b..e52177e 100644 --- a/cmd/go-judge/main.go +++ b/cmd/go-judge/main.go @@ -510,6 +510,7 @@ func generateHandleVersion(conf *config.Config, builderParam map[string]any) fun "pipeProxy": true, "symlink": true, "addressSpaceLimit": true, + "stream": true, }) } } @@ -521,6 +522,7 @@ func generateHandleConfig(conf *config.Config, builderParam map[string]any) func "pipeProxy": true, "symlink": true, "addressSpaceLimit": true, + "stream": true, "fileStorePath": conf.Dir, "runnerConfig": builderParam, }) diff --git a/cmd/go-judge/model/model.go b/cmd/go-judge/model/model.go index 5357c20..1cee630 100644 --- a/cmd/go-judge/model/model.go +++ b/cmd/go-judge/model/model.go @@ -12,7 +12,10 @@ import ( "github.com/criyle/go-judge/worker" ) +// FileError defines the location, file name and the detailed message for a failed file operation type FileError = envexec.FileError + +// FileErrorType defines the location that file operation fails type FileErrorType = envexec.FileErrorType // CmdFile defines file from multiple source including local / memory / cached or pipe collector diff --git a/cmd/go-judge/stream/stream.go b/cmd/go-judge/stream/stream.go index 4ac4354..832b1a8 100644 --- a/cmd/go-judge/stream/stream.go +++ b/cmd/go-judge/stream/stream.go @@ -40,12 +40,12 @@ type Response struct { // ResizeRequest defines resize operation to the virtual terminal type ResizeRequest struct { - Index int - Fd int - Rows int - Cols int - X int - Y int + Index int `json:"index,omitempty"` + Fd int `json:"fd,omitempty"` + Rows int `json:"rows,omitempty"` + Cols int `json:"cols,omitempty"` + X int `json:"x,omitempty"` + Y int `json:"y,omitempty"` } // InputRequest defines input operation from the remote @@ -97,15 +97,13 @@ func Start(baseCtx context.Context, s Stream, w worker.Worker, srcPrefix []strin defer cancel() // stream in - if len(streamIn) > 0 { - wg.Go(func() error { - if err := streamInput(ctx, s, streamIn, execCancel); err != nil { - cancel() - return err - } - return nil - }) - } + wg.Go(func() error { + if err := streamInput(ctx, s, streamIn, execCancel); err != nil { + cancel() + return err + } + return nil + }) // stream out outCh := make(chan *OutputResponse, len(streamOut)) diff --git a/cmd/go-judge/ws_executor/stream.go b/cmd/go-judge/ws_executor/stream.go new file mode 100644 index 0000000..db1d85a --- /dev/null +++ b/cmd/go-judge/ws_executor/stream.go @@ -0,0 +1,120 @@ +package wsexecutor + +import ( + "context" + "encoding/json" + "fmt" + "io" + "time" + + "github.com/criyle/go-judge/cmd/go-judge/model" + "github.com/criyle/go-judge/cmd/go-judge/stream" + "github.com/gorilla/websocket" +) + +var _ stream.Stream = &streamWrapper{} + +type streamWrapper struct { + ctx context.Context + conn *websocket.Conn + sendCh chan stream.Response +} + +func (w *streamWrapper) sendLoop() { + conn := w.conn + defer conn.Close() + + ticker := time.NewTicker(pingPeriod) + defer ticker.Stop() + for { + select { + case <-w.ctx.Done(): + return + case r := <-w.sendCh: + conn.SetWriteDeadline(time.Now().Add(writeWait)) + switch { + case r.Response != nil: + w, err := conn.NextWriter(websocket.BinaryMessage) + if err != nil { + return + } + if _, err := w.Write([]byte{1}); err != nil { + return + } + if err := json.NewEncoder(w).Encode(r.Response); err != nil { + return + } + if err := w.Close(); err != nil { + return + } + conn.Close() + return + case r.Output != nil: + w, err := conn.NextWriter(websocket.BinaryMessage) + if err != nil { + return + } + if _, err := w.Write([]byte{2, byte(r.Output.Index<<4 | r.Output.Fd)}); err != nil { + return + } + if _, err := w.Write(r.Output.Content); err != nil { + return + } + if err := w.Close(); err != nil { + return + } + } + case <-ticker.C: + conn.SetWriteDeadline(time.Now().Add(writeWait)) + if err := conn.WriteMessage(websocket.PingMessage, nil); err != nil { + return + } + } + } +} + +func (w *streamWrapper) Send(resp stream.Response) error { + select { + case <-w.ctx.Done(): + return w.ctx.Err() + case w.sendCh <- resp: + return nil + } +} + +func (w *streamWrapper) Recv() (*stream.Request, error) { + conn := w.conn + _, buf, err := conn.ReadMessage() + if err != nil { + return nil, err + } + if len(buf) == 0 { + return nil, io.ErrUnexpectedEOF + } + var req stream.Request + switch buf[0] { + case 1: + req.Request = new(model.Request) + if err := json.Unmarshal(buf[1:], req.Request); err != nil { + return nil, err + } + case 2: + req.Resize = new(stream.ResizeRequest) + if err := json.Unmarshal(buf[1:], req.Resize); err != nil { + return nil, err + } + case 3: + if len(buf) < 2 { + return nil, io.ErrUnexpectedEOF + } + req.Input = new(stream.InputRequest) + req.Input.Index = int(buf[1]>>4) & 0xf + req.Input.Fd = int(buf[1]) & 0xf + req.Input.Content = buf[2:] + case 4: + req.Cancel = new(struct{}) + default: + return nil, fmt.Errorf("invalid type code: %d", buf[0]) + } + return &req, nil +} diff --git a/cmd/go-judge/ws_executor/websocket.go b/cmd/go-judge/ws_executor/websocket.go index 8872ddd..9618348 100644 --- a/cmd/go-judge/ws_executor/websocket.go +++ b/cmd/go-judge/ws_executor/websocket.go @@ -8,6 +8,7 @@ import ( "time" "github.com/criyle/go-judge/cmd/go-judge/model" + "github.com/criyle/go-judge/cmd/go-judge/stream" "github.com/criyle/go-judge/worker" "github.com/gin-gonic/gin" "github.com/gorilla/websocket" @@ -55,6 +56,7 @@ type wsRequest struct { func (h *wsHandle) Register(r *gin.Engine) { r.GET("/ws", h.handleWS) + r.GET("/stream", h.handleStream) } func (h *wsHandle) handleWS(c *gin.Context) { @@ -178,6 +180,26 @@ func (h *wsHandle) handleWS(c *gin.Context) { }() } +func (h *wsHandle) handleStream(c *gin.Context) { + conn, err := upgrader.Upgrade(c.Writer, c.Request, nil) + if err != nil { + c.Error(err) + c.AbortWithStatusJSON(http.StatusBadRequest, err.Error()) + return + } + conn.SetReadDeadline(time.Now().Add(pongWait)) + conn.SetPongHandler(func(string) error { + conn.SetReadDeadline(time.Now().Add(pongWait)) + return nil + }) + ctx, cancel := context.WithCancel(context.TODO()) + defer cancel() + + w := &streamWrapper{ctx: ctx, conn: conn, sendCh: make(chan stream.Response)} + go w.sendLoop() + stream.Start(ctx, w, h.worker, h.srcPrefix, h.logger) +} + type contextMap struct { m map[string]context.CancelFunc mu sync.Mutex