ollie-core/pkg/tools/execute/server.go

package execute

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"
	"sync"
	"syscall"
	"time"

	"ollie/internal/sandbox"
	"ollie/pkg/paths"
	"ollie/pkg/tools"
)

const (
	failureWindow = 1 * time.Minute
	maxFailures   = 5
	blockDuration = 5 * time.Minute
)

// Server runs code in a sandboxed environment.
type Server struct {
	// cwd is the working directory for sandboxed commands. If empty,
	// the process working directory is used.
	wdMu sync.RWMutex
	cwd  string

	// envExtra holds per-session environment variables injected via SetEnv.
	envMu    sync.RWMutex
	envExtra map[string]string

	// lockDir is the directory for advisory flock files used during parallel
	// step dispatch. Set once at session init via SetLockDir; empty disables locking.
	lockDir string

	// mountProc is the ollie-9p-mount subprocess for this session's FUSE proxy.
	mountProc    *os.Process
	mountPath    string
	mountEnabled bool

	// Strict rejects inline {code} steps; only {tool} steps are allowed.
	Strict bool

	// Yolo skips the landrun sandbox for all execution.
	Yolo bool

	// allowExecutors restricts which executors (execute_code, call_tool, pipe)
	// are available. Empty means all are allowed.
	allowExecutors map[string]bool

	// allowTools restricts which named tool scripts can be invoked via call_tool/pipe.
	// Empty means all are allowed.
	allowTools map[string]bool

	// rate limiting state (per-Server)
	rateLimitMu        sync.Mutex
	validationFailures int
	lastFailure        time.Time
	blockedUntil       time.Time
}

// Option configures a Server.
type Option func(*Server)

// WithStrict rejects inline {code} steps; only {tool} steps are allowed.
func WithStrict() Option { return func(s *Server) { s.Strict = true } }

// WithYolo skips the landrun sandbox.
func WithYolo() Option { return func(s *Server) { s.Yolo = true } }

// WithMount enables per-session FUSE mount for 9P identity.
func WithMount() Option { return func(s *Server) { s.mountEnabled = true } }

// WithAllowExecutors restricts which executors are available.
func WithAllowExecutors(names []string) Option {
	return func(s *Server) {
		if len(names) > 0 {
			s.allowExecutors = make(map[string]bool, len(names))
			for _, n := range names {
				s.allowExecutors[n] = true
			}
		}
	}
}

// WithAllowTools restricts which tool scripts can be invoked.
func WithAllowTools(names []string) Option {
	return func(s *Server) {
		if len(names) > 0 {
			s.allowTools = make(map[string]bool, len(names))
			for _, n := range names {
				s.allowTools[n] = true
			}
		}
	}
}

// AllowTools returns the set of allowed tool names, or nil if unrestricted.
func (e *Server) AllowTools() []string {
	if len(e.allowTools) == 0 {
		return nil
	}
	out := make([]string, 0, len(e.allowTools))
	for k := range e.allowTools {
		out = append(out, k)
	}
	return out
}

// Decl returns a factory for an execute Server with the given working directory.
func Decl(cwd string, opts ...Option) func() tools.Server {
	return func() tools.Server {
		s := New(cwd)
		for _, o := range opts {
			o(s)
		}
		return s
	}
}

