judge/runner/runner.go

package runner

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"
	"time"

	"github.com/Mond1c/judge/dsl"
)

const MaxOutputBytes = 16 * 1024 * 1024

type Config struct {
	WorkDir    string
	BinaryName string
	Wrapper    string

	TargetBuild string
}

type Runner struct {
	cfg    Config
	file   *dsl.File
	binary string
}

func New(f *dsl.File, cfg Config) *Runner {
	name := cfg.BinaryName
	if name == "" {
		name = f.Binary
	}
	if name == "" {
		name = "solution"
	}
	absWork, err := filepath.Abs(cfg.WorkDir)
	if err != nil {
		absWork = cfg.WorkDir
	}
	cfg.WorkDir = absWork
	return &Runner{file: f, cfg: cfg, binary: resolveBinary(absWork, name)}
}

func resolveBinary(workDir, name string) string {
	primary := filepath.Join(workDir, name)
	if runtime.GOOS == "windows" && !strings.HasSuffix(strings.ToLower(name), ".exe") {
		if _, err := os.Stat(primary); err != nil {
			withExe := primary + ".exe"
			if _, err2 := os.Stat(withExe); err2 == nil {
				return withExe
			}
		}
	}
	return primary
}

func (r *Runner) Run() *SuiteResult {
	defer cleanupCgroupRoot()

	result := &SuiteResult{}

	if len(r.file.Builds) == 0 {
		run := r.runLegacyBuild()
		result.Builds = append(result.Builds, run)
	} else {
		result.Builds = r.runStructuredBuilds()
	}

	result.TotalScore = result.AggregateScore()
	return result
}

func (r *Runner) runLegacyBuild() *BuildRun {
	run := &BuildRun{Name: "default"}

	if r.cfg.TargetBuild != "" && r.cfg.TargetBuild != "default" {
		run.Skipped = true
		run.SkipReason = fmt.Sprintf("--build=%q selected, but this suite has no structured builds", r.cfg.TargetBuild)
		return run
	}

	buildLog, err := r.legacyBuild()
	run.BuildLog = buildLog
	if err != nil {
		r.fillBuildError(run)
		return run
	}

	r.binary = resolveBinary(r.cfg.WorkDir, filepath.Base(r.binary))
	r.runGroups(run)
	return run
}

func (r *Runner) resolveRuntimeToolchain() (Toolchain, string) {
	goos := runtime.GOOS
	wanted := os.Getenv("JUDGE_TOOLCHAIN")
	if wanted == "" {
		wanted = os.Getenv("JUDGE_CC")
	}
	for _, spec := range r.file.Toolchains {
		if spec.Name == wanted {
			return ResolveToolchainSpec(spec), goos
		}
	}
	return ResolveToolchain(wanted), goos
}

func (r *Runner) runStructuredBuilds() []*BuildRun {
	tc, goos := r.resolveRuntimeToolchain()
	var runs []*BuildRun

	for _, b := range r.file.Builds {
		run := &BuildRun{Name: b.Name, Toolchain: tc.Name}

		if r.cfg.TargetBuild != "" && r.cfg.TargetBuild != b.Name {
			continue
		}

		effective := b.Resolve(r.file.BuildDefaults, goos)
		if !effective.AppliesTo(goos, tc.Name) {
			run.Skipped = true
			run.SkipReason = fmt.Sprintf("not applicable to %s/%s (platforms=%v, compilers=%v)", goos, tc.Name, effective.Platforms, effective.Compilers)
			runs = append(runs, run)
			continue
		}

		log, binaryPath, err := r.compileStructured(b.Name, effective, tc)
		run.BuildLog = log
		if err != nil {
			run.Groups = r.synthesizeBuildError()
			run.TotalScore = 0
			runs = append(runs, run)
			continue
		}

		prevBinary := r.binary
		prevWrapper := r.cfg.Wrapper
		r.binary = binaryPath
		if r.cfg.Wrapper == "" && effective.Wrapper != "" {
			r.cfg.Wrapper = effective.Wrapper
		}
		r.runGroups(run)
		r.binary = prevBinary
		r.cfg.Wrapper = prevWrapper

		runs = append(runs, run)
	}

	return runs
}

func (r *Runner) runGroups(run *BuildRun) {
	for _, g := range r.file.Groups {
		gr := r.runGroup(g)
		run.Groups = append(run.Groups, gr)
		run.TotalScore += gr.Score
	}
}

func (r *Runner) fillBuildError(run *BuildRun) {
	run.Groups = r.synthesizeBuildError()
}

