Cyrene_v2/internal/runtime/exec/terminal.go

package exec

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"os"
	"os/exec"
	"sort"
	"strings"
	"sync"
	"time"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/config"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/logger"
)

// TerminalExecutor is the default command Provider. It runs allowlisted commands
// inside a grant-controlled terminal sandbox (DevDocs/design/04-运行时.md §8):
// command allowlist, working-directory confinement, environment allowlist,
// timeout, and output truncation. GUI and virtual-machine providers are reserved
// for later milestones and surface only through the Provider interface.
type TerminalExecutor struct {
	cfg   config.ExecSandbox
	log   *slog.Logger
	roots []string // absolute, clean allowed roots for workdir confinement
}

// New builds a TerminalExecutor from cfg and log. The sandbox is closed unless
// cfg.Enabled is true; a closed executor still constructs fine and refuses every
// request with ErrExecutionDisabled. When enabled, at least one allowed root is
// required so workdir confinement has an anchor.
func New(cfg config.ExecSandbox, log *slog.Logger) (*TerminalExecutor, error) {
	if log == nil {
		return nil, ErrNilLogger
	}
	roots := absCleaners(cfg.AllowedRoots)
	if cfg.Enabled && len(roots) == 0 {
		return nil, fmt.Errorf("%w: an enabled sandbox needs at least one allowed root", ErrNoAllowedRoot)
	}
	return &TerminalExecutor{
		cfg:   cfg,
		log:   logger.WithSource(log, "runtime/exec"),
		roots: roots,
	}, nil
}

// Execute runs the command in req under the sandbox. It returns an
// *[ExecResult] on success (including a non-zero exit code) or a sentinel error
// (wrapped) on refusal or an underlying context failure.
func (t *TerminalExecutor) Execute(ctx context.Context, req ExecRequest) (*ExecResult, error) {
	if !t.cfg.Enabled {
		t.log.Warn("sandbox disabled; refusing to execute command")
		return nil, ErrExecutionDisabled
	}
	if len(req.Command) == 0 {
		return nil, fmt.Errorf("%w: no argv", ErrEmptyCommand)
	}
	if err := t.checkCommand(req.Command); err != nil {
		return nil, err
	}
	workdir, err := t.resolveWorkdir(req.WorkDir)
	if err != nil {
		return nil, err
	}
	env, err := t.buildEnv(req.Env)
	if err != nil {
		return nil, err
	}
	return t.runCmd(ctx, req, workdir, env)
}

// runCmd launches the already-validated command, streams its output through a
// capped writer, honors the deadline, and classifies the outcome.
func (t *TerminalExecutor) runCmd(ctx context.Context, req ExecRequest, workdir string, env []string) (*ExecResult, error) {
	timeout := req.Timeout
	if timeout <= 0 {
		timeout = DefaultTimeout
	}
	maxOut := req.MaxOutputBytes
	if maxOut <= 0 {
		maxOut = DefaultMaxOutputBytes
	}

	runCtx := ctx
	cancel := func() {}
	if timeout > 0 {
		runCtx, cancel = context.WithTimeout(ctx, timeout)
	}
	defer cancel()

	cmd := exec.CommandContext(runCtx, req.Command[0], req.Command[1:]...)
	cmd.Dir = workdir
	cmd.Env = env
	if len(req.Stdin) > 0 {
		cmd.Stdin = bytes.NewReader(req.Stdin)
	}
	out := &cappedOutput{max: maxOut}
	cmd.Stdout = out.stdoutWriter()
	cmd.Stderr = out.stderrWriter()

	t.log.Debug("starting command", "cmd", strings.Join(req.Command, " "),
		"workdir", workdir, "timeout", timeout.String())
	startedAt := time.Now()
	if err := cmd.Start(); err != nil {
		return nil, fmt.Errorf("%w: %s: %w", ErrExecFailed, strings.Join(req.Command, " "), err)
	}
	result := &ExecResult{StartedAt: startedAt, PID: cmd.Process.Pid}

	waitErr := cmd.Wait()
	result.EndedAt = time.Now()
	result.Stdout = out.stdout.String()
	result.Stderr = out.stderr.String()
	result.Truncated = out.truncated

	if waitErr != nil {
		switch {
		case errors.Is(runCtx.Err(), context.DeadlineExceeded):
			t.log.Warn("command timed out; killed", "cmd", strings.Join(req.Command, " "), "pid", result.PID)
			return nil, fmt.Errorf("%w: %s", ErrTimeout, strings.Join(req.Command, " "))
		case errors.Is(runCtx.Err(), context.Canceled):
			return nil, fmt.Errorf("%w: %s", context.Canceled, strings.Join(req.Command, " "))
		}
		var exitErr *exec.ExitError
		if errors.As(waitErr, &exitErr) {
			result.ExitCode = exitErr.ExitCode()
		} else {
			return nil, fmt.Errorf("%w: %s: %w", ErrExecFailed, strings.Join(req.Command, " "), waitErr)
		}
	}
	t.log.Debug("command finished", "cmd", strings.Join(req.Command, " "),
		"exit", result.ExitCode, "truncated", result.Truncated)
	return result, nil
}