// ListTools implements tools.Server, returning ToolInfo for execute_code,
// call_tool, and pipe.
func (e *Server) ListTools() ([]tools.ToolInfo, error) {
	all := []tools.ToolInfo{
		{
			Name: "execute_code",
			Description: `Run one or more inline code steps in a sandboxed environment.

Steps run in parallel when consecutive steps carry an ollie:parallel annotation;
otherwise they run serially. Outputs are concatenated in submission order.
No stdout chaining between steps — use pipe for that.

Each step is one of:
- {code, language}              — inline code (default language: bash)
- {elevated: true, code}        — run outside sandbox via elevation backend (bash only)
- {parallel: [{code/language}...]} — concurrent fan-out; outputs concatenated in submission order

Supported languages: bash, python3, perl, lua, awk, sed, jq, ed, expect, bc.
timeout applies to each step independently (default: 30s). A failed step aborts.

Examples:
- Single step:   steps=[{code: "date"}]
- Two steps:     steps=[{code: "echo hello"}, {code: "echo world"}]
- Fan-out:       steps=[{parallel: [{code: "cat a.txt"}, {code: "cat b.txt"}]}]`,
			InputSchema: json.RawMessage(`{
				"type": "object",
				"required": ["steps"],
				"properties": {
					"steps": {
						"type": "array",
						"description": "Inline code steps. Run in parallel when safe (ollie:parallel annotation), serially otherwise.",
						"items": {
							"type": "object",
							"properties": {
								"code":     {"type": "string", "description": "Inline code to execute."},
								"language": {"type": "string", "description": "Language interpreter (default: bash)."},
								"elevated": {"type": "boolean", "description": "Run outside the sandbox via the elevation backend. Only bash is supported."},
								"parallel": {
									"type": "array",
									"description": "Fan-out: steps run concurrently, outputs concatenated in submission order.",
									"items": {
										"type": "object",
										"properties": {
											"code":     {"type": "string"},
											"language": {"type": "string"},
											"elevated": {"type": "boolean"}
										}
									}
								}
							}
						}
					},
					"timeout": {"type": "integer", "description": "Timeout in seconds per step (default: 30)."},
					"sandbox": {"type": "string",  "description": "Sandbox name (default: default)."}
				}
			}`),
		},
		{
			Name: "call_tool",
			Description: `Run one or more named scripts from $OLLIE/s/$SID/t (the tools directory).

Tool names and their arguments are discoverable via:
  grep -iA2 'keyword' $OLLIE/s/$SID/t/idx

call_tool fans out in parallel when consecutive calls carry an ollie:parallel read
annotation in their script header; write-annotated or unannotated tools run serially.
Outputs are concatenated in submission order.
No stdout chaining between calls — use pipe for that.

Each call is one of:
- {tool, args}                   — named script with arguments
- {tool, args, elevated: true}   — run outside sandbox (bash tools only)
- {parallel: [{tool, args}...]}  — explicit concurrent fan-out

Examples:
- Single call:    calls=[{tool: "file_read", args: ["README.md"]}]
- Parallel reads: calls=[{tool: "file_read", args: ["a.txt"]}, {tool: "file_read", args: ["b.txt"]}]`,
			InputSchema: json.RawMessage(`{
				"type": "object",
				"required": ["calls"],
				"properties": {
					"calls": {
						"type": "array",
						"description": "Tool calls. Fanned out in parallel when safe (ollie:parallel annotation), serially otherwise.",
						"items": {
							"type": "object",
							"properties": {
								"tool":     {"type": "string", "description": "Named tool script from the tools directory."},
								"args":     {"type": "array", "items": {"type": "string"}, "description": "Arguments for the tool script."},
								"elevated": {"type": "boolean", "description": "Run outside the sandbox via the elevation backend."},
								"parallel": {
									"type": "array",
									"description": "Explicit fan-out: calls run concurrently, outputs concatenated in submission order.",
									"items": {
										"type": "object",
										"properties": {
											"tool":     {"type": "string"},
											"args":     {"type": "array", "items": {"type": "string"}},
											"elevated": {"type": "boolean"}
										}
									}
								}
							}
						}
					},
					"timeout": {"type": "integer", "description": "Timeout in seconds per call (default: 30)."},
					"sandbox": {"type": "string",  "description": "Sandbox name (default: default)."}
				}
			}`),
		},
		{
			Name: "pipe",
			Description: `Compose execute_code and call_tool steps into a sequential pipeline.

Stages run in order; each stage's stdout becomes the next stage's stdin.
Use this when you need to chain heterogeneous code and tool steps together.

Each stage is one of:
- {code, language}              — inline code (default: bash)
- {tool, args}                  — named script from $OLLIE/s/$SID/t
- {elevated: true, code/tool}   — run outside sandbox
- {parallel: [{code/tool}...]}  — fan-out within a stage; outputs concatenated, fed to next stage

Supported inline languages: bash, python3, perl, lua, awk, sed, jq, ed, expect, bc.
timeout applies to each stage independently (default: 30s). A failed stage aborts.

Examples:
- Code → code:    stages=[{code: "grep error app.log"}, {code: "wc -l"}]
- Tool → code:    stages=[{tool: "fetch.sh", args: ["--last=1h"]}, {code: "jq .result"}]
- Fan-out → code: stages=[{parallel: [{code: "cat a.txt"}, {code: "cat b.txt"}]}, {code: "sort"}]`,
			InputSchema: json.RawMessage(`{
				"type": "object",
				"required": ["stages"],
				"properties": {
					"stages": {
						"type": "array",
						"description": "Pipeline stages. Each stage's stdout feeds the next stage's stdin.",
						"items": {
							"type": "object",
							"properties": {
								"code":     {"type": "string", "description": "Inline code to execute."},
								"language": {"type": "string", "description": "Language interpreter (default: bash). Ignored when tool or parallel is set."},
								"tool":     {"type": "string", "description": "Named tool script from the tools directory."},
								"args":     {"type": "array", "items": {"type": "string"}, "description": "Arguments for the tool script."},
								"elevated": {"type": "boolean", "description": "Run outside the sandbox via the elevation backend."},
								"parallel": {
									"type": "array",
									"description": "Fan-out within a stage: steps run concurrently, outputs concatenated, then fed to next stage.",
									"items": {
										"type": "object",
										"properties": {
											"code":     {"type": "string"},
											"language": {"type": "string"},
											"tool":     {"type": "string"},
											"args":     {"type": "array", "items": {"type": "string"}},
											"elevated": {"type": "boolean"}
										}
									}
								}
							}
						}
					},
					"timeout": {"type": "integer", "description": "Timeout in seconds per stage (default: 30)."},
					"sandbox": {"type": "string",  "description": "Sandbox name (default: default)."}
				}
			}`),
		},
	}
	if len(e.allowExecutors) > 0 {
		filtered := all[:0]
		for _, t := range all {
			if e.allowExecutors[t.Name] {
				filtered = append(filtered, t)
			}
		}
		return filtered, nil
	}
	return all, nil
}