func (r *Runner) synthesizeBuildError() []*GroupResult {
	var out []*GroupResult
	for _, g := range r.file.Groups {
		gr := &GroupResult{
			Name:   g.Name,
			Weight: g.Weight,
			Score:  0,
		}
		total := len(g.Tests)
		if g.Pattern != nil {
			total = -1
		}
		gr.Total = total
		for _, t := range g.Tests {
			gr.Tests = append(gr.Tests, &TestResult{
				Name:   t.Name,
				Status: StatusBuildError,
			})
		}
		out = append(out, gr)
	}
	return out
}

func (r *Runner) legacyBuild() (string, error) {
	buildCmd := r.buildCommand()

	sources, err := r.findSources()
	if err != nil {
		return "", err
	}
	if sources != "" {
		buildCmd = strings.ReplaceAll(buildCmd, "$SOURCES", sources)
	}

	ctx := context.Background()
	if r.file.Timeout > 0 {
		var cancel context.CancelFunc
		ctx, cancel = context.WithTimeout(ctx, r.file.Timeout)
		defer cancel()
	}

	cmd := shellCommand(ctx, buildCmd)
	cmd.Dir = r.cfg.WorkDir
	setProcessGroup(cmd)
	cmd.Env = os.Environ()

	var out bytes.Buffer
	cmd.Stdout = &out
	cmd.Stderr = &out

	if err := cmd.Run(); err != nil {
		killProcessGroup(cmd)
		return out.String(), fmt.Errorf("build failed: %w\n%s", err, out.String())
	}
	return out.String(), nil
}

func (r *Runner) compileStructured(name string, cfg dsl.BuildConfig, tc Toolchain) (string, string, error) {
	sources, err := expandSources(r.cfg.WorkDir, cfg.Sources)
	if err != nil {
		return "", "", err
	}
	if len(sources) == 0 {
		return "", "", fmt.Errorf("build %q: no sources", name)
	}
	cfg.Sources = sources

	outputName := cfg.Output
	if outputName == "" {
		outputName = "solution"
	}
	if runtime.GOOS == "windows" && !strings.HasSuffix(strings.ToLower(outputName), ".exe") {
		outputName += ".exe"
	}
	buildDir := filepath.Join(r.cfg.WorkDir, "build", name)
	if err := os.MkdirAll(buildDir, 0755); err != nil {
		return "", "", fmt.Errorf("mkdir %s: %w", buildDir, err)
	}
	outputPath := filepath.Join(buildDir, outputName)

	argv, err := Compile(cfg, tc, outputPath)
	if err != nil {
		return "", "", err
	}

	ctx := context.Background()
	if r.file.Timeout > 0 {
		var cancel context.CancelFunc
		ctx, cancel = context.WithTimeout(ctx, r.file.Timeout)
		defer cancel()
	}

	cmd := exec.CommandContext(ctx, argv[0], argv[1:]...)
	cmd.Dir = r.cfg.WorkDir
	setProcessGroup(cmd)
	cmd.Env = os.Environ()

	var out bytes.Buffer
	cmd.Stdout = &out
	cmd.Stderr = &out

	logPrefix := fmt.Sprintf("$ %s\n", strings.Join(argv, " "))
	if err := cmd.Run(); err != nil {
		killProcessGroup(cmd)
		return logPrefix + out.String(), "", fmt.Errorf("build %q failed: %w\n%s", name, err, out.String())
	}
	return logPrefix + out.String(), outputPath, nil
}

func expandSources(workDir string, patterns []string) ([]string, error) {
	var out []string
	seen := map[string]bool{}
	for _, pat := range patterns {
		matches, err := filepath.Glob(filepath.Join(workDir, pat))
		if err != nil {
			return nil, fmt.Errorf("glob %q: %w", pat, err)
		}
		if len(matches) == 0 {
			if _, statErr := os.Stat(filepath.Join(workDir, pat)); statErr == nil {
				matches = []string{filepath.Join(workDir, pat)}
			} else {
				return nil, fmt.Errorf("source glob %q matched no files", pat)
			}
		}
		for _, m := range matches {
			rel, err := filepath.Rel(workDir, m)
			if err != nil {
				rel = m
			}
			rel = filepath.ToSlash(rel)
			if !seen[rel] {
				seen[rel] = true
				out = append(out, rel)
			}
		}
	}
	return out, nil
}

func (r *Runner) buildCommand() string {
	switch runtime.GOOS {
	case "windows":
		if r.file.BuildWindows != "" {
			return r.file.BuildWindows
		}
	case "linux":
		if r.file.BuildLinux != "" {
			return r.file.BuildLinux
		}
	case "darwin":
		if r.file.BuildDarwin != "" {
			return r.file.BuildDarwin
		}
	}
	if r.file.Build != "" {
		return r.file.Build
	}
	return "go build -o solution ."
}

