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
}

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 {
	result := &SuiteResult{}

	buildLog, err := r.build()
	result.BuildLog = buildLog
	if err != nil {
		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,
				})
			}
			result.Groups = append(result.Groups, gr)
		}
		return result
	}

	r.binary = resolveBinary(r.cfg.WorkDir, filepath.Base(r.binary))

	for _, g := range r.file.Groups {
		gr := r.runGroup(g)
		result.Groups = append(result.Groups, gr)
		result.TotalScore += gr.Score
	}

	return result
}

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) build() (string, error) {
	buildCmd := r.buildCommand()

	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 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.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

	start := time.Now()
	runErr := cmd.Run()
	tr.Elapsed = time.Since(start)

	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
	}

	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 applyMatcher("stdout", t.Stdout, tr.ActualStdout) {
		tr.addFailure("%s", f)
	}
	for _, f := range applyMatcher("stderr", t.Stderr, tr.ActualStderr) {
		tr.addFailure("%s", f)
	}

	for name, expected := range t.OutFiles {
		path := filepath.Join(tmpDir, name)
		content, err := os.ReadFile(path)
		if err != nil {
			tr.addFailure("output file %q not found: %v", name, err)
			continue
		}
		actual := normalizeOutput(string(content), r.file)
		for _, f := range applyMatcher(fmt.Sprintf("file(%s)", name), dsl.ExactMatcher{Value: expected}, actual) {
			tr.addFailure("%s", f)
		}
	}

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

	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)