feat(copy_out_truncate): add support to copy out truncate

Author: criyle

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() {

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,
 		})

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)

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

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
 }
 

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

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
 }