func (r *Runner) findSources() (string, error) {
	if r.file.Sources == "" {
		return "", nil
	}
	var files []string
	err := filepath.Walk(r.cfg.WorkDir, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if info.IsDir() {
			if info.Name() == ".git" || info.Name() == ".gitea" {
				return filepath.SkipDir
			}
			return nil
		}
		matched, _ := filepath.Match(r.file.Sources, info.Name())
		if matched {
			rel, _ := filepath.Rel(r.cfg.WorkDir, path)
			files = append(files, filepath.ToSlash(rel))
		}
		return nil
	})
	if err != nil {
		return "", fmt.Errorf("source discovery: %w", err)
	}
	if len(files) == 0 {
		return "", fmt.Errorf("no files matching %q found", r.file.Sources)
	}
	return strings.Join(files, " "), nil
}

func shellCommand(ctx context.Context, cmdline string) *exec.Cmd {
	if runtime.GOOS == "windows" {
		return exec.CommandContext(ctx, "cmd", "/C", cmdline)
	}
	return exec.CommandContext(ctx, "sh", "-c", cmdline)
}

func (r *Runner) runGroup(g *dsl.Group) *GroupResult {
	gr := &GroupResult{
		Name:   g.Name,
		Weight: g.Weight,
	}

	tests := g.Tests

	if g.Pattern != nil {
		expanded, err := expandPattern(g.Pattern)
		if err != nil {
			gr.Tests = append(gr.Tests, &TestResult{
				Name:     "pattern_expand",
				Status:   StatusFail,
				Failures: []string{fmt.Sprintf("pattern expand error: %v", err)},
			})
			gr.Total = 1
			gr.Score = 0
			return gr
		}
		tests = append(tests, expanded...)
	}

	gr.Total = len(tests)

	for _, t := range tests {
		for k, v := range g.Env {
			if _, ok := t.Env[k]; !ok {
				t.Env[k] = v
			}
		}
		if t.Timeout == 0 {
			t.Timeout = g.Timeout
		}
		if t.MemoryLimit == 0 {
			t.MemoryLimit = g.MemoryLimit
		}
		if t.Wrapper == "" {
			t.Wrapper = g.Wrapper
		}

		tr := r.runTest(t)
		gr.Tests = append(gr.Tests, tr)
		if tr.Status == StatusPass {
			gr.Passed++
		}
	}

	switch g.Scoring {
	case dsl.ScoringAllOrNone:
		if gr.Passed == gr.Total {
			gr.Score = g.Weight
		} else {
			gr.Score = 0
		}
	default:
		if gr.Total > 0 {
			gr.Score = g.Weight * float64(gr.Passed) / float64(gr.Total)
		}
	}

	return gr
}

