mirror of
https://github.com/criyle/go-judge.git
synced 2025-11-04 14:50:02 +08:00
176 lines
4.8 KiB
Go
176 lines
4.8 KiB
Go
package grpcexecutor
|
|
|
|
import (
|
|
"errors"
|
|
|
|
"github.com/criyle/go-judge/cmd/go-judge/model"
|
|
"github.com/criyle/go-judge/cmd/go-judge/stream"
|
|
"github.com/criyle/go-judge/pb"
|
|
"google.golang.org/grpc/codes"
|
|
"google.golang.org/grpc/status"
|
|
)
|
|
|
|
var _ stream.Stream = &streamWrapper{}
|
|
|
|
type streamWrapper struct {
|
|
es pb.Executor_ExecStreamServer
|
|
}
|
|
|
|
func (sw *streamWrapper) Send(r stream.Response) error {
|
|
res := &pb.StreamResponse{}
|
|
switch {
|
|
case r.Response != nil:
|
|
resp, err := convertPBResponse(*r.Response)
|
|
if err != nil {
|
|
return status.Errorf(codes.Aborted, "response: %v", err)
|
|
}
|
|
res.Response = &pb.StreamResponse_ExecResponse{ExecResponse: resp}
|
|
case r.Output != nil:
|
|
res.Response = &pb.StreamResponse_ExecOutput{ExecOutput: &pb.StreamResponse_Output{
|
|
Index: uint32(r.Output.Index),
|
|
Fd: uint32(r.Output.Fd),
|
|
Content: r.Output.Content,
|
|
}}
|
|
}
|
|
return sw.es.Send(res)
|
|
}
|
|
|
|
func (sw *streamWrapper) Recv() (*stream.Request, error) {
|
|
req, err := sw.es.Recv()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
switch i := req.Request.(type) {
|
|
case *pb.StreamRequest_ExecRequest:
|
|
return &stream.Request{Request: convertPBStreamRequest(i.ExecRequest)}, nil
|
|
case *pb.StreamRequest_ExecInput:
|
|
return &stream.Request{Input: &stream.InputRequest{
|
|
Index: int(i.ExecInput.Index),
|
|
Fd: int(i.ExecInput.Fd),
|
|
Content: i.ExecInput.Content,
|
|
}}, nil
|
|
case *pb.StreamRequest_ExecResize:
|
|
return &stream.Request{Resize: &stream.ResizeRequest{
|
|
Index: int(i.ExecResize.Index),
|
|
Fd: int(i.ExecResize.Fd),
|
|
Rows: int(i.ExecResize.Rows),
|
|
Cols: int(i.ExecResize.Cols),
|
|
X: int(i.ExecResize.X),
|
|
Y: int(i.ExecResize.Y),
|
|
}}, nil
|
|
case *pb.StreamRequest_ExecCancel:
|
|
return &stream.Request{Cancel: &struct{}{}}, nil
|
|
}
|
|
return nil, errors.ErrUnsupported
|
|
}
|
|
|
|
func convertPBStreamRequest(req *pb.Request) *model.Request {
|
|
ret := &model.Request{
|
|
RequestID: req.RequestID,
|
|
}
|
|
for _, cmd := range req.Cmd {
|
|
ret.Cmd = append(ret.Cmd, model.Cmd{
|
|
Args: cmd.Args,
|
|
Env: cmd.Env,
|
|
TTY: cmd.Tty,
|
|
Files: convertPBStreamFiles(cmd.Files),
|
|
CPULimit: cmd.CpuTimeLimit,
|
|
ClockLimit: cmd.ClockTimeLimit,
|
|
MemoryLimit: cmd.MemoryLimit,
|
|
StackLimit: cmd.StackLimit,
|
|
ProcLimit: cmd.ProcLimit,
|
|
CPURateLimit: cmd.CpuRateLimit,
|
|
CPUSetLimit: cmd.CpuSetLimit,
|
|
DataSegmentLimit: cmd.DataSegmentLimit,
|
|
AddressSpaceLimit: cmd.AddressSpaceLimit,
|
|
CopyIn: convertPBStreamCopyIn(cmd),
|
|
CopyOut: convertStreamCopyOut(cmd.CopyOut),
|
|
CopyOutCached: convertStreamCopyOut(cmd.CopyOutCached),
|
|
CopyOutMax: cmd.CopyOutMax,
|
|
CopyOutDir: cmd.CopyOutDir,
|
|
})
|
|
}
|
|
for _, p := range req.PipeMapping {
|
|
ret.PipeMapping = append(ret.PipeMapping, model.PipeMap{
|
|
In: convertPBStreamPipeIndex(p.In),
|
|
Out: convertPBStreamPipeIndex(p.Out),
|
|
Max: int64(p.Max),
|
|
Name: p.Name,
|
|
Proxy: p.Proxy,
|
|
})
|
|
}
|
|
return ret
|
|
}
|
|
|
|
func convertPBStreamPipeIndex(pi *pb.Request_PipeMap_PipeIndex) model.PipeIndex {
|
|
return model.PipeIndex{Index: int(pi.Index), Fd: int(pi.Fd)}
|
|
}
|
|
|
|
func convertPBStreamFiles(files []*pb.Request_File) []*model.CmdFile {
|
|
var rt []*model.CmdFile
|
|
for _, f := range files {
|
|
if f == nil {
|
|
rt = append(rt, nil)
|
|
} else {
|
|
m := convertPBStreamFile(f)
|
|
rt = append(rt, &m)
|
|
}
|
|
}
|
|
return rt
|
|
}
|
|
|
|
func convertPBStreamCopyIn(cmd *pb.Request_CmdType) map[string]model.CmdFile {
|
|
rt := make(map[string]model.CmdFile, len(cmd.CopyIn)+len(cmd.Symlinks))
|
|
for k, i := range cmd.CopyIn {
|
|
if i.File == nil {
|
|
continue
|
|
}
|
|
rt[k] = convertPBStreamFile(i)
|
|
}
|
|
for k, v := range cmd.Symlinks {
|
|
rt[k] = model.CmdFile{Symlink: &v}
|
|
}
|
|
return rt
|
|
}
|
|
|
|
func convertPBStreamFile(i *pb.Request_File) model.CmdFile {
|
|
switch c := i.File.(type) {
|
|
case *pb.Request_File_Local:
|
|
return model.CmdFile{Src: &c.Local.Src}
|
|
case *pb.Request_File_Memory:
|
|
s := byteArrayToString(c.Memory.Content)
|
|
return model.CmdFile{Content: &s}
|
|
case *pb.Request_File_Cached:
|
|
return model.CmdFile{FileID: &c.Cached.FileID}
|
|
case *pb.Request_File_Pipe:
|
|
return model.CmdFile{Name: &c.Pipe.Name, Max: &c.Pipe.Max, Pipe: c.Pipe.Pipe}
|
|
case *pb.Request_File_StreamIn:
|
|
return model.CmdFile{StreamIn: true}
|
|
case *pb.Request_File_StreamOut:
|
|
return model.CmdFile{StreamOut: true}
|
|
}
|
|
return model.CmdFile{}
|
|
}
|
|
|
|
func convertStreamCopyOut(copyOut []*pb.Request_CmdCopyOutFile) []string {
|
|
rt := make([]string, 0, len(copyOut))
|
|
for _, n := range copyOut {
|
|
name := n.Name
|
|
if n.Optional {
|
|
name += "?"
|
|
}
|
|
rt = append(rt, name)
|
|
}
|
|
return rt
|
|
}
|
|
|
|
func (e *execServer) ExecStream(es pb.Executor_ExecStreamServer) error {
|
|
w := &streamWrapper{
|
|
es: es,
|
|
}
|
|
if err := stream.Start(es.Context(), w, e.worker, e.srcPrefix, e.logger); err != nil {
|
|
return status.Error(codes.Internal, err.Error())
|
|
}
|
|
return nil
|
|
}
|