1029 lines
27 KiB
Go
1029 lines
27 KiB
Go
package execute
|
|
|
|
import (
|
|
"context"
|
|
"encoding/binary"
|
|
"bytes"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"net"
|
|
"os"
|
|
"os/exec"
|
|
"path/filepath"
|
|
"regexp"
|
|
"strings"
|
|
"sync"
|
|
"syscall"
|
|
"time"
|
|
|
|
"ollie/internal/sandbox"
|
|
"ollie/pkg/paths"
|
|
"ollie/pkg/tools"
|
|
"ollie/pkg/skills"
|
|
)
|
|
|
|
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
|
|
|
|
// 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
|
|
|
|
// allowTools restricts which named tool scripts can be invoked via call_tool/pipe.
|
|
// Empty means all are allowed.
|
|
allowTools map[string]bool
|
|
|
|
toolRegistry *Registry
|
|
skillsRegistry *skills.Registry
|
|
sessionID string
|
|
|
|
// OnInjection is called when a skill is loaded and its content
|
|
// should be injected into the agent's context.
|
|
OnInjection func(content string)
|
|
|
|
// 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 } }
|
|
|
|
// 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
|
|
}
|
|
|
|
|
|
// WithToolRegistry attaches a tool registry and session ID to the Server.
|
|
func WithToolRegistry(r *Registry, sessionID string) Option {
|
|
return func(s *Server) {
|
|
s.toolRegistry = r
|
|
s.sessionID = sessionID
|
|
}
|
|
}
|
|
|
|
// WithSkillsRegistry attaches a skills registry to the Server.
|
|
func WithSkillsRegistry(r *skills.Registry) Option {
|
|
return func(s *Server) { s.skillsRegistry = r }
|
|
}
|
|
|
|
// 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 shell plus any
|
|
// tools promoted in the session's tool registry.
|
|
func (e *Server) ListTools() ([]tools.ToolInfo, error) {
|
|
all := []tools.ToolInfo{
|
|
{
|
|
Name: "shell",
|
|
Description: `Execute a single bash command in a sandboxed environment.
|
|
|
|
Usage: {"cmd": "your command here"}
|
|
|
|
Sandbox prevents dangerous operations. Use for computation, builds, scripting.
|
|
timeout applies to each call (default: 30s). A non-zero exit is an error.`,
|
|
InputSchema: json.RawMessage(`{
|
|
"type": "object",
|
|
"required": ["cmd"],
|
|
"properties": {
|
|
"cmd": {"type": "string", "description": "Bash command to execute."},
|
|
"timeout": {"type": "integer", "description": "Timeout in seconds (default: 30). Use 0 for no timeout."},
|
|
"sandbox": {"type": "string", "description": "Sandbox profile name (default: default)."},
|
|
"elevated": {"type": "boolean", "description": "Run outside the sandbox via elevation broker."}
|
|
}
|
|
}`),
|
|
},
|
|
{
|
|
Name: "tool_list",
|
|
Description: `List all available tools that can be loaded.
|
|
|
|
Usage: {"name": "toolname"} — if name is provided, loads that tool.
|
|
Otherwise lists all tools with descriptions.`,
|
|
InputSchema: json.RawMessage(`{
|
|
"type": "object",
|
|
"properties": {}
|
|
}`),
|
|
},
|
|
{
|
|
Name: "tool_load",
|
|
Description: `Load a tool by name into the current session.
|
|
|
|
Usage: {"name": "toolname"}
|
|
|
|
After loading, the tool becomes a native callable function.`,
|
|
InputSchema: json.RawMessage(`{
|
|
"type": "object",
|
|
"required": ["name"],
|
|
"properties": {
|
|
"name": {"type": "string", "description": "Tool name to load."}
|
|
}
|
|
}`),
|
|
},
|
|
{
|
|
Name: "tool_active",
|
|
Description: `List tools currently loaded (promoted) in this session.
|
|
|
|
Usage: (no arguments)
|
|
|
|
Returns tools with descriptions, one per line.`,
|
|
InputSchema: json.RawMessage(`{
|
|
"type": "object",
|
|
"properties": {}
|
|
}`),
|
|
},
|
|
}
|
|
|
|
if e.toolRegistry != nil && e.sessionID != "" {
|
|
all = append(all, e.toolRegistry.Loaded(e.sessionID)...)
|
|
}
|
|
// Append skill tools if skills registry is available
|
|
skillTools := ListSkillsTools(e.skillsRegistry, e.sessionID)
|
|
all = append(all, skillTools...)
|
|
return all, nil
|
|
}
|
|
|
|
// CallTool implements tools.Server.
|
|
func (e *Server) CallTool(ctx context.Context, tool string, args json.RawMessage) (json.RawMessage, error) {
|
|
if e.toolRegistry != nil && e.sessionID != "" {
|
|
if _, promoted := e.toolRegistry.Lookup(e.sessionID, tool); promoted {
|
|
return e.callPromotedTool(ctx, tool, args)
|
|
}
|
|
}
|
|
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()
|
|
}
|
|
|
|
// 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
|
|
}
|
|
}
|
|
|
|
// 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)
|
|
}
|
|
}
|
|
|
|
// SetToolRegistry attaches a session-local tool registry.
|
|
func (e *Server) SetToolRegistry(r *Registry, sessionID string) {
|
|
e.toolRegistry = r
|
|
e.sessionID = sessionID
|
|
}
|
|
|
|
// callPromotedTool executes a tool promoted via the registry by running the
|
|
// script file inside the sandbox, piping the JSON args to stdin.
|
|
func (e *Server) callPromotedTool(ctx context.Context, tool string, args json.RawMessage) (json.RawMessage, error) {
|
|
// Resolve script path.
|
|
if strings.Contains(tool, "/") || strings.Contains(tool, "..") {
|
|
return nil, fmt.Errorf("invalid tool name")
|
|
}
|
|
path := filepath.Join(ToolsPath(), tool)
|
|
|
|
// Extract elevated flag (dispatch-level concern, not passed to tool).
|
|
elevated := false
|
|
var argMap map[string]interface{}
|
|
if err := json.Unmarshal(args, &argMap); err == nil {
|
|
if e, ok := argMap["elevated"]; ok {
|
|
switch v := e.(type) {
|
|
case bool:
|
|
elevated = v
|
|
case string:
|
|
elevated = v == "true" || v == "1"
|
|
}
|
|
// Remove elevated from the args passed to the tool.
|
|
delete(argMap, "elevated")
|
|
args, _ = json.Marshal(argMap)
|
|
}
|
|
}
|
|
|
|
// The tool script receives JSON args on stdin.
|
|
code := path
|
|
stdinData := string(args)
|
|
|
|
var result string
|
|
var err error
|
|
if elevated {
|
|
e.wdMu.RLock()
|
|
workDir := e.cwd
|
|
e.wdMu.RUnlock()
|
|
// Broker protocol has no stdin support; pipe JSON via heredoc.
|
|
elevatedCode := fmt.Sprintf("cat <<'OLLIE_EOF' | %s\n%s\nOLLIE_EOF", code, stdinData)
|
|
result, err = e.executeElevated(ctx, elevatedCode, workDir, 30)
|
|
} else {
|
|
result, err = e.executeWithStdin(ctx, code, "bash", 30, "default", false, stdinData)
|
|
}
|
|
if err != nil {
|
|
return json.Marshal(map[string]interface{}{
|
|
"isError": true,
|
|
"content": []map[string]string{{"type": "text", "text": result + ": " + err.Error()}},
|
|
})
|
|
}
|
|
return json.Marshal(map[string]interface{}{
|
|
"content": []map[string]string{{"type": "text", "text": result}},
|
|
})
|
|
}
|
|
|
|
|
|
|
|
|
|
// 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 the integrated elevation broker.
|
|
// Connects to the broker socket, sends the request with the current env,
|
|
// and streams the framed response back.
|
|
func (e *Server) executeElevated(ctx context.Context, cmd, dir string, timeout int, detach ...bool) (string, error) {
|
|
sockPath := os.Getenv("OLLIE_ELEVATE_SOCKET")
|
|
if sockPath == "" {
|
|
xdg := os.Getenv("XDG_RUNTIME_DIR")
|
|
if xdg == "" {
|
|
return "", fmt.Errorf("elevation not available: no XDG_RUNTIME_DIR")
|
|
}
|
|
sockPath = filepath.Join(xdg, "ollie", "elevate.sock")
|
|
}
|
|
|
|
wantDetach := len(detach) > 0 && detach[0]
|
|
|
|
var cancel context.CancelFunc
|
|
if wantDetach || timeout <= 0 {
|
|
ctx, cancel = context.WithCancel(ctx)
|
|
} else {
|
|
ctx, cancel = context.WithTimeout(ctx, time.Duration(timeout)*time.Second)
|
|
}
|
|
|
|
// Connect to broker
|
|
conn, err := net.DialTimeout("unix", sockPath, 5*time.Second)
|
|
if err != nil {
|
|
cancel()
|
|
return "", fmt.Errorf("elevation not available: %w", err)
|
|
}
|
|
|
|
// Send request
|
|
envMap := make(map[string]string, len(os.Environ()))
|
|
for _, kv := range os.Environ() {
|
|
if k, v, ok := strings.Cut(kv, "="); ok {
|
|
envMap[k] = v
|
|
}
|
|
}
|
|
reqJSON, _ := json.Marshal(struct {
|
|
Cmd string `json:"cmd"`
|
|
Cwd string `json:"cwd"`
|
|
Env map[string]string `json:"env"`
|
|
Session string `json:"session,omitempty"`
|
|
}{Cmd: cmd, Cwd: dir, Env: envMap, Session: os.Getenv("OLLIE_SESSION_ID")})
|
|
reqJSON = append(reqJSON, '\n')
|
|
if _, err := conn.Write(reqJSON); err != nil {
|
|
conn.Close()
|
|
cancel()
|
|
return "", fmt.Errorf("elevated execution failed: write: %w", err)
|
|
}
|
|
|
|
// Set up detach channel
|
|
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 wantDetach {
|
|
close(detachCh)
|
|
}
|
|
|
|
// readFrames reads from the connection, writing to w and streaming.
|
|
// Returns exit code when 'x' frame arrives or -1 on error.
|
|
readFrames := func(w io.Writer, stream func(string)) int {
|
|
header := make([]byte, 5)
|
|
for {
|
|
conn.SetReadDeadline(time.Now().Add(1 * time.Second)) //nolint:errcheck
|
|
_, err := io.ReadFull(conn, header)
|
|
if err != nil {
|
|
if ne, ok := err.(net.Error); ok && ne.Timeout() {
|
|
// Check for context cancellation
|
|
select {
|
|
case <-ctx.Done():
|
|
return -1
|
|
default:
|
|
continue
|
|
}
|
|
}
|
|
return -1
|
|
}
|
|
|
|
frameType := header[0]
|
|
length := binary.BigEndian.Uint32(header[1:5])
|
|
|
|
payload := make([]byte, length)
|
|
if length > 0 {
|
|
conn.SetReadDeadline(time.Now().Add(30 * time.Second)) //nolint:errcheck
|
|
if _, err := io.ReadFull(conn, payload); err != nil {
|
|
return -1
|
|
}
|
|
}
|
|
|
|
switch frameType {
|
|
case 'd':
|
|
w.Write(payload) //nolint:errcheck
|
|
if stream != nil {
|
|
stream(string(payload))
|
|
}
|
|
case 'x':
|
|
var code int
|
|
fmt.Sscanf(string(payload), "%d", &code)
|
|
return code
|
|
}
|
|
}
|
|
}
|
|
|
|
// Check if detach was requested
|
|
select {
|
|
case <-detachCh:
|
|
// Detach: background the socket read into a goroutine
|
|
cmdStr := "elevated: " + cmd
|
|
if len(cmdStr) > 80 {
|
|
cmdStr = cmdStr[:77] + "..."
|
|
}
|
|
ring := newRingBuffer(ringBufSize)
|
|
pid := int(time.Now().UnixNano() & 0x7FFFFFFF) // synthetic PID
|
|
|
|
proc := &DetachedProcess{
|
|
PID: pid,
|
|
Command: cmdStr,
|
|
Started: time.Now(),
|
|
ring: ring,
|
|
done: make(chan struct{}),
|
|
}
|
|
e.detachMu.Lock()
|
|
e.detached = append(e.detached, proc)
|
|
e.detachMu.Unlock()
|
|
|
|
go func() {
|
|
defer conn.Close()
|
|
defer cancel()
|
|
exitCode := readFrames(ring, nil)
|
|
proc.mu.Lock()
|
|
proc.Exited = true
|
|
proc.ExitCode = exitCode
|
|
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)
|
|
}
|
|
return fmt.Sprintf("[detached: pid %d]", pid), nil
|
|
|
|
default:
|
|
// Normal (foreground) execution — must also handle manual detach.
|
|
var outputBuf bytes.Buffer
|
|
streamFn := tools.StreamFunc(ctx)
|
|
lw := &limitedWriter{
|
|
w: &outputBuf,
|
|
limit: 10 * 1024 * 1024,
|
|
stream: streamFn,
|
|
}
|
|
|
|
// Run readFrames in a goroutine so we can select on detach.
|
|
type frameResult struct{ exitCode int }
|
|
frameCh := make(chan frameResult, 1)
|
|
go func() {
|
|
code := readFrames(lw, nil)
|
|
frameCh <- frameResult{code}
|
|
}()
|
|
|
|
select {
|
|
case fr := <-frameCh:
|
|
// Normal completion
|
|
conn.Close()
|
|
cancel()
|
|
combined := outputBuf.String()
|
|
if fr.exitCode != 0 {
|
|
if combined == "" {
|
|
return "", fmt.Errorf("elevated execution failed (exit %d)", fr.exitCode)
|
|
}
|
|
errOutput := combined
|
|
const maxErrOutput = 8192
|
|
if len(errOutput) > maxErrOutput {
|
|
errOutput = errOutput[:maxErrOutput] + fmt.Sprintf("\n[...truncated %d bytes in error]", len(combined)-maxErrOutput)
|
|
}
|
|
return combined, fmt.Errorf("elevated execution failed (exit %d)\nOutput: %s", fr.exitCode, errOutput)
|
|
}
|
|
return combined, nil
|
|
|
|
case <-ctx.Done():
|
|
// Timeout or external cancellation
|
|
conn.Close()
|
|
cancel()
|
|
combined := outputBuf.String()
|
|
if ctx.Err() == context.DeadlineExceeded {
|
|
return combined, fmt.Errorf("elevated execution timeout after %d seconds", timeout)
|
|
}
|
|
return combined, fmt.Errorf("elevated execution interrupted")
|
|
|
|
case <-detachCh:
|
|
// Manual detach: background the connection read
|
|
cmdStr := "elevated: " + cmd
|
|
if len(cmdStr) > 80 {
|
|
cmdStr = cmdStr[:77] + "..."
|
|
}
|
|
ring := newRingBuffer(ringBufSize)
|
|
pid := int(time.Now().UnixNano() & 0x7FFFFFFF)
|
|
|
|
// Splice: future output goes to ring buffer, stop streaming
|
|
lw.mu.Lock()
|
|
lw.w = ring
|
|
lw.stream = nil
|
|
lw.mu.Unlock()
|
|
|
|
proc := &DetachedProcess{
|
|
PID: pid,
|
|
Command: cmdStr,
|
|
Started: time.Now(),
|
|
ring: ring,
|
|
done: make(chan struct{}),
|
|
}
|
|
e.detachMu.Lock()
|
|
e.detached = append(e.detached, proc)
|
|
e.detachMu.Unlock()
|
|
|
|
go func() {
|
|
fr := <-frameCh
|
|
conn.Close()
|
|
cancel()
|
|
proc.mu.Lock()
|
|
proc.Exited = true
|
|
proc.ExitCode = fr.exitCode
|
|
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]", pid), 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
|
|
}
|
|
|
|
if !trusted {
|
|
if err := e.ValidateCode(code, language); err != nil {
|
|
return "", err
|
|
}
|
|
}
|
|
|
|
var cfg *sandbox.Config
|
|
var err error
|
|
if !e.Yolo {
|
|
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}
|
|
default:
|
|
cancel()
|
|
return "", fmt.Errorf("unsupported language: %s (only bash is supported)", 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 {
|
|
cancel()
|
|
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,
|
|
stream: tools.StreamFunc(ctx),
|
|
}
|
|
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 {
|
|
cancel()
|
|
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 {
|
|
// Cap output in error message to avoid flooding the agent context.
|
|
// The full output is still returned as the first return value.
|
|
errOutput := string(output)
|
|
const maxErrOutput = 8192
|
|
if len(errOutput) > maxErrOutput {
|
|
errOutput = errOutput[:maxErrOutput] + fmt.Sprintf("\n[...truncated %d bytes in error]", len(output)-maxErrOutput)
|
|
}
|
|
return string(output), fmt.Errorf("execution failed: %v\nOutput: %s", err, errOutput)
|
|
}
|
|
return string(output), nil
|
|
|
|
case <-ctx.Done():
|
|
// Context cancelled (interrupt or timeout)
|
|
cancel()
|
|
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.stream = nil
|
|
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
|