func (r *Runner) runTest(t *dsl.Test) *TestResult {
	tr := &TestResult{Name: t.Name, Status: StatusPass}

	tmpDir, err := os.MkdirTemp("", "judge-test-*")
	if err != nil {
		tr.Status = StatusRuntimeError
		tr.addFailure("failed to create temp dir: %v", err)
		return tr
	}
	defer os.RemoveAll(tmpDir)

	for name, content := range t.InFiles {
		path := filepath.Join(tmpDir, name)
		if err := os.MkdirAll(filepath.Dir(path), 0755); err != nil {
			tr.Status = StatusRuntimeError
			tr.addFailure("mkdir for file %q: %v", name, err)
			return tr
		}
		if err := os.WriteFile(path, []byte(content), 0644); err != nil {
			tr.Status = StatusRuntimeError
			tr.addFailure("write input file %q: %v", name, err)
			return tr
		}
	}

	timeout := t.Timeout
	if timeout == 0 {
		timeout = 10 * time.Second
	}

	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()

	wrapper := r.cfg.Wrapper
	if wrapper == "" {
		wrapper = t.Wrapper
	}

	args := t.Args
	if wrapper != "" {
		args = absoluteArgs(tmpDir, args)
	}
	cmd := buildExecCmd(ctx, wrapper, r.binary, args)
	cmd.Dir = tmpDir
	setProcessGroup(cmd)

	cmd.Env = os.Environ()
	cmd.Env = append(cmd.Env, "LC_ALL=C", "LANG=C")
	for k, v := range t.Env {
		cmd.Env = append(cmd.Env, fmt.Sprintf("%s=%s", k, v))
	}

	if t.Stdin != nil {
		cmd.Stdin = strings.NewReader(*t.Stdin)
	}

	stdout := &cappedBuffer{limit: MaxOutputBytes}
	stderr := &cappedBuffer{limit: MaxOutputBytes}
	cmd.Stdout = stdout
	cmd.Stderr = stderr

	tr.MemoryLimit = t.MemoryLimit
	lim := newLimiter(t.MemoryLimit)
	if err := lim.prepare(cmd); err != nil {
		tr.Status = StatusRuntimeError
		tr.addFailure("memory limiter setup: %v", err)
		return tr
	}
	defer lim.cleanup()

	start := time.Now()
	if err := cmd.Start(); err != nil {
		tr.Status = StatusRuntimeError
		tr.addFailure("start: %v", err)
		return tr
	}
	if err := lim.afterStart(cmd); err != nil {
		killProcessGroup(cmd)
		_ = cmd.Wait()
		tr.Status = StatusRuntimeError
		tr.addFailure("memory limiter attach: %v", err)
		return tr
	}
	runErr := cmd.Wait()
	tr.Elapsed = time.Since(start)

	stats := lim.collect()
	tr.PeakMemory = stats.PeakMemory

	if ctx.Err() == context.DeadlineExceeded {
		killProcessGroup(cmd)
	}

	tr.ActualStdout = normalizeOutput(stdout.String(), r.file)
	tr.ActualStderr = normalizeOutput(stderr.String(), r.file)

	if stdout.truncated || stderr.truncated {
		tr.addFailure("output truncated at %d bytes (possible runaway output)", MaxOutputBytes)
	}

	if ctx.Err() == context.DeadlineExceeded {
		tr.Status = StatusTLE
		tr.addFailure("time limit exceeded (%v)", timeout)
		return tr
	}

	if stats.MemoryExceeded {
		tr.Status = StatusMLE
		tr.addFailure("memory limit exceeded (limit %d bytes, peak %d bytes)", t.MemoryLimit, stats.PeakMemory)
		return tr
	}

	actualCode := 0
	if runErr != nil {
		if exitErr, ok := runErr.(*exec.ExitError); ok {
			actualCode = exitErr.ExitCode()
		} else {
			tr.Status = StatusRuntimeError
			tr.addFailure("runtime error: %v", runErr)
			return tr
		}
	}
	tr.ActualCode = actualCode

	if t.ExitCode != nil && actualCode != *t.ExitCode {
		tr.addFailure("exit code: expected %d, got %d", *t.ExitCode, actualCode)
	}

	for _, f := range t.Stdout.Match("stdout", tr.ActualStdout) {
		tr.addFailure("%s", f)
	}
	for _, f := range t.Stderr.Match("stderr", tr.ActualStderr) {
		tr.addFailure("%s", f)
	}

	if len(tr.Failures) > 0 {
		tr.Status = StatusFail
	}

	for name, expected := range t.OutFiles {
		actualPath := filepath.Join(tmpDir, name)
		data, err := os.ReadFile(actualPath)
		if err != nil {
			tr.Status = StatusFail
			tr.addFailure("output file %q missing: %v", name, err)
			continue
		}
		actual := normalizeOutput(string(data), r.file)
		if actual != expected {
			tr.Status = StatusFail
			tr.addFailure("output file %q mismatch\n  expected: %q\n  actual:   %q", name, expected, actual)
		}
	}

	return tr
}

func absoluteArgs(dir string, args []string) []string {
	out := make([]string, len(args))
	for i, a := range args {
		if !filepath.IsAbs(a) {
			out[i] = filepath.Join(dir, a)
		} else {
			out[i] = a
		}
	}
	return out
}

func buildExecCmd(ctx context.Context, wrapper, binary string, args []string) *exec.Cmd {
	if wrapper == "" {
		return exec.CommandContext(ctx, binary, args...)
	}
	wrapperFields := strings.Fields(wrapper)
	full := append(wrapperFields, binary)
	full = append(full, args...)
	return exec.CommandContext(ctx, full[0], full[1:]...)
}

func normalizeOutput(s string, f *dsl.File) string {
	if f.NormalizeCRLF {
		s = strings.ReplaceAll(s, "\r\n", "\n")
		s = strings.ReplaceAll(s, "\r", "\n")
	}
	if f.TrimTrailingWS {
		lines := strings.Split(s, "\n")
		for i := range lines {
			lines[i] = strings.TrimRight(lines[i], " \t\r")
		}
		s = strings.Join(lines, "\n")
	}
	return s
}

type cappedBuffer struct {
	buf       bytes.Buffer
	limit     int
	truncated bool
}

func (c *cappedBuffer) Write(p []byte) (int, error) {
	room := c.limit - c.buf.Len()
	if room <= 0 {
		c.truncated = true
		return len(p), nil
	}
	if len(p) > room {
		c.buf.Write(p[:room])
		c.truncated = true
		return len(p), nil
	}
	return c.buf.Write(p)
}

func (c *cappedBuffer) String() string { return c.buf.String() }

var _ io.Writer = (*cappedBuffer)(nil)