mirror of
https://github.com/criyle/go-judge.git
synced 2025-09-26 22:39:12 +08:00
105 lines
2.2 KiB
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
|
|
}
|