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 }