go-judge/cmd/go-judge-shell/websocket.go

105 lines
2.2 KiB
Go

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
}