fix(executor/sandbox): stream in tar file to avoid SOCK_SEQPACKET size limit

internal/executor/sandbox/executor.go

@@ -1,16 +1,26 @@
 package sandbox
 
 import (
+	"archive/tar"
+	"bytes"
 	"context"
 	"fmt"
+	"io"
 	"log/slog"
 	"maps"
+	"os"
+	"path/filepath"
 
 	"github.com/criyle/go-judge/pb"
 	"github.com/joint-online-judge/JOJ3/internal/stage"
 	"google.golang.org/protobuf/proto"
 )
 
+const (
+	tarStreamThreshold = 128 * 1024
+	streamChunkSize    = 32 * 1024
+)
+
 func (e *Sandbox) Run(cmds []stage.Cmd) ([]stage.ExecutorResult, error) {
 	var err error
 	if e.execClient == nil {
@@ -32,6 +42,34 @@ func (e *Sandbox) Run(cmds []stage.Cmd) ([]stage.ExecutorResult, error) {
 			}
 		}
 	}
+	if needStream, tarData := prepareTarStream(cmds); needStream {
+		return e.runWithTarStream(cmds, tarData)
+	}
+	return e.runUnary(cmds)
+}
+
+func prepareTarStream(cmds []stage.Cmd) (bool, []byte) {
+	if len(cmds) == 0 {
+		return false, nil
+	}
+	for i := range cmds {
+		if cmds[i].CopyInDir != "" &&
+			estimateCopyInSize(&cmds[i]) >= tarStreamThreshold {
+			tarData, err := createCopyInTar(&cmds[i])
+			if err != nil {
+				slog.Error("create copyIn tar", "error", err)
+				return false, nil
+			}
+			for j := range cmds {
+				cmds[j].CopyInDir = ""
+			}
+			return true, tarData
+		}
+	}
+	return false, nil
+}
+
+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 +91,205 @@ func (e *Sandbox) Run(cmds []stage.Cmd) ([]stage.ExecutorResult, error) {
 	return results, nil
 }
 
+func (e *Sandbox) runWithTarStream(cmds []stage.Cmd, tarData []byte) ([]stage.ExecutorResult, error) {
+	tarCmd := stage.Cmd{
+		// FIXME: `/w` won't work if go-judge change default working directory
+		Args:         []string{"/bin/tar", "xf", "-", "-C", "/w"},
+		Env:          cmds[0].Env,
+		Stdin:        &stage.CmdFile{StreamIn: true},
+		CPULimit:     cmds[0].CPULimit,
+		ClockLimit:   cmds[0].ClockLimit,
+		MemoryLimit:  cmds[0].MemoryLimit,
+		StackLimit:   cmds[0].StackLimit,
+		ProcLimit:    cmds[0].ProcLimit,
+		CPURateLimit: cmds[0].CPURateLimit,
+		CPUSetLimit:  cmds[0].CPUSetLimit,
+	}
+
+	allCmds := append([]stage.Cmd{tarCmd}, cmds...)
+	pbCmds := convertPBCmd(allCmds)
+
+	pbReq := &pb.Request{}
+	pbReq.SetCmd(pbCmds)
+	slog.Debug(
+		"sandbox stream execute",
+		"pbReq size", proto.Size(pbReq),
+		"tarSize", len(tarData),
+	)
+
+	stream, err := e.execClient.ExecStream(context.TODO())
+	if err != nil {
+		return nil, fmt.Errorf("exec stream: %w", err)
+	}
+
+	sr := &pb.StreamRequest{}
+	sr.SetExecRequest(pbReq)
+	if err := stream.Send(sr); err != nil {
+		return nil, fmt.Errorf("stream send request: %w", err)
+	}
+
+	for offset := 0; offset < len(tarData); offset += streamChunkSize {
+		end := offset + streamChunkSize
+		if end > len(tarData) {
+			end = len(tarData)
+		}
+		input := pb.StreamRequest_builder{
+			ExecInput: (&pb.StreamRequest_Input_builder{
+				Index:   0,
+				Fd:      0,
+				Content: tarData[offset:end],
+			}).Build(),
+		}.Build()
+		if err := stream.Send(input); err != nil {
+			return nil, fmt.Errorf("stream send input: %w", err)
+		}
+	}
+
+	if err := stream.CloseSend(); err != nil {
+		return nil, fmt.Errorf("stream close send: %w", err)
+	}
+
+	var finalResponse *pb.Response
+	for {
+		resp, err := stream.Recv()
+		if err == io.EOF {
+			break
+		}
+		if err != nil {
+			return nil, fmt.Errorf("stream recv: %w", err)
+		}
+		if resp.HasExecResponse() {
+			finalResponse = resp.GetExecResponse()
+		}
+	}
+
+	if finalResponse == nil {
+		return nil, fmt.Errorf("stream: no exec response received")
+	}
+	if finalResponse.GetError() != "" {
+		return nil, fmt.Errorf("stream execute error: %s", finalResponse.GetError())
+	}
+
+	results := convertPBResult(finalResponse.GetResults())
+	if len(results) == 0 {
+		return nil, fmt.Errorf("stream: empty results")
+	}
+	tarResult := results[0]
+	if tarResult.Status != stage.StatusAccepted || tarResult.ExitStatus != 0 {
+		return nil, fmt.Errorf(
+			"tar extraction failed: status=%v, exit=%d, error=%s",
+			tarResult.Status, tarResult.ExitStatus, tarResult.Error,
+		)
+	}
+
+	results = results[1:]
+	for _, result := range results {
+		maps.Copy(e.cachedMap, result.FileIDs)
+	}
+	return results, nil
+}
+
+func estimateCopyInSize(cmd *stage.Cmd) int {
+	total := 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 {
+					total += int(info.Size())
+				}
+				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 {
+				total += int(fi.Size())
+			}
+		} else if f.Content != nil {
+			total += len(*f.Content)
+		}
+	}
+	return total
+}
+
+func createCopyInTar(cmd *stage.Cmd) ([]byte, error) {
+	var buf bytes.Buffer
+	tw := tar.NewWriter(&buf)
+
+	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 = relPath
+					return tw.WriteHeader(hdr)
+				}
+				hdr, err := tar.FileInfoHeader(info, "")
+				if err != nil {
+					return err
+				}
+				hdr.Name = 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 {
+			return nil, fmt.Errorf("walk copyInDir: %w", err)
+		}
+	}
+
+	if err := tw.Close(); err != nil {
+		return nil, fmt.Errorf("close tar: %w", err)
+	}
+	return buf.Bytes(), nil
+}
+
 func (e *Sandbox) Cleanup() error {
 	for k, fileID := range e.cachedMap {
 		req := &pb.FileID{}