// checkCommand enforces the command allowlist and the per-argument injection
// rules. The command runs via exec.Command (never a shell), so metacharacters
// are rejected rather than interpreted.
func (t *TerminalExecutor) checkCommand(argv []string) error {
	if !t.allowed(argv[0]) {
		return fmt.Errorf("%w: %q", ErrCommandNotAllowed, argv[0])
	}
	for _, a := range argv {
		if err := validateArg(a); err != nil {
			return fmt.Errorf("%w: %q: %w", ErrCommandInjection, a, err)
		}
	}
	return nil
}

func (t *TerminalExecutor) allowed(name string) bool {
	for _, c := range t.cfg.CommandList {
		if c == name {
			return true
		}
	}
	return false
}

func (t *TerminalExecutor) envAllowed(name string) bool {
	for _, a := range t.cfg.AllowedEnv {
		if a == name {
			return true
		}
	}
	return false
}

// resolveWorkdir confines the process working directory to an allowed root. An
// empty request uses the first root; a relative path is resolved against it.
func (t *TerminalExecutor) resolveWorkdir(workdir string) (string, error) {
	if workdir == "" {
		if len(t.roots) == 0 {
			return "", ErrNoAllowedRoot
		}
		real, err := resolveReal(t.roots[0])
		if err != nil {
			return "", fmt.Errorf("%w: %s", ErrWorkDirOutsideRoot, t.roots[0])
		}
		return real, nil
	}
	abs, err := resolveInsideRoot(workdir, t.roots)
	if err != nil {
		if errors.Is(err, ErrPathOutsideRoot) || errors.Is(err, ErrSymlinkEscape) {
			return "", fmt.Errorf("%w: %s", ErrWorkDirOutsideRoot, workdir)
		}
		return "", err
	}
	return abs, nil
}

// buildEnv returns the child environment: the base process environment plus any
// allowlisted overrides from req.Env. Names outside the allowlist are rejected.
func (t *TerminalExecutor) buildEnv(extra map[string]string) ([]string, error) {
	base := os.Environ()
	if len(extra) == 0 {
		return base, nil
	}

	overlays := make(map[string]string, len(extra))
	for k, v := range extra {
		if !t.envAllowed(k) {
			return nil, fmt.Errorf("%w: %s", ErrEnvNotAllowed, k)
		}
		overlays[k] = v
	}

	out := make([]string, 0, len(base)+len(overlays))
	seen := make(map[string]struct{}, len(base)+len(overlays))
	for _, e := range base {
		k, _, _ := strings.Cut(e, "=")
		if v, ok := overlays[k]; ok {
			out = append(out, k+"="+v)
			delete(overlays, k)
		} else {
			out = append(out, e)
		}
		seen[k] = struct{}{}
	}

	rest := make([]string, 0, len(overlays))
	for k := range overlays {
		if _, ok := seen[k]; !ok {
			rest = append(rest, k)
		}
	}
	sort.Strings(rest)
	for _, k := range rest {
		out = append(out, k+"="+overlays[k])
	}
	return out, nil
}

// validateArg rejects anything unsafe in a command argument: empty value, shell
// metacharacters, control characters, or a path-traversal ("..") component.
func validateArg(s string) error {
	if s == "" {
		return errors.New("empty argument")
	}
	if strings.ContainsAny(s, `;&|`+"`$()<>") {
		return errors.New("contains shell metacharacters")
	}
	if strings.ContainsAny(s, "\x00\r\n") {
		return errors.New("contains control characters")
	}
	for _, sep := range []string{"/", `\`} {
		for _, part := range strings.Split(s, sep) {
			if part == ".." {
				return errors.New("contains path traversal")
			}
		}
	}
	return nil
}

// cappedOutput captures stdout and stderr while enforcing a combined byte cap.
// Writes past the cap are dropped and the truncated flag is raised. The shared
// mutex makes the two streams safe to write concurrently by os/exec.
type cappedOutput struct {
	mu        sync.Mutex
	stdout    bytes.Buffer
	stderr    bytes.Buffer
	max       int
	used      int
	truncated bool
}

func (o *cappedOutput) stdoutWriter() io.Writer { return &capWriter{o: o, buf: &o.stdout} }
func (o *cappedOutput) stderrWriter() io.Writer { return &capWriter{o: o, buf: &o.stderr} }

type capWriter struct {
	o   *cappedOutput
	buf *bytes.Buffer
}

func (w *capWriter) Write(p []byte) (int, error) {
	w.o.mu.Lock()
	defer w.o.mu.Unlock()
	n := len(p)
	remaining := w.o.max - w.o.used
	if remaining <= 0 {
		w.o.truncated = true
		return n, nil
	}
	take := remaining
	if take > n {
		take = n
	}
	w.buf.Write(p[:take])
	w.o.used += take
	if take < n {
		w.o.truncated = true
	}
	return n, nil
}