woj-server/internal/service/runner/nsjail.go

package runner

import (
	"bufio"
	"fmt"
	"git.0x7f.app/WOJ/woj-server/pkg/file"
	"git.0x7f.app/WOJ/woj-server/pkg/utils"
	"go.uber.org/zap"
	"io"
	"os"
	"os/exec"
	"runtime"
	"strconv"
	"time"
)

type ProgramArgs struct {
	Args []string
	Env  []string
}

func (p *ProgramArgs) EnvArgs() []string {
	var env []string
	for _, e := range p.Env {
		env = append(env, "-E", e)
	}
	return env
}

type MountInfo struct {
	Source      string
	Destination string
	Readonly    bool
}

func (m *MountInfo) Args() []string {
	mapping := m.Source + ":" + m.Destination
	if m.Readonly {
		return []string{"-R", mapping}
	} else {
		return []string{"-B", mapping}
	}
}

type RuntimeArgs struct {
	Rootfs  string
	CPU     int
	Pid     int64
	Memory  uint64 // Memory is in bytes
	Timeout time.Duration
	Mount   []MountInfo
}

func (r *RuntimeArgs) Args() []string {
	cpus := runtime.NumCPU()
	cpus = utils.If(r.CPU < cpus, r.CPU, cpus)
	cpus = utils.If(r.CPU == 0, 1, cpus)

	args := []string{
		"-c", r.Rootfs,
		"--cgroup_pids_max", strconv.FormatInt(r.Pid, 10),
		"--cgroup_mem_max", strconv.FormatUint(r.Memory, 10),
		"--max_cpus", strconv.FormatInt(int64(cpus), 10),
		"-t", strconv.FormatInt(int64(r.Timeout.Seconds()), 10),
	}
	for _, m := range r.Mount {
		args = append(args, m.Args()...)
	}
	return args
}

type IOArgs struct {
	Output *os.File
	// Limit is the max size of output in chars.
	// if Limit = 0, output to stderr if verbose, discard output if not.
	// if Limit < 0, discard output.
	Limit int64
}

type RunArgs struct {
	Program ProgramArgs
	Runtime RuntimeArgs
	IO      IOArgs
}

func (s *service) JailRun(arg *RunArgs) (RuntimeStatus, error) {
	// check cgroup creation
	if arg.Runtime.Pid == 0 && arg.Runtime.Memory == 0 {
		s.log.Warn("cgroup pid and memory not set, resource tracing by cgroup will not work")
	}

	// create stats file
	statFile, err := os.CreateTemp("", "jail-stats-*")
	if err != nil {
		s.log.Warn("create stats file failed", zap.Error(err))
		return RuntimeStatus{}, err
	}
	_ = statFile.Close()

	// prepare output
	var writer io.Writer = nil
	if arg.IO.Limit == 0 && s.verbose {
		writer = os.Stderr
	} else if arg.IO.Limit > 0 && arg.IO.Output != nil {
		writer = &file.LimitedWriter{
			File:  arg.IO.Output,
			Limit: arg.IO.Limit,
		}
	}

	// build args
	args := []string{
		"--quiet",
		"--use_cgroupv2",
		"-T", "/tmp",
		"-E", "PATH=/bin:/sbin:/usr/bin:/usr/sbin:/usr/local/bin:/usr/local/sbin",
		// following envs must sync with resource/runner
		"-E", "WOJ_LAUNCHER=/woj/framework/scripts/woj_launcher",
		"-E", "TEMPLATE=/woj/framework/template",
		"-E", "TESTLIB=/woj/framework/template/testlib",
		"-E", "PREFIX=/woj",
	}
	args = append(args, "--cgroupv2_mount", s.cgroup)
	args = append(args, "--dump_stats", "--dump_stats_file", statFile.Name())
	args = append(args, arg.Program.EnvArgs()...)
	args = append(args, arg.Runtime.Args()...)
	args = append(args, "--")
	args = append(args, arg.Program.Args...)

	// run
	s.log.Debug("jail run", zap.Strings("args", args))
	cmd := exec.Command(NSJailFile, args...)
	cmd.Dir = Prefix
	if s.verbose {
		cmd.Stdout = writer
		cmd.Stderr = writer
	}
	err = cmd.Run()
	if err != nil {
		s.log.Warn("jail run failed", zap.Error(err))
		return RuntimeStatus{}, err
	}

	// re-open stat file
	statFile, err = os.Open(statFile.Name())
	if err != nil {
		s.log.Error("open stats file failed", zap.Error(err))
		return RuntimeStatus{}, err
	}
	defer func(statFile *os.File) {
		_ = statFile.Close()
		_ = os.Remove(statFile.Name())
	}(statFile)

	// collect metrics
	status := RuntimeStatus{}
	scanner := bufio.NewScanner(statFile)
	for scanner.Scan() {
		var key string
		var value int
		_, _ = fmt.Sscanf(scanner.Text(), "%s %d", &key, &value)
		switch key {
		case "usage_usec":
			status.RealTime = value / 1000
		case "user_usec":
			status.CpuTime = value / 1000
		case "memory.peak":
			status.Memory = value / 1024
		}
	}

	return status, nil
}

func (s *service) JailRunPool(arg *RunArgs) uint64 {
	return s.pool.AddTask(func() (interface{}, error) { return s.JailRun(arg) })
}