package main import ( "context" "errors" "log" "os" "strings" "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" "google.golang.org/grpc/credentials" "google.golang.org/grpc/credentials/insecure" "google.golang.org/protobuf/types/known/emptypb" ) var _ Stream = &grpcWrapper{} type grpcWrapper struct { sc pb.Executor_ExecStreamClient } 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.NewClient(srvAddr, opts...) if err != nil { log.Fatalln("client", err) } client := pb.NewExecutorClient(conn) sc, err := client.ExecStream(context.TODO()) if err != nil { log.Fatalln("exec_stream", err) } log.Println("start", args) return &grpcWrapper{sc: sc} } func (w *grpcWrapper) Send(req *stream.Request) error { switch { case req.Request != nil: w.sc.Send(convertPBRequest(req.Request)) case req.Input != nil: w.sc.Send(pb.StreamRequest_builder{ExecInput: pb.StreamRequest_Input_builder{ Index: uint32(req.Input.Index), Fd: uint32(req.Input.Fd), Content: req.Input.Content, }.Build()}.Build()) case req.Resize != nil: w.sc.Send(pb.StreamRequest_builder{ExecResize: pb.StreamRequest_Resize_builder{ Index: uint32(req.Resize.Index), Fd: uint32(req.Resize.Fd), Rows: uint32(req.Resize.Rows), Cols: uint32(req.Resize.Cols), X: uint32(req.Resize.X), Y: uint32(req.Resize.Y), }.Build()}.Build()) case req.Cancel != nil: w.sc.Send(pb.StreamRequest_builder{ExecCancel: &emptypb.Empty{}}.Build()) default: return errors.New("send: unknown operation") } return nil } func (w *grpcWrapper) Recv() (*stream.Response, error) { resp, err := w.sc.Recv() if err != nil { return nil, err } switch resp.WhichResponse() { case pb.StreamResponse_ExecOutput_case: return &stream.Response{Output: &stream.OutputResponse{ Index: int(resp.GetExecOutput().GetIndex()), Fd: int(resp.GetExecOutput().GetFd()), Content: resp.GetExecOutput().GetContent(), }}, nil case pb.StreamResponse_ExecResponse_case: return &stream.Response{Response: &model.Response{ RequestID: resp.GetExecResponse().GetRequestID(), Results: convertPBResult(resp.GetExecResponse().GetResults()), ErrorMsg: resp.GetExecResponse().GetError(), }}, nil } return nil, errors.New("recv: invalid response") } func convertPBResult(res []*pb.Response_Result) []model.Result { var ret []model.Result for _, r := range res { ret = append(ret, model.Result{ Status: model.Status(r.GetStatus()), ExitStatus: int(r.GetExitStatus()), Error: r.GetError(), Time: r.GetTime(), RunTime: r.GetRunTime(), Memory: r.GetMemory(), Files: convertFiles(r.GetFiles()), Buffs: r.GetFiles(), FileIDs: r.GetFileIDs(), FileError: convertPBFileError(r.GetFileError()), }) } return ret } func convertFiles(buf map[string][]byte) map[string]string { ret := make(map[string]string, len(buf)) for k, v := range buf { ret[k] = byteArrayToString(v) } return ret } func convertPBRequest(req *model.Request) *pb.StreamRequest { ret := pb.StreamRequest_builder{ ExecRequest: pb.Request_builder{ RequestID: req.RequestID, Cmd: convertPBCmd(req.Cmd), PipeMapping: convertPBPipeMapping(req.PipeMapping), }.Build(), }.Build() return ret } func convertPBFileError(fe []*pb.Response_FileError) []model.FileError { var ret []model.FileError for _, v := range fe { ret = append(ret, model.FileError{ Name: v.GetName(), Type: model.FileErrorType(v.GetType()), Message: v.GetMessage(), }) } return ret } func convertPBCmd(cmd []model.Cmd) []*pb.Request_CmdType { var ret []*pb.Request_CmdType for _, c := range cmd { ret = append(ret, pb.Request_CmdType_builder{ Args: c.Args, Env: c.Env, Tty: c.TTY, Files: convertPBFiles(c.Files), CpuTimeLimit: c.CPULimit, ClockTimeLimit: c.ClockLimit, MemoryLimit: c.MemoryLimit, StackLimit: c.StackLimit, ProcLimit: c.ProcLimit, CpuRateLimit: c.CPURateLimit, CpuSetLimit: c.CPUSetLimit, DataSegmentLimit: c.DataSegmentLimit, AddressSpaceLimit: c.AddressSpaceLimit, CopyIn: convertPBCopyIn(c.CopyIn), CopyOut: convertPBCopyOut(c.CopyOut), CopyOutCached: convertPBCopyOut(c.CopyOutCached), CopyOutMax: c.CopyOutMax, CopyOutDir: c.CopyOutDir, Symlinks: convertSymlink(c.CopyIn), }.Build()) } return ret } func convertPBCopyIn(copyIn map[string]model.CmdFile) map[string]*pb.Request_File { rt := make(map[string]*pb.Request_File, len(copyIn)) for k, i := range copyIn { if i.Symlink != nil { continue } rt[k] = convertPBFile(i) } return rt } func convertPBCopyOut(copyOut []string) []*pb.Request_CmdCopyOutFile { rt := make([]*pb.Request_CmdCopyOutFile, 0, len(copyOut)) for _, n := range copyOut { optional := false if strings.HasSuffix(n, "?") { optional = true n = strings.TrimSuffix(n, "?") } rt = append(rt, pb.Request_CmdCopyOutFile_builder{ Name: n, Optional: optional, }.Build()) } return rt } func convertSymlink(copyIn map[string]model.CmdFile) map[string]string { ret := make(map[string]string) for k, v := range copyIn { if v.Symlink == nil { continue } ret[k] = *v.Symlink } return ret } func convertPBFiles(files []*model.CmdFile) []*pb.Request_File { var ret []*pb.Request_File for _, f := range files { if f == nil { ret = append(ret, nil) } else { ret = append(ret, convertPBFile(*f)) } } return ret } func convertPBFile(i model.CmdFile) *pb.Request_File { switch { case i.Src != nil: return pb.Request_File_builder{Local: pb.Request_LocalFile_builder{Src: *i.Src}.Build()}.Build() case i.Content != nil: s := strToBytes(*i.Content) return pb.Request_File_builder{Memory: pb.Request_MemoryFile_builder{Content: s}.Build()}.Build() case i.FileID != nil: return pb.Request_File_builder{Cached: pb.Request_CachedFile_builder{FileID: *i.FileID}.Build()}.Build() case i.Name != nil && i.Max != nil: return pb.Request_File_builder{Pipe: pb.Request_PipeCollector_builder{Name: *i.Name, Max: *i.Max, Pipe: i.Pipe}.Build()}.Build() case i.StreamIn: return pb.Request_File_builder{StreamIn: &emptypb.Empty{}}.Build() case i.StreamOut: return pb.Request_File_builder{StreamOut: &emptypb.Empty{}}.Build() } return nil } func convertPBPipeMapping(pm []model.PipeMap) []*pb.Request_PipeMap { var ret []*pb.Request_PipeMap for _, p := range pm { ret = append(ret, pb.Request_PipeMap_builder{ In: convertPBPipeIndex(p.In), Out: convertPBPipeIndex(p.Out), Name: p.Name, Proxy: p.Proxy, Max: uint64(p.Max), }.Build()) } return ret } func convertPBPipeIndex(pi model.PipeIndex) *pb.Request_PipeMap_PipeIndex { return pb.Request_PipeMap_PipeIndex_builder{Index: int32(pi.Index), Fd: int32(pi.Fd)}.Build() } type tokenAuth struct { token string } func newTokenAuth(token string) credentials.PerRPCCredentials { return &tokenAuth{token: token} } // Return value is mapped to request headers. func (t *tokenAuth) GetRequestMetadata(ctx context.Context, in ...string) (map[string]string, error) { return map[string]string{ "authorization": "Bearer " + t.token, }, nil } func (*tokenAuth) RequireTransportSecurity() bool { return false }