// CallTool implements tools.Server.
func (e *Server) CallTool(ctx context.Context, tool string, args json.RawMessage) (json.RawMessage, error) {
	result, err := e.Dispatch(ctx, tool, args)
	if err != nil {
		return json.Marshal(map[string]any{
			"isError": true,
			"content": []map[string]string{{"type": "text", "text": err.Error()}},
		})
	}
	return json.Marshal(map[string]any{
		"content": []map[string]string{{"type": "text", "text": result}},
	})
}

// New creates a new Server with the given working directory.
func New(cwd string) *Server { return &Server{cwd: paths.ExpandHome(cwd)} }

// SetCWD updates the working directory used for subsequent command executions.
func (e *Server) SetCWD(dir string) {
	e.wdMu.Lock()
	e.cwd = paths.ExpandHome(dir)
	e.wdMu.Unlock()
}

// SetAllowExecutors restricts which executors are available.
func (e *Server) SetAllowExecutors(names []string) {
	if len(names) > 0 {
		e.allowExecutors = make(map[string]bool, len(names))
		for _, n := range names {
			e.allowExecutors[n] = true
		}
	} else {
		e.allowExecutors = nil
	}
}

// SetAllowTools restricts which tool scripts can be invoked.
func (e *Server) SetAllowTools(names []string) {
	if len(names) > 0 {
		e.allowTools = make(map[string]bool, len(names))
		for _, n := range names {
			e.allowTools[n] = true
		}
	} else {
		e.allowTools = nil
	}
}

// SetLockDir sets the directory used for advisory flock files during parallel
// step dispatch. Must be called before the Server handles concurrent requests.
func (e *Server) SetLockDir(dir string) { e.lockDir = dir }

// SetEnv adds a session-scoped environment variable injected into all
// subsequent subprocess invocations for this session.
func (e *Server) SetEnv(key, value string) {
	e.envMu.Lock()
	if e.envExtra == nil {
		e.envExtra = make(map[string]string)
	}
	e.envExtra[key] = value
	// Start mount once we have the principal.
	if key == "OLLIE_UNAME" && value != "" {
		e.startMount(value)
	}
	e.envMu.Unlock()
}

