feat(copy_out_truncate): add support to copy out truncate

This commit is contained in:
criyle 2025-10-06 22:41:03 +00:00
parent b613b1f747
commit ad661b1ad2
7 changed files with 27 additions and 12 deletions

View File

@ -211,6 +211,7 @@ func convertPBCmd(c *pb.Request_CmdType, srcPrefix []string) (cm worker.Cmd, err
CopyOutCached: convertCopyOut(c.GetCopyOutCached()), CopyOutCached: convertCopyOut(c.GetCopyOutCached()),
CopyOutMax: c.GetCopyOutMax(), CopyOutMax: c.GetCopyOutMax(),
CopyOutDir: c.GetCopyOutDir(), CopyOutDir: c.GetCopyOutDir(),
CopyOutTruncate: c.GetCopyOutTruncate(),
Symlinks: c.GetSymlinks(), Symlinks: c.GetSymlinks(),
} }
for _, f := range c.GetFiles() { for _, f := range c.GetFiles() {

View File

@ -571,6 +571,7 @@ func generateHandleVersion(_ *config.Config, _ map[string]any) func(*gin.Context
"addressSpaceLimit": true, "addressSpaceLimit": true,
"stream": true, "stream": true,
"procPeak": true, "procPeak": true,
"copyOutTruncate": true,
}) })
} }
} }
@ -584,6 +585,7 @@ func generateHandleConfig(conf *config.Config, builderParam map[string]any) func
"addressSpaceLimit": true, "addressSpaceLimit": true,
"stream": true, "stream": true,
"procPeak": true, "procPeak": true,
"copyOutTruncate": true,
"fileStorePath": conf.Dir, "fileStorePath": conf.Dir,
"runnerConfig": builderParam, "runnerConfig": builderParam,
}) })

View File

@ -52,6 +52,7 @@ type Cmd struct {
CopyOutCached []string `json:"copyOutCached"` CopyOutCached []string `json:"copyOutCached"`
CopyOutMax uint64 `json:"copyOutMax"` CopyOutMax uint64 `json:"copyOutMax"`
CopyOutDir string `json:"copyOutDir"` CopyOutDir string `json:"copyOutDir"`
CopyOutTruncate bool `json:"copyOutTruncate"`
TTY bool `json:"tty,omitempty"` TTY bool `json:"tty,omitempty"`
StrictMemoryLimit bool `json:"strictMemoryLimit"` StrictMemoryLimit bool `json:"strictMemoryLimit"`
@ -317,6 +318,7 @@ func convertCmd(c Cmd, srcPrefix []string) (worker.Cmd, error) {
CopyOutCached: convertCopyOut(c.CopyOutCached), CopyOutCached: convertCopyOut(c.CopyOutCached),
CopyOutMax: c.CopyOutMax, CopyOutMax: c.CopyOutMax,
CopyOutDir: c.CopyOutDir, CopyOutDir: c.CopyOutDir,
CopyOutTruncate: c.CopyOutTruncate,
} }
for _, f := range c.Files { for _, f := range c.Files {
cf, err := convertCmdFile(f, srcPrefix) cf, err := convertCmdFile(f, srcPrefix)

View File

@ -52,6 +52,7 @@ type Cmd struct {
// file names to copyout after exec // file names to copyout after exec
CopyOut []CmdCopyOutFile CopyOut []CmdCopyOutFile
CopyOutMax Size // file size limit CopyOutMax Size // file size limit
CopyOutTruncate bool
// CopyOutDir specifies a dir to dump all /w content // CopyOutDir specifies a dir to dump all /w content
CopyOutDir string CopyOutDir string

View File

@ -98,8 +98,11 @@ func copyOutFile(
s := stat.Size() s := stat.Size()
if c.CopyOutMax > 0 && s > int64(c.CopyOutMax) { if c.CopyOutMax > 0 && s > int64(c.CopyOutMax) {
t = ErrCopyOutSizeExceeded t = ErrCopyOutSizeExceeded
if !c.CopyOutTruncate {
return fmt.Errorf("copyout: %q size (%d) exceeds limit (%d)", n.Name, s, c.CopyOutMax) return fmt.Errorf("copyout: %q size (%d) exceeds limit (%d)", n.Name, s, c.CopyOutMax)
} }
s = int64(c.CopyOutMax) + 1
}
// create store file // create store file
buf, err := newStoreFile() buf, err := newStoreFile()
if err != nil { if err != nil {
@ -108,13 +111,17 @@ func copyOutFile(
} }
// Ensure not copy over file size // Ensure not copy over file size
_, err = buf.ReadFrom(io.LimitReader(cf, s)) s, err = buf.ReadFrom(io.LimitReader(cf, s))
if err != nil { if err != nil {
t = ErrCopyOutCopyContent t = ErrCopyOutCopyContent
buf.Close() buf.Close()
return fmt.Errorf("copyout: failed to copy content for %q: %w", n.Name, err) return fmt.Errorf("copyout: failed to copy content for %q: %w", n.Name, err)
} }
put(buf, n.Name) 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 return nil
} }

View File

@ -38,6 +38,7 @@ type Cmd struct {
CopyOutCached []CmdCopyOutFile CopyOutCached []CmdCopyOutFile
CopyOutMax uint64 CopyOutMax uint64
CopyOutDir string CopyOutDir string
CopyOutTruncate bool
TTY bool TTY bool
DataSegmentLimit bool DataSegmentLimit bool

View File

@ -391,6 +391,7 @@ func (w *worker) prepareCmd(rc Cmd, pipeFileName map[string]bool) (*envexec.Cmd,
CopyOut: copyOut, CopyOut: copyOut,
CopyOutDir: copyOutDir, CopyOutDir: copyOutDir,
CopyOutMax: copyOutMax, CopyOutMax: copyOutMax,
CopyOutTruncate: rc.CopyOutTruncate,
Waiter: wait.Wait, Waiter: wait.Wait,
}, nil }, nil
} }