319 lines
7.5 KiB
Go
319 lines
7.5 KiB
Go
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 {
|
|
slog.Debug("create exec client", "server", e.execServer)
|
|
e.execClient, err = createExecClient(e.execServer, e.token)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
}
|
|
for i := 0; i < len(cmds); i += 1 {
|
|
if cmd := &cmds[i]; cmd.CopyIn == nil {
|
|
cmd.CopyIn = make(map[string]stage.CmdFile)
|
|
}
|
|
for k, v := range cmds[i].CopyInCached {
|
|
if fileID, ok := e.cachedMap[v]; ok {
|
|
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))
|
|
}
|
|
pbReq := &pb.Request{}
|
|
pbReq.SetCmd(pbCmds)
|
|
slog.Debug("sandbox execute", "pbReq size", proto.Size(pbReq))
|
|
pbRet, err := e.execClient.Exec(context.TODO(), pbReq)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if pbRet.GetError() != "" {
|
|
return nil, fmt.Errorf("sandbox execute error: %s", pbRet.GetError())
|
|
}
|
|
results := convertPBResult(pbRet.GetResults())
|
|
for _, result := range results {
|
|
maps.Copy(e.cachedMap, result.FileIDs)
|
|
}
|
|
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{}
|
|
req.SetFileID(fileID)
|
|
_, err := e.execClient.FileDelete(context.TODO(), req)
|
|
if err != nil {
|
|
slog.Error("sandbox cleanup", "error", err)
|
|
}
|
|
delete(e.cachedMap, k)
|
|
}
|
|
return nil
|
|
}
|