pkg/tools/sandbox.go

package tools

import (
	"context"
	"fmt"
	"io"
	"os"
	"os/exec"
	"regexp"
	"runtime"
	"strconv"
	"strings"
	"syscall"
	"time"

	"github.com/sipeed/picoclaw/pkg/config"
	"github.com/sipeed/picoclaw/pkg/logger"
	"github.com/sipeed/picoclaw/pkg/metrics"
	"github.com/sipeed/picoclaw/pkg/utils"
)

// LimitedWriter discards output after Limit bytes
type LimitedWriter struct {
	Limit     int
	Written   int
	Writer    io.Writer
	Truncated bool
}

func (w *LimitedWriter) Write(p []byte) (n int, err error) {
	if w.Written >= w.Limit {
		w.Truncated = true
		return len(p), nil
	}
	remaining := w.Limit - w.Written
	if len(p) > remaining {
		w.Truncated = true
		n, err = w.Writer.Write(p[:remaining])
		w.Written += n
		return len(p), err
	}
	n, err = w.Writer.Write(p)
	w.Written += n
	return n, err
}

// Sandbox wraps command execution with resource limits and security checks
type Sandbox struct {
	Limits config.ResourceLimits
}

func NewSandbox(limits config.ResourceLimits) *Sandbox {
	return &Sandbox{
		Limits: limits,
	}
}

// ExecuteCommand executes a command within the sandbox
func (s *Sandbox) ExecuteCommand(ctx context.Context, command string, workingDir string) ([]byte, error) {
	// Apply execution time limit
	timeout := time.Duration(s.Limits.MaxExecutionTime) * time.Second
	if timeout == 0 {
		timeout = 60 * time.Second // Default timeout
	}

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

	// Construct command with limits
	var cmd *exec.Cmd

	if runtime.GOOS == "windows" {
		// Windows: No ulimit support, just basic execution with timeout
		cmd = exec.CommandContext(ctx, "powershell", "-NoProfile", "-NonInteractive", "-Command", command)
	} else {
		// Linux/macOS: Use ulimit for resource limits
		wrappedCommand := s.wrapWithUlimit(command)
		cmd = exec.CommandContext(ctx, "sh", "-c", wrappedCommand)
		
		// Set process group to allow killing the entire tree
		cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
	}

	cmd.Dir = workingDir
	
	// Environment filtering
	if len(s.Limits.AllowedEnvVars) > 0 {
		cmd.Env = s.filterEnv(s.Limits.AllowedEnvVars)
	} else {
		// Default safe environment variables if none specified
		// or inherit based on policy. For now, inherit if empty to avoid breaking changes,
		// but typically we should be restrictive.
		cmd.Env = os.Environ()
	}

	logger.DebugCF("sandbox", "Executing command", map[string]interface{}{
		"command": command,
		"dir":     workingDir,
		"limits":  s.Limits,
	})

	// Output limit
	limit := s.Limits.MaxOutputSize
	if limit <= 0 {
		limit = 10 * 1024 * 1024 // 10MB default
	}

	buf := utils.GetBuffer()
	defer utils.PutBuffer(buf)
	
	limitedWriter := &LimitedWriter{Limit: limit, Writer: buf}

	cmd.Stdout = limitedWriter
	cmd.Stderr = limitedWriter

	err := cmd.Run()
	
	// Copy buffer content to new slice since buffer will be returned to pool
	output := make([]byte, buf.Len())
	copy(output, buf.Bytes())

	if err != nil {
		metrics.SandboxExecs.WithLabelValues("error").Inc()
	} else {
		metrics.SandboxExecs.WithLabelValues("success").Inc()
	}

	if limitedWriter.Truncated {
		output = append(output, []byte(fmt.Sprintf("\n... (truncated, max output size %d bytes exceeded)", limit))...)
	}

	// Check for timeout or other errors
	if ctx.Err() == context.DeadlineExceeded {
		// Kill process group if timed out
		if cmd.Process != nil {
			if runtime.GOOS != "windows" {
				syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL)
			} else {
				cmd.Process.Kill()
			}
		}
		return output, fmt.Errorf("command timed out after %v", timeout)
	}

	return output, err
}

func (s *Sandbox) wrapWithUlimit(command string) string {
	var limits []string

	// CPU time limit (soft limit)
	if s.Limits.MaxExecutionTime > 0 {
		// Add a buffer to the OS limit so context timeout hits first usually
		limits = append(limits, fmt.Sprintf("ulimit -t %d", s.Limits.MaxExecutionTime+5))
	}

	// Memory limit
	if s.Limits.MaxMemory != "" {
		bytes, err := parseMemory(s.Limits.MaxMemory)
		if err == nil && bytes > 0 {
			// ulimit -v takes kbytes
			kbytes := bytes / 1024
			limits = append(limits, fmt.Sprintf("ulimit -v %d", kbytes))
		}
	}

	if len(limits) == 0 {
		return command
	}

	// Chain limits and command
	return fmt.Sprintf("%s; %s", strings.Join(limits, "; "), command)
}

func (s *Sandbox) filterEnv(allowed []string) []string {
	var env []string
	allowedMap := make(map[string]bool)
	for _, k := range allowed {
		allowedMap[k] = true
	}

	for _, e := range os.Environ() {
		pair := strings.SplitN(e, "=", 2)
		if len(pair) > 0 && allowedMap[pair[0]] {
			env = append(env, e)
		}
	}
	return env
}

func parseMemory(s string) (int64, error) {
	re := regexp.MustCompile(`^(\d+)([KMG]B?)?$`)
	matches := re.FindStringSubmatch(strings.ToUpper(s))
	if matches == nil {
		return 0, fmt.Errorf("invalid memory format")
	}

	val, err := strconv.ParseInt(matches[1], 10, 64)
	if err != nil {
		return 0, err
	}

	unit := matches[2]
	switch unit {
	case "K", "KB":
		val *= 1024
	case "M", "MB":
		val *= 1024 * 1024
	case "G", "GB":
		val *= 1024 * 1024 * 1024
	}

	return val, nil
}