// MountPath returns the per-session FUSE mount path for a given uname.
func MountPath(uname string) string {
	return filepath.Join(os.TempDir(), "ollie-mount-"+uname)
}

// startMount launches ollie-9p-mount for this session. Must be called with envMu held.
func (e *Server) startMount(sessionID string) {
	if !e.mountEnabled {
		return
	}
	bin, err := exec.LookPath("ollie-9p-mount")
	if err != nil {
		return
	}
	mnt := MountPath(sessionID)
	os.MkdirAll(mnt, 0755) //nolint:errcheck
	cmd := exec.Command(bin, "-u", sessionID, mnt)
	cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
	if err := cmd.Start(); err != nil {
		return
	}
	// Reap in background so Wait works later.
	go cmd.Wait() //nolint:errcheck
	// Wait for mount to be ready.
	for i := 0; i < 20; i++ {
		time.Sleep(50 * time.Millisecond)
		if _, err := os.Stat(filepath.Join(mnt, "s")); err == nil {
			break
		}
	}
	e.mountProc = cmd.Process
	e.mountPath = mnt
	e.envExtra["OLLIE"] = mnt
}

// remountIfStale checks the per-session FUSE mount health and re-establishes it if broken.
// Must be called without envMu held.
func (e *Server) remountIfStale() {
	e.envMu.RLock()
	mnt := e.mountPath
	proc := e.mountProc
	e.envMu.RUnlock()
	if mnt == "" || proc == nil {
		return
	}
	if _, err := os.Stat(filepath.Join(mnt, "s")); err == nil {
		return
	}
	// Mount is stale — tear down and restart.
	e.envMu.Lock()
	defer e.envMu.Unlock()
	exec.Command("fusermount", "-u", e.mountPath).Run() //nolint:errcheck
	syscall.Kill(e.mountProc.Pid, syscall.SIGKILL)      //nolint:errcheck
	e.mountProc = nil
	e.mountPath = ""
	if uname := e.envExtra["OLLIE_UNAME"]; uname != "" {
		e.startMount(uname)
	}
}

// Close tears down the per-session FUSE mount.
func (e *Server) Close() {
	e.envMu.Lock()
	proc := e.mountProc
	mnt := e.mountPath
	e.mountProc = nil
	e.mountPath = ""
	e.envMu.Unlock()
	if proc != nil {
		ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
		exec.CommandContext(ctx, "fusermount", "-u", mnt).Run() //nolint:errcheck
		cancel()
		syscall.Kill(proc.Pid, syscall.SIGKILL) //nolint:errcheck
		os.Remove(mnt)                          //nolint:errcheck
	}
}

// executeElevated runs cmd outside the sandbox via x/elevate.
// Returns (output, error); a non-zero exit code is treated as an error.
func (e *Server) executeElevated(ctx context.Context, cmd, dir string, timeout int) (string, error) {
	script := filepath.Join(PluginsPath(), "elevate")
	if _, err := os.Stat(script); err != nil {
		return "", fmt.Errorf("elevation not available: x/elevate not found")
	}
	ctx, cancel := context.WithTimeout(ctx, time.Duration(timeout)*time.Second)
	defer cancel()

	c := exec.CommandContext(ctx, script, "--", cmd)
	c.Dir = dir
	var outBuf, errBuf bytes.Buffer
	lw := &limitedWriter{w: &outBuf, limit: 10 * 1024 * 1024}
	c.Stdout = lw
	c.Stderr = &errBuf
	err := c.Run()
	combined := outBuf.String()
	if errBuf.Len() > 0 {
		combined += errBuf.String()
	}
	if lw.truncated {
		combined += "\n[output truncated at 10MB]"
	}
	if ctx.Err() == context.DeadlineExceeded {
		return combined, fmt.Errorf("execution timeout after %d seconds", timeout)
	}
	if err != nil {
		var exitErr *exec.ExitError
		if errors.As(err, &exitErr) {
			return combined, fmt.Errorf("elevated execution failed (exit %d)\nOutput: %s", exitErr.ExitCode(), combined)
		}
		return combined, fmt.Errorf("elevated execution failed: %v", err)
	}
	return combined, nil
}

