diff --git a/cmd/go-judge/grpc_executor/grpc.go b/cmd/go-judge/grpc_executor/grpc.go index 25e2169..135c0e4 100644 --- a/cmd/go-judge/grpc_executor/grpc.go +++ b/cmd/go-judge/grpc_executor/grpc.go @@ -211,6 +211,7 @@ func convertPBCmd(c *pb.Request_CmdType, srcPrefix []string) (cm worker.Cmd, err CopyOutCached: convertCopyOut(c.GetCopyOutCached()), CopyOutMax: c.GetCopyOutMax(), CopyOutDir: c.GetCopyOutDir(), + CopyOutTruncate: c.GetCopyOutTruncate(), Symlinks: c.GetSymlinks(), } for _, f := range c.GetFiles() { diff --git a/cmd/go-judge/main.go b/cmd/go-judge/main.go index e28e1cb..427a433 100644 --- a/cmd/go-judge/main.go +++ b/cmd/go-judge/main.go @@ -571,6 +571,7 @@ func generateHandleVersion(_ *config.Config, _ map[string]any) func(*gin.Context "addressSpaceLimit": true, "stream": true, "procPeak": true, + "copyOutTruncate": true, }) } } @@ -584,6 +585,7 @@ func generateHandleConfig(conf *config.Config, builderParam map[string]any) func "addressSpaceLimit": true, "stream": true, "procPeak": true, + "copyOutTruncate": true, "fileStorePath": conf.Dir, "runnerConfig": builderParam, }) diff --git a/cmd/go-judge/model/model.go b/cmd/go-judge/model/model.go index 74b8446..7ef450f 100644 --- a/cmd/go-judge/model/model.go +++ b/cmd/go-judge/model/model.go @@ -48,10 +48,11 @@ type Cmd struct { CopyIn map[string]CmdFile `json:"copyIn"` - CopyOut []string `json:"copyOut"` - CopyOutCached []string `json:"copyOutCached"` - CopyOutMax uint64 `json:"copyOutMax"` - CopyOutDir string `json:"copyOutDir"` + CopyOut []string `json:"copyOut"` + CopyOutCached []string `json:"copyOutCached"` + CopyOutMax uint64 `json:"copyOutMax"` + CopyOutDir string `json:"copyOutDir"` + CopyOutTruncate bool `json:"copyOutTruncate"` TTY bool `json:"tty,omitempty"` StrictMemoryLimit bool `json:"strictMemoryLimit"` @@ -317,6 +318,7 @@ func convertCmd(c Cmd, srcPrefix []string) (worker.Cmd, error) { CopyOutCached: convertCopyOut(c.CopyOutCached), CopyOutMax: c.CopyOutMax, CopyOutDir: c.CopyOutDir, + CopyOutTruncate: c.CopyOutTruncate, } for _, f := range c.Files { cf, err := convertCmdFile(f, srcPrefix) diff --git a/envexec/cmd.go b/envexec/cmd.go index ba0d768..26c8a01 100644 --- a/envexec/cmd.go +++ b/envexec/cmd.go @@ -50,8 +50,9 @@ type Cmd struct { Waiter func(context.Context, Process) bool // file names to copyout after exec - CopyOut []CmdCopyOutFile - CopyOutMax Size // file size limit + CopyOut []CmdCopyOutFile + CopyOutMax Size // file size limit + CopyOutTruncate bool // CopyOutDir specifies a dir to dump all /w content CopyOutDir string diff --git a/envexec/file_collect.go b/envexec/file_collect.go index 35318db..f6991c6 100644 --- a/envexec/file_collect.go +++ b/envexec/file_collect.go @@ -98,7 +98,10 @@ func copyOutFile( s := stat.Size() if c.CopyOutMax > 0 && s > int64(c.CopyOutMax) { t = ErrCopyOutSizeExceeded - return fmt.Errorf("copyout: %q size (%d) exceeds limit (%d)", n.Name, s, c.CopyOutMax) + if !c.CopyOutTruncate { + return fmt.Errorf("copyout: %q size (%d) exceeds limit (%d)", n.Name, s, c.CopyOutMax) + } + s = int64(c.CopyOutMax) + 1 } // create store file buf, err := newStoreFile() @@ -108,13 +111,17 @@ func copyOutFile( } // Ensure not copy over file size - _, err = buf.ReadFrom(io.LimitReader(cf, s)) + s, err = buf.ReadFrom(io.LimitReader(cf, s)) if err != nil { t = ErrCopyOutCopyContent buf.Close() return fmt.Errorf("copyout: failed to copy content for %q: %w", n.Name, err) } put(buf, n.Name) + if s > int64(c.CopyOutMax) { + t = ErrCopyOutSizeExceeded + return fmt.Errorf("copyout: %q size (%d) exceeds limit (%d)", n.Name, s, c.CopyOutMax) + } return nil } diff --git a/worker/model.go b/worker/model.go index deba650..4240b20 100644 --- a/worker/model.go +++ b/worker/model.go @@ -34,10 +34,11 @@ type Cmd struct { CopyIn map[string]CmdFile Symlinks map[string]string - CopyOut []CmdCopyOutFile - CopyOutCached []CmdCopyOutFile - CopyOutMax uint64 - CopyOutDir string + CopyOut []CmdCopyOutFile + CopyOutCached []CmdCopyOutFile + CopyOutMax uint64 + CopyOutDir string + CopyOutTruncate bool TTY bool DataSegmentLimit bool diff --git a/worker/worker.go b/worker/worker.go index 4f8ea13..a3676b0 100644 --- a/worker/worker.go +++ b/worker/worker.go @@ -391,6 +391,7 @@ func (w *worker) prepareCmd(rc Cmd, pipeFileName map[string]bool) (*envexec.Cmd, CopyOut: copyOut, CopyOutDir: copyOutDir, CopyOutMax: copyOutMax, + CopyOutTruncate: rc.CopyOutTruncate, Waiter: wait.Wait, }, nil }