orca.ai/pkg/sandbox/process.go

package sandbox

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path/filepath"
)

const (
	// DefaultOutputLimit is the maximum number of bytes captured from stdout/stderr (512 KB).
	DefaultOutputLimit = 512 * 1024

	// DefaultWorkingDir is the default working directory for sandboxed commands.
	DefaultWorkingDir = "."
)

// AllowedEnvVars is the whitelist of environment variables accessible inside the sandbox.
// Only these variables are passed through from the parent process.
var AllowedEnvVars = []string{
	"HOME",
	"USER",
	"PATH",
	"LANG",
	"SHELL",
	"TMPDIR",
	"ORCA_HOME",
}

// ProcessSandbox is a Sandbox implementation that uses os/exec to run commands
// as child processes with resource restrictions.
type ProcessSandbox struct {
	// WorkingDir is the directory in which commands execute.
	WorkingDir string
	// OutputLimit is the maximum bytes to capture from stdout/stderr.
	OutputLimit int
	// EnvWhitelist controls which environment variables are passed to child processes.
	// If nil, AllowedEnvVars is used. If empty, no env vars are passed.
	EnvWhitelist []string
}

// NewProcessSandbox creates a ProcessSandbox with sensible defaults.
func NewProcessSandbox() *ProcessSandbox {
	return &ProcessSandbox{
		WorkingDir:   DefaultWorkingDir,
		OutputLimit:  DefaultOutputLimit,
		EnvWhitelist: nil, // uses AllowedEnvVars
	}
}

// Execute runs a command as a child process with resource restrictions.
func (ps *ProcessSandbox) Execute(ctx context.Context, cmd string, args ...string) (*Result, error) {
	// Ensure working directory exists
	if err := os.MkdirAll(ps.WorkingDir, 0755); err != nil {
		return nil, fmt.Errorf("sandbox: failed to create working directory %q: %w", ps.WorkingDir, err)
	}

	// Build the command
	c := exec.CommandContext(ctx, cmd, args...)
	c.Dir = ps.WorkingDir

	// Set up environment variable whitelist
	env := ps.buildEnv()
	c.Env = env

	// Capture stdout and stderr with size limits
	stdoutBuf := newLimitedBuffer(ps.outputLimit())
	stderrBuf := newLimitedBuffer(ps.outputLimit())

	stdoutPipe, err := c.StdoutPipe()
	if err != nil {
		return nil, fmt.Errorf("sandbox: failed to create stdout pipe: %w", err)
	}
	stderrPipe, err := c.StderrPipe()
	if err != nil {
		return nil, fmt.Errorf("sandbox: failed to create stderr pipe: %w", err)
	}

	// Start the command
	if err := c.Start(); err != nil {
		return nil, fmt.Errorf("sandbox: failed to start command: %w", err)
	}

	// Read stdout and stderr concurrently
	var readStdout, readStderr error
	var wg syncWaitGroup
	wg.Add(2)

	go func() {
		defer wg.Done()
		_, readStdout = io.Copy(stdoutBuf, stdoutPipe)
		// ErrShortWrite is expected when the output limit is reached — not a real error.
		if readStdout != nil && readStdout != io.EOF && readStdout != io.ErrShortWrite {
			readStdout = fmt.Errorf("sandbox: stdout read error: %w", readStdout)
		} else {
			readStdout = nil
		}
	}()

	go func() {
		defer wg.Done()
		_, readStderr = io.Copy(stderrBuf, stderrPipe)
		if readStderr != nil && readStderr != io.EOF && readStderr != io.ErrShortWrite {
			readStderr = fmt.Errorf("sandbox: stderr read error: %w", readStderr)
		} else {
			readStderr = nil
		}
	}()

	wg.Wait()

	// Wait for the command to finish
	err = c.Wait()
	exitCode := 0

	if err != nil {
		// Check if the process was killed due to context cancellation (timeout)
		if ctx.Err() != nil {
			return nil, fmt.Errorf("sandbox: command timed out: %w", ctx.Err())
		}

		// Normal non-zero exit
		if exitErr, ok := err.(*exec.ExitError); ok {
			exitCode = exitErr.ExitCode()
			err = nil
		}
	}

	// Combine errors: prefer command error, then read errors
	if err != nil {
		return nil, err
	}
	if readStdout != nil {
		return nil, readStdout
	}
	if readStderr != nil {
		return nil, readStderr
	}

	return &Result{
		Stdout:   stdoutBuf.String(),
		Stderr:   stderrBuf.String(),
		ExitCode: exitCode,
	}, nil
}

// buildEnv constructs the environment variable list for the child process
// based on the whitelist configuration.
func (ps *ProcessSandbox) buildEnv() []string {
	whitelist := ps.EnvWhitelist
	if whitelist == nil {
		whitelist = AllowedEnvVars
	}

	env := make([]string, 0, len(whitelist))
	for _, key := range whitelist {
		if val, ok := os.LookupEnv(key); ok {
			env = append(env, key+"="+val)
		}
	}
	return env
}

// outputLimit returns the effective output size limit.
func (ps *ProcessSandbox) outputLimit() int {
	if ps.OutputLimit <= 0 {
		return DefaultOutputLimit
	}
	return ps.OutputLimit
}

// WorkingDirPath returns the absolute path of the sandbox working directory.
func (ps *ProcessSandbox) WorkingDirPath() string {
	abs, err := filepath.Abs(ps.WorkingDir)
	if err != nil {
		return ps.WorkingDir
	}
	return abs
}

// ---------------------------------------------------------------------------
// limitedBuffer — a writer that stops accepting data after MaxSize bytes.
// Uses a named field (not embedded) to avoid promoting bytes.Buffer.ReadFrom
// which would bypass the size limit when used with io.Copy.
// ---------------------------------------------------------------------------

type limitedBuffer struct {
	buf     bytes.Buffer
	MaxSize int
}

func newLimitedBuffer(maxSize int) *limitedBuffer {
	return &limitedBuffer{MaxSize: maxSize}
}

func (lb *limitedBuffer) Write(p []byte) (int, error) {
	remaining := lb.MaxSize - lb.buf.Len()
	if remaining <= 0 {
		return len(p), nil // silently drop excess; io.Copy sees nw==nr, continues draining pipe
	}
	if len(p) > remaining {
		p = p[:remaining]
		n, err := lb.buf.Write(p)
		// Return n < original len(p) so io.Copy stops with ErrShortWrite.
		return n, err
	}
	return lb.buf.Write(p)
}

func (lb *limitedBuffer) String() string {
	return lb.buf.String()
}

func (lb *limitedBuffer) Len() int {
	return lb.buf.Len()
}

// ---------------------------------------------------------------------------
// syncWaitGroup — a simple goroutine synchronization mechanism.
// ---------------------------------------------------------------------------

type syncWaitGroup struct {
	ch chan struct{}
}

func (wg *syncWaitGroup) Add(n int) {
	if wg.ch == nil {
		wg.ch = make(chan struct{}, n)
	}
}

func (wg *syncWaitGroup) Done() {
	wg.ch <- struct{}{}
}

func (wg *syncWaitGroup) Wait() {
	for i := 0; i < cap(wg.ch); i++ {
		<-wg.ch
	}
}

// Compile-time interface check.
var _ Sandbox = (*ProcessSandbox)(nil)