var whitespacePattern = regexp.MustCompile(`\s+`)

func (e *Server) checkRateLimit() error {
	e.rateLimitMu.Lock()
	defer e.rateLimitMu.Unlock()

	now := time.Now()
	if now.Before(e.blockedUntil) {
		remaining := e.blockedUntil.Sub(now).Round(time.Second)
		return fmt.Errorf("rate limited: too many validation failures, blocked for %v", remaining)
	}
	return nil
}

func (e *Server) recordValidationFailure() {
	e.rateLimitMu.Lock()
	defer e.rateLimitMu.Unlock()

	now := time.Now()
	if now.Sub(e.lastFailure) > failureWindow {
		e.validationFailures = 0
	}

	e.validationFailures++
	e.lastFailure = now

	if e.validationFailures >= maxFailures {
		e.blockedUntil = now.Add(blockDuration)
		e.validationFailures = 0
	}
}

// Execute runs code in a sandbox and returns combined stdout+stderr.
func (e *Server) Execute(ctx context.Context, code, language string, timeout int, sandboxName string, trusted bool) (string, error) {
	return e.executeWithStdin(ctx, code, language, timeout, sandboxName, trusted, "")
}

// executeWithStdin is like Execute but feeds stdinData to the command's stdin.
// For languages where code is itself passed via stdin (ed, expect, bc), stdinData is ignored.
func (e *Server) executeWithStdin(ctx context.Context, code, language string, timeout int, sandboxName string, trusted bool, stdinData string) (string, error) {
	if timeout <= 0 {
		timeout = 30
	}

	// Domain-specific tools (awk, sed, jq, ed, bc, expect) skip validation —
	// they are constrained by design and the sandbox handles OS-level restrictions.
	// General-purpose languages get universal + language-specific pattern checks.
	_, isDomainSpecific := map[string]struct{}{
		"awk": {}, "sed": {}, "jq": {}, "ed": {}, "bc": {}, "expect": {},
	}[language]
	if !trusted && !isDomainSpecific {
		if err := e.ValidateCode(code, language); err != nil {
			return "", err
		}
	}

	cfg, err := loadSandboxConfig(sandboxName)
	if err != nil {
		return "", err
	}

	e.wdMu.RLock()
	workDir := e.cwd
	e.wdMu.RUnlock()
	if workDir == "" {
		workDir, _ = os.Getwd()
	}

	ctx, cancel := context.WithTimeout(ctx, time.Duration(timeout)*time.Second)
	defer cancel()

	var cmd *exec.Cmd
	var interpreter []string
	// codeStdin: non-empty means the code itself is fed via stdin (ed, expect, bc).
	// In these cases stdinData cannot be used simultaneously.
	var codeStdin string
	switch language {
	case "bash", "":
		interpreter = []string{"bash", "-c", code}
	case "python3", "python":
		interpreter = []string{"python3", "-c", code}
	case "perl":
		interpreter = []string{"perl", "-e", code}
	case "lua":
		interpreter = []string{"lua", "-e", code}
	case "awk":
		interpreter = []string{"gawk", code}
	case "sed":
		interpreter = []string{"sed", "-e", code}
	case "jq":
		interpreter = []string{"jq", code}
	case "ed":
		interpreter = []string{"ed", "-s"}
		codeStdin = code
	case "expect":
		interpreter = []string{"expect", "-"}
		codeStdin = code
	case "bc":
		interpreter = []string{"bc", "-l"}
		codeStdin = code
	default:
		return "", fmt.Errorf("unsupported language: %s (supported: bash, python3, perl, awk, sed, ed, jq, expect, bc, lua)", language)
	}
	e.envMu.RLock()
	envMap := make(map[string]string, len(e.envExtra))
	for k, v := range e.envExtra {
		envMap[k] = v
	}
	e.envMu.RUnlock()
	for _, ev := range os.Environ() {
		k, v, _ := strings.Cut(ev, "=")
		if _, ok := envMap[k]; !ok {
			envMap[k] = v
		}
	}
	if _, ok := envMap["NAMESPACE"]; !ok {
		if ns := plan9Namespace(envMap); ns != "" {
			envMap["NAMESPACE"] = ns
		}
	}
	getenv := func(key string) string { return envMap[key] }

	if e.Yolo {
		cmd = exec.CommandContext(ctx, interpreter[0], interpreter[1:]...)
	} else {
		wrapped, wrapErr := sandbox.WrapCommand(cfg, interpreter, workDir, getenv)
		if wrapErr != nil {
			return "", wrapErr
		}
		cmd = exec.CommandContext(ctx, wrapped[0], wrapped[1:]...)
	}
	cmd.Dir = workDir
	switch {
	case codeStdin != "":
		cmd.Stdin = strings.NewReader(codeStdin)
	case stdinData != "":
		cmd.Stdin = strings.NewReader(stdinData)
	}

	e.envMu.RLock()
	baseEnv := os.Environ()
	// Remove keys that envExtra overrides so duplicates don't leak through.
	filtered := baseEnv[:0]
	for _, ev := range baseEnv {
		k, _, _ := strings.Cut(ev, "=")
		if _, overridden := e.envExtra[k]; !overridden {
			filtered = append(filtered, ev)
		}
	}
	olliePath := e.envExtra["OLLIE"]
	if olliePath == "" {
		olliePath = os.Getenv("OLLIE")
	}
	cmd.Env = prependOlliePath(filtered, olliePath)
	cmd.Env = append(cmd.Env, "OLLIE_TOOLS_PATH="+ToolsPath())
	for k, v := range e.envExtra {
		cmd.Env = append(cmd.Env, k+"="+v)
	}
	e.envMu.RUnlock()

	cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
	cmd.Cancel = func() error {
		if cmd.Process != nil {
			syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL)
		}
		return nil
	}

	var outputBuf bytes.Buffer
	lw := &limitedWriter{w: &outputBuf, limit: 10 * 1024 * 1024}
	cmd.Stdout = lw
	cmd.Stderr = lw

	err = cmd.Run()
	output := outputBuf.Bytes()
	if lw.truncated {
		output = append(output, []byte("\n[output truncated at 10MB]")...)
	}

	if ctx.Err() == context.DeadlineExceeded {
		return "", fmt.Errorf("execution timeout after %d seconds", timeout)
	}
	if err != nil {
		return string(output), fmt.Errorf("execution failed: %v\nOutput: %s", err, string(output))
	}
	return string(output), nil
}

