rony-llm-agent/pkg/tools/sandbox/sandbox.go

package sandbox

import (
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"github.com/VictorVargas/rony-llm-agent/pkg/llm"
	"github.com/VictorVargas/rony-llm-agent/pkg/tools"
)

// ErrSandboxViolation is returned when a tool call violates sandbox rules.
var ErrSandboxViolation = fmt.Errorf("sandbox violation")

// Sandbox wraps os.Root to enforce filesystem boundaries for tool calls.
type Sandbox struct {
	root *os.Root
	path string
}

// NewSandbox creates a new Sandbox rooted at the given directory.
// The directory is created if it does not exist.
func NewSandbox(rootDir string) (*Sandbox, error) {
	if err := os.MkdirAll(rootDir, 0755); err != nil {
		return nil, fmt.Errorf("creating sandbox root: %w", err)
	}

	absRoot, err := filepath.Abs(rootDir)
	if err != nil {
		return nil, fmt.Errorf("resolving absolute path: %w", err)
	}

	root, err := os.OpenRoot(absRoot)
	if err != nil {
		return nil, fmt.Errorf("opening sandbox root: %w", err)
	}

	return &Sandbox{root: root, path: absRoot}, nil
}

// ValidateToolCall checks if a tool call's arguments reference paths outside the sandbox.
// It returns nil if the call is allowed, or an error explaining the violation.
func (s *Sandbox) ValidateToolCall(tool tools.Tool, call llm.ToolCall) error {
	args := make(map[string]interface{})
	if len(call.Arguments) > 0 {
		if err := json.Unmarshal(call.Arguments, &args); err != nil {
			return fmt.Errorf("parsing arguments: %w", err)
		}
	}

	for _, path := range extractPaths(args) {
		if err := s.validatePath(path); err != nil {
			return fmt.Errorf("%w: %s", ErrSandboxViolation, err)
		}
	}

	return nil
}

// validatePath checks that the given path resolves inside the sandbox root.
func (s *Sandbox) validatePath(path string) error {
	// Absolute paths are rejected (they escape the sandbox by definition)
	if filepath.IsAbs(path) {
		return fmt.Errorf("absolute paths not allowed: %s", path)
	}

	// Resolve to absolute path relative to sandbox root
	abs := filepath.Join(s.path, path)

	// Clean the path to normalize it
	abs = filepath.Clean(abs)

	// Check if the path is inside the root
	if !strings.HasPrefix(abs, s.path+string(filepath.Separator)) && abs != s.path {
		return fmt.Errorf("path escapes sandbox: %s (root: %s)", abs, s.path)
	}

	return nil
}

// extractPaths collects all path-like values from the arguments map.
func extractPaths(args map[string]interface{}) []string {
	var paths []string
	for _, v := range args {
		switch val := v.(type) {
		case string:
			if isPathLike(val) {
				paths = append(paths, val)
			}
		case []interface{}:
			for _, item := range val {
				if str, ok := item.(string); ok && isPathLike(str) {
					paths = append(paths, str)
				}
			}
		}
	}
	return paths
}

// isPathLike checks if a string looks like a filesystem path.
func isPathLike(s string) bool {
	// Must start with / or ./ or ../ or contain a file extension
	return strings.HasPrefix(s, "/") ||
		strings.HasPrefix(s, "./") ||
		strings.HasPrefix(s, "../") ||
		strings.Contains(s, ".") && strings.Contains(s, "/")
}

// Open returns the underlying os.Root for testing.
func (s *Sandbox) Open() *os.Root {
	return s.root
}