From 62b1538b6ae181f9a15101d28345d930eba0566a Mon Sep 17 00:00:00 2001 From: Boming Zhang Date: Fri, 24 Jul 2026 13:04:09 +0800 Subject: [PATCH] fix: sandbox executor pack size limit (#101) Reviewed-on: https://focs.gc.sjtu.edu.cn/git/JOJ/JOJ3/pulls/101 Co-authored-by: Boming Zhang Co-committed-by: Boming Zhang --- internal/executor/sandbox/executor.go | 261 +++++++++++++++++++++++++- 1 file changed, 256 insertions(+), 5 deletions(-) diff --git a/internal/executor/sandbox/executor.go b/internal/executor/sandbox/executor.go index 987382f..65c3369 100644 --- a/internal/executor/sandbox/executor.go +++ b/internal/executor/sandbox/executor.go @@ -1,16 +1,27 @@ package sandbox import ( + "archive/tar" + "bytes" "context" "fmt" + "io" "log/slog" "maps" + "os" + "path/filepath" + "strings" "github.com/criyle/go-judge/pb" "github.com/joint-online-judge/JOJ3/internal/stage" "google.golang.org/protobuf/proto" ) +const ( + tarSizeThreshold = 128 * 1024 // 128 KB + tarCountThreshold = 100 // 100 files +) + func (e *Sandbox) Run(cmds []stage.Cmd) ([]stage.ExecutorResult, error) { var err error if e.execClient == nil { @@ -20,18 +31,51 @@ func (e *Sandbox) Run(cmds []stage.Cmd) ([]stage.ExecutorResult, error) { return nil, err } } - // cannot use range loop since we need to change the value for i := 0; i < len(cmds); i += 1 { - cmd := &cmds[i] - if cmd.CopyIn == nil { + if cmd := &cmds[i]; cmd.CopyIn == nil { cmd.CopyIn = make(map[string]stage.CmdFile) } - for k, v := range cmd.CopyInCached { + for k, v := range cmds[i].CopyInCached { if fileID, ok := e.cachedMap[v]; ok { - cmd.CopyIn[k] = stage.CmdFile{FileID: &fileID} + cmds[i].CopyIn[k] = stage.CmdFile{FileID: &fileID} } } } + if needTar, tarData := prepareTar(cmds); needTar { + return e.runWithTar(cmds, tarData) + } + return e.runUnary(cmds) +} + +func prepareTar(cmds []stage.Cmd) (bool, []byte) { + if len(cmds) == 0 { + return false, nil + } + for i := range cmds { + if shouldTar(&cmds[i]) { + tarData, keysInTar := createCopyInTar(&cmds[i]) + if tarData == nil { + return false, nil + } + slog.Debug("prepareTar", "tarSize", len(tarData), "keysInTar", len(keysInTar)) + for j := range cmds { + cmds[j].CopyInDir = "" + for _, k := range keysInTar { + delete(cmds[j].CopyIn, k) + } + } + return true, tarData + } + } + return false, nil +} + +func shouldTar(cmd *stage.Cmd) bool { + size, count := estimateCopyIn(cmd) + return size >= tarSizeThreshold || count >= tarCountThreshold +} + +func (e *Sandbox) runUnary(cmds []stage.Cmd) ([]stage.ExecutorResult, error) { pbCmds := convertPBCmd(cmds) for i, pbCmd := range pbCmds { slog.Debug("sandbox execute", "i", i, "pbCmd size", proto.Size(pbCmd)) @@ -53,6 +97,213 @@ func (e *Sandbox) Run(cmds []stage.Cmd) ([]stage.ExecutorResult, error) { return results, nil } +func (e *Sandbox) runWithTar(cmds []stage.Cmd, tarData []byte) ([]stage.ExecutorResult, error) { + fc := &pb.FileContent{} + fc.SetContent(tarData) + fileIDResp, err := e.execClient.FileAdd(context.TODO(), fc) + if err != nil { + return nil, fmt.Errorf("file add tar: %w", err) + } + fid := fileIDResp.GetFileID() + slog.Debug("sandbox tar uploaded", "fileID", fid, "tarSize", len(tarData)) + + defer func() { + deleteReq := &pb.FileID{} + deleteReq.SetFileID(fid) + if _, err := e.execClient.FileDelete(context.TODO(), deleteReq); err != nil { + slog.Warn("sandbox tar file delete", "fileID", fid, "error", err) + } + }() + + tarFileName := "/w/__joj3_copyin.tar" + script := fmt.Sprintf( + "/bin/tar xf %s -C / --no-same-owner && rm -f %s && exec \"$@\"", + tarFileName, tarFileName, + ) + + for i := range cmds { + if cmds[i].CopyIn == nil { + cmds[i].CopyIn = make(map[string]stage.CmdFile) + } + cmds[i].CopyIn[tarFileName] = stage.CmdFile{FileID: &fid} + cmds[i].Args = append([]string{ + "/bin/sh", "-c", script, "_", + }, cmds[i].Args...) + } + + slog.Debug("sandbox tar exec", "cmd", cmds[0].Args[:3]) + + return e.runUnary(cmds) +} + +func estimateCopyIn(cmd *stage.Cmd) (int, int) { + totalSize := 0 + totalCount := 0 + if cmd.CopyInDir != "" { + _ = filepath.Walk(cmd.CopyInDir, + func(path string, info os.FileInfo, err error) error { + if err != nil || info.IsDir() { + return nil + } + relPath, err := filepath.Rel(cmd.CopyInDir, path) + if err != nil { + return nil + } + if _, exists := cmd.CopyIn[relPath]; !exists { + totalSize += int(info.Size()) + totalCount++ + } + return nil + }) + } + for _, f := range cmd.CopyIn { + if f.Symlink != nil { + continue + } + if f.Src != nil { + if fi, err := os.Stat(*f.Src); err == nil { + totalSize += int(fi.Size()) + totalCount++ + } + } else if f.Content != nil { + totalSize += len(*f.Content) + totalCount++ + } + } + return totalSize, totalCount +} + +func formatTarPath(p string) string { + if filepath.IsAbs(p) { + return strings.TrimPrefix(p, "/") + } + return "w/" + p +} + +func createCopyInTar(cmd *stage.Cmd) ([]byte, []string) { + var buf bytes.Buffer + tw := tar.NewWriter(&buf) + tarKeys := make([]string, 0, len(cmd.CopyIn)) + + if cmd.CopyInDir != "" { + err := filepath.Walk(cmd.CopyInDir, + func(path string, info os.FileInfo, err error) error { + if err != nil { + return err + } + relPath, err := filepath.Rel(cmd.CopyInDir, path) + if err != nil { + return err + } + if relPath == "." { + return nil + } + if _, exists := cmd.CopyIn[relPath]; exists { + if info.IsDir() { + return filepath.SkipDir + } + return nil + } + if info.Mode()&os.ModeSymlink != 0 { + link, err := os.Readlink(path) + if err != nil { + return err + } + hdr, err := tar.FileInfoHeader(info, link) + if err != nil { + return err + } + hdr.Name = formatTarPath(relPath) + return tw.WriteHeader(hdr) + } + hdr, err := tar.FileInfoHeader(info, "") + if err != nil { + return err + } + hdr.Name = formatTarPath(relPath) + if info.IsDir() { + hdr.Name += "/" + } + if err := tw.WriteHeader(hdr); err != nil { + return err + } + if info.IsDir() { + return nil + } + f, err := os.Open(path) + if err != nil { + return err + } + _, err = io.Copy(tw, f) + f.Close() + return err + }) + if err != nil { + slog.Error("create copyIn tar walk", "error", err) + return nil, nil + } + } + + for k, f := range cmd.CopyIn { + if f.FileID != nil || f.Symlink != nil || f.StreamIn || f.StreamOut || f.Pipe { + continue + } + if f.Content != nil { + hdr := &tar.Header{ + Name: formatTarPath(k), + Mode: 0o644, + Size: int64(len(*f.Content)), + } + if err := tw.WriteHeader(hdr); err != nil { + slog.Error("create copyIn tar write header", "key", k, "error", err) + continue + } + if _, err := tw.Write([]byte(*f.Content)); err != nil { + slog.Error("create copyIn tar write content", "key", k, "error", err) + continue + } + tarKeys = append(tarKeys, k) + } else if f.Src != nil { + fi, err := os.Stat(*f.Src) + if err != nil { + slog.Error("create copyIn tar stat", "key", k, "src", *f.Src, "error", err) + continue + } + if fi.IsDir() { + continue + } + hdr, err := tar.FileInfoHeader(fi, "") + if err != nil { + slog.Error("create copyIn tar file info header", "key", k, "error", err) + continue + } + hdr.Name = formatTarPath(k) + if err := tw.WriteHeader(hdr); err != nil { + slog.Error("create copyIn tar write header", "key", k, "error", err) + continue + } + srcFile, err := os.Open(*f.Src) + if err != nil { + slog.Error("create copyIn tar open src", "key", k, "src", *f.Src, "error", err) + continue + } + _, err = io.Copy(tw, srcFile) + srcFile.Close() + if err != nil { + slog.Error("create copyIn tar copy", "key", k, "error", err) + continue + } + tarKeys = append(tarKeys, k) + } + } + + if err := tw.Close(); err != nil { + slog.Error("create copyIn tar close", "error", err) + return nil, nil + } + return buf.Bytes(), tarKeys +} + func (e *Sandbox) Cleanup() error { for k, fileID := range e.cachedMap { req := &pb.FileID{}