// prependOlliePath returns env with $OLLIE/x/ prepended to PATH so wrapper
// scripts placed there (e.g. ~/.config/ollie/scripts/x/) shadow system binaries.
func prependOlliePath(env []string, ollie string) []string {
	if ollie == "" {
		return env
	}
	xdir := filepath.Join(ollie, "x")
	result := make([]string, 0, len(env))
	for _, e := range env {
		if strings.HasPrefix(e, "PATH=") {
			e = "PATH=" + xdir + string(filepath.ListSeparator) + e[5:]
		}
		result = append(result, e)
	}
	return result
}

// plan9Namespace computes the plan9port namespace directory from the environment.
// Convention: $NAMESPACE if set, else /tmp/ns.<user>.<display> where display is
// $DISPLAY with trailing .0 stripped and / replaced by _.
func plan9Namespace(env map[string]string) string {
	if ns := env["NAMESPACE"]; ns != "" {
		return ns
	}
	user := env["USER"]
	if user == "" {
		return ""
	}
	disp := env["DISPLAY"]
	if disp == "" {
		return ""
	}
	// Canonicalize: strip trailing .0
	if strings.HasSuffix(disp, ".0") {
		disp = disp[:len(disp)-2]
	}
	// Replace / with _
	disp = strings.ReplaceAll(disp, "/", "_")
	return "/tmp/ns." + user + "." + disp
}

var _ tools.Server = (*Server)(nil) // compile-time interface check