887 lines
27 KiB
Go
887 lines
27 KiB
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
|
|
|
|
// Hooks for lifecycle events (harnesses like 9P can inject mount logic here)
|
|
OnPreDispatch func()
|
|
OnClose func()
|
|
OnEnvSet func(key, value string)
|
|
|
|
// 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
|
|
|
|
// Detached process management
|
|
detachMu sync.Mutex
|
|
detachCh chan struct{} // signal to detach the currently running process
|
|
detached []*DetachedProcess
|
|
OnDetach func(pid int, cmd string) // hook: called when a process is detached
|
|
OnExit func(pid int, exitCode int) // hook: called when a detached process exits
|
|
}
|
|
|
|
// 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 } }
|
|
|
|
// WithOnPreDispatch registers a hook called before each Dispatch.
|
|
func WithOnPreDispatch(fn func()) Option { return func(s *Server) { s.OnPreDispatch = fn } }
|
|
|
|
// WithOnClose registers a hook called during Close.
|
|
func WithOnClose(fn func()) Option { return func(s *Server) { s.OnClose = fn } }
|
|
|
|
// SetOnExit sets the hook called when a detached process exits.
|
|
func (e *Server) SetOnExit(fn func(pid, exitCode int)) { e.OnExit = fn }
|
|
|
|
// WithOnEnvSet registers a hook called each time SetEnv is called.
|
|
func WithOnEnvSet(fn func(key, value string)) Option { return func(s *Server) { s.OnEnvSet = fn } }
|
|
|
|
// 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)
|
|
- {detach: true, code} — start process and immediately background it; returns PID
|
|
- {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."},
|
|
"detach": {"type": "boolean", "description": "Start the process and immediately background it. Returns the PID. Use process_list/process_output to inspect."},
|
|
"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). Use 0 for no timeout."},
|
|
"sandbox": {"type": "string", "description": "Sandbox name (default: default)."}
|
|
}
|
|
}`),
|
|
},
|
|
{
|
|
Name: "call_tool",
|
|
Description: `Run one or more named scripts from the tools directory.
|
|
|
|
Tool names and their arguments are discoverable via:
|
|
grep -iA2 'keyword' $OLLIE_TOOLS_PATH/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). Use 0 for no timeout."},
|
|
"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 the tools directory
|
|
- {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). Use 0 for no timeout."},
|
|
"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()}},
|
|
})
|
|
}
|
|
// If the tool output is already structured content-block JSON, pass it through.
|
|
var structured map[string]json.RawMessage
|
|
if json.Unmarshal([]byte(result), &structured) == nil {
|
|
if _, ok := structured["content"]; ok {
|
|
return []byte(result), nil
|
|
}
|
|
}
|
|
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
|
|
e.envMu.Unlock()
|
|
if e.OnEnvSet != nil {
|
|
e.OnEnvSet(key, value)
|
|
}
|
|
}
|
|
// Close is called when the session ends. Calls OnClose hook if registered.
|
|
func (e *Server) Close() {
|
|
e.cleanupDetached()
|
|
if e.OnClose != nil {
|
|
e.OnClose()
|
|
}
|
|
}
|
|
|
|
// 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
|
|
c.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
|
|
c.Cancel = func() error {
|
|
if c.Process != nil {
|
|
syscall.Kill(-c.Process.Pid, syscall.SIGKILL)
|
|
}
|
|
return nil
|
|
}
|
|
c.WaitDelay = time.Second
|
|
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, detach ...bool) (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()
|
|
}
|
|
|
|
var cancel context.CancelFunc
|
|
if timeout > 0 {
|
|
ctx, cancel = context.WithTimeout(ctx, time.Duration(timeout)*time.Second)
|
|
} else {
|
|
// timeout=0 means no timeout (run indefinitely, e.g. for daemons)
|
|
ctx, cancel = context.WithCancel(ctx)
|
|
}
|
|
// NOTE: do NOT defer cancel() here — the detach path must avoid cancelling
|
|
// the context (which would kill the detached process via cmd.Cancel).
|
|
// Each exit path calls cancel() explicitly except detach.
|
|
|
|
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)
|
|
}
|
|
}
|
|
cmd.Env = prependOlliePath(filtered, paths.CfgDir())
|
|
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
|
|
}
|
|
cmd.WaitDelay = time.Second
|
|
|
|
var outputBuf bytes.Buffer
|
|
lw := &limitedWriter{w: &outputBuf, limit: 10 * 1024 * 1024}
|
|
cmd.Stdout = lw
|
|
cmd.Stderr = lw
|
|
|
|
// Set up detach channel for this execution
|
|
detachCh := make(chan struct{}, 1)
|
|
e.detachMu.Lock()
|
|
e.detachCh = detachCh
|
|
e.detachMu.Unlock()
|
|
defer func() {
|
|
e.detachMu.Lock()
|
|
e.detachCh = nil
|
|
e.detachMu.Unlock()
|
|
}()
|
|
|
|
if err = cmd.Start(); err != nil {
|
|
return "", fmt.Errorf("execution failed: %v", err)
|
|
}
|
|
|
|
// If detach requested, immediately signal the detach channel
|
|
if len(detach) > 0 && detach[0] {
|
|
close(detachCh)
|
|
}
|
|
|
|
// Wait for completion, context cancellation, or detach signal
|
|
waitCh := make(chan error, 1)
|
|
go func() { waitCh <- cmd.Wait() }()
|
|
|
|
select {
|
|
case err = <-waitCh:
|
|
// Normal completion
|
|
cancel()
|
|
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
|
|
|
|
case <-ctx.Done():
|
|
// Context cancelled (interrupt or timeout) — cancel() already fired via ctx
|
|
syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL)
|
|
<-waitCh // reap
|
|
output := outputBuf.Bytes()
|
|
if ctx.Err() == context.DeadlineExceeded {
|
|
return string(output), fmt.Errorf("execution timeout after %d seconds", timeout)
|
|
}
|
|
return string(output), fmt.Errorf("execution interrupted")
|
|
|
|
case <-detachCh:
|
|
// Detach: move process to registry, continue running
|
|
// Neutralize cmd.Cancel and WaitDelay so context cancellation
|
|
// won't kill the process after the timeout fires.
|
|
cmd.Cancel = nil
|
|
cmd.WaitDelay = 0
|
|
cancel()
|
|
|
|
cmdStr := strings.Join(interpreter, " ")
|
|
if len(cmdStr) > 80 {
|
|
cmdStr = cmdStr[:77] + "..."
|
|
}
|
|
ring := newRingBuffer(ringBufSize)
|
|
// Splice: future output goes to ring buffer instead of outputBuf.
|
|
lw.mu.Lock()
|
|
lw.w = ring
|
|
lw.mu.Unlock()
|
|
|
|
proc := &DetachedProcess{
|
|
PID: cmd.Process.Pid,
|
|
Command: cmdStr,
|
|
Started: time.Now(),
|
|
ring: ring,
|
|
cmd: cmd.Process,
|
|
done: make(chan struct{}),
|
|
}
|
|
e.detachMu.Lock()
|
|
e.detached = append(e.detached, proc)
|
|
e.detachMu.Unlock()
|
|
|
|
// Monitor for exit in background
|
|
go func() {
|
|
waitErr := <-waitCh
|
|
proc.mu.Lock()
|
|
proc.Exited = true
|
|
if waitErr != nil {
|
|
if exitErr, ok := waitErr.(*exec.ExitError); ok {
|
|
proc.ExitCode = exitErr.ExitCode()
|
|
} else {
|
|
proc.ExitCode = -1
|
|
}
|
|
}
|
|
proc.mu.Unlock()
|
|
close(proc.done)
|
|
if e.OnExit != nil {
|
|
e.OnExit(proc.PID, proc.ExitCode)
|
|
}
|
|
}()
|
|
|
|
if e.OnDetach != nil {
|
|
e.OnDetach(proc.PID, cmdStr)
|
|
}
|
|
|
|
partial := outputBuf.String()
|
|
return partial + fmt.Sprintf("\n[detached: pid %d]", proc.PID), nil
|
|
}
|
|
}
|
|
|
|
// prependOlliePath returns env with $OLLIE_CFG_PATH/scripts/x/ prepended to PATH so wrapper
|
|
// scripts placed there shadow system binaries.
|
|
func prependOlliePath(env []string, cfgDir string) []string {
|
|
if cfgDir == "" {
|
|
return env
|
|
}
|
|
xdir := filepath.Join(cfgDir, "scripts", "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
|
|
}
|
|
|
|
// Detach signals the currently running process to be detached from the agent.
|
|
// The process continues running; its output is captured in a ring buffer.
|
|
// Returns false if no process is currently running.
|
|
func (e *Server) Detach() bool {
|
|
e.detachMu.Lock()
|
|
ch := e.detachCh
|
|
e.detachMu.Unlock()
|
|
if ch == nil {
|
|
return false
|
|
}
|
|
select {
|
|
case ch <- struct{}{}:
|
|
return true
|
|
default:
|
|
return false
|
|
}
|
|
}
|
|
|
|
// ListDetached returns all detached processes (running and exited).
|
|
func (e *Server) ListDetached() []*DetachedProcess {
|
|
e.detachMu.Lock()
|
|
defer e.detachMu.Unlock()
|
|
out := make([]*DetachedProcess, len(e.detached))
|
|
copy(out, e.detached)
|
|
return out
|
|
}
|
|
|
|
// ListDetachedInfo returns plain-data snapshots of all detached processes.
|
|
func (e *Server) ListDetachedInfo() []DetachedInfoData {
|
|
e.detachMu.Lock()
|
|
defer e.detachMu.Unlock()
|
|
out := make([]DetachedInfoData, len(e.detached))
|
|
for i, p := range e.detached {
|
|
out[i] = p.Info()
|
|
}
|
|
return out
|
|
}
|
|
|
|
// ListDetachedRaw returns detached process info as []any (each element is map[string]any)
|
|
// for consumption by packages that can't import this package directly.
|
|
func (e *Server) ListDetachedRaw() []any {
|
|
e.detachMu.Lock()
|
|
defer e.detachMu.Unlock()
|
|
out := make([]any, len(e.detached))
|
|
for i, p := range e.detached {
|
|
info := p.Info()
|
|
out[i] = map[string]any{
|
|
"pid": info.PID,
|
|
"command": info.Command,
|
|
"started": info.Started,
|
|
"exited": info.Exited,
|
|
"exit_code": info.ExitCode,
|
|
}
|
|
}
|
|
return out
|
|
}
|
|
|
|
// SignalDetached sends a signal to a detached process by PID.
|
|
func (e *Server) SignalDetached(pid int, sig syscall.Signal) error {
|
|
e.detachMu.Lock()
|
|
defer e.detachMu.Unlock()
|
|
for _, p := range e.detached {
|
|
if p.PID == pid {
|
|
return p.Signal(sig)
|
|
}
|
|
}
|
|
return fmt.Errorf("no detached process with pid %d", pid)
|
|
}
|
|
|
|
// GetDetachedOutput returns the ring buffer contents for a detached process.
|
|
func (e *Server) GetDetachedOutput(pid int) (string, error) {
|
|
e.detachMu.Lock()
|
|
defer e.detachMu.Unlock()
|
|
for _, p := range e.detached {
|
|
if p.PID == pid {
|
|
return p.Output(), nil
|
|
}
|
|
}
|
|
return "", fmt.Errorf("no detached process with pid %d", pid)
|
|
}
|
|
|
|
// DismissDetached removes an exited process from the list.
|
|
func (e *Server) DismissDetached(pid int) bool {
|
|
e.detachMu.Lock()
|
|
defer e.detachMu.Unlock()
|
|
for i, p := range e.detached {
|
|
if p.PID == pid && p.Exited {
|
|
e.detached = append(e.detached[:i], e.detached[i+1:]...)
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
// cleanupDetached sends SIGTERM to all running detached processes.
|
|
func (e *Server) cleanupDetached() {
|
|
e.detachMu.Lock()
|
|
defer e.detachMu.Unlock()
|
|
for _, p := range e.detached {
|
|
p.mu.Lock()
|
|
if !p.Exited && p.cmd != nil {
|
|
syscall.Kill(-p.PID, syscall.SIGTERM)
|
|
}
|
|
p.mu.Unlock()
|
|
}
|
|
}
|
|
|
|
var _ tools.Server = (*Server)(nil) // compile-time interface check
|