From 9cdd71420ff9c707b37b4328d01d3d22c6e0ff14 Mon Sep 17 00:00:00 2001 From: dj-oyu <68707227+dj-oyu@users.noreply.github.com> Date: Fri, 13 Mar 2026 16:06:33 +0900 Subject: [PATCH] refactor: extract shell.go fork additions into separate files Split ExecTool fork-specific code to reduce upstream merge conflicts: - shell_bg.go: background process system (ringBuffer, bgProcess, executeBg, handleBgAction, BgProcesses, RuntimeStatus, Shutdown) - shell_net.go: network restriction functions (SetLocalNetOnly, isCurlOrWget, checkCurlLocalNet, isLocalHost) - shell_ext.go: fork-specific fields (execToolExt embedded struct) Co-Authored-By: Claude Opus 4.6 --- pkg/tools/shell.go | 668 +---------------------------------------- pkg/tools/shell_bg.go | 573 +++++++++++++++++++++++++++++++++++ pkg/tools/shell_ext.go | 26 ++ pkg/tools/shell_net.go | 82 +++++ 4 files changed, 689 insertions(+), 660 deletions(-) create mode 100644 pkg/tools/shell_bg.go create mode 100644 pkg/tools/shell_ext.go create mode 100644 pkg/tools/shell_net.go diff --git a/pkg/tools/shell.go b/pkg/tools/shell.go index e2c8a5c11..16bb82b64 100644 --- a/pkg/tools/shell.go +++ b/pkg/tools/shell.go @@ -5,191 +5,32 @@ import ( "context" "errors" "fmt" - "io" - "net" - "net/url" "os" "os/exec" "path/filepath" "regexp" "runtime" - "sort" "strings" - "sync" "time" "github.com/sipeed/picoclaw/pkg/config" "github.com/sipeed/picoclaw/pkg/constants" ) -const ( - bgMaxLifetime = 45 * time.Minute - - bgRingBufSize = 32 * 1024 // 32KB - - bgInitCapture = 3 * time.Second - - bgMaxProcesses = 10 -) - -// ringBuffer is a thread-safe circular buffer that retains the most recent bytes. - -type ringBuffer struct { - mu sync.Mutex - - buf []byte - - size int -} - -func newRingBuffer(size int) *ringBuffer { - return &ringBuffer{size: size} -} - -// Write appends data to the ring buffer, dropping oldest bytes if capacity is exceeded. - -func (rb *ringBuffer) Write(p []byte) (int, error) { - rb.mu.Lock() - - defer rb.mu.Unlock() - - rb.buf = append(rb.buf, p...) - - if len(rb.buf) > rb.size { - rb.buf = rb.buf[len(rb.buf)-rb.size:] - } - - return len(p), nil -} - -// String returns the current buffer contents. - -func (rb *ringBuffer) String() string { - rb.mu.Lock() - - defer rb.mu.Unlock() - - return string(rb.buf) -} - -// Lines returns the last n lines from the buffer. - -func (rb *ringBuffer) Lines(n int) []string { - rb.mu.Lock() - - defer rb.mu.Unlock() - - if len(rb.buf) == 0 { - return nil - } - - all := strings.Split(string(rb.buf), "\n") - - // Remove trailing empty element from final newline - - if len(all) > 0 && all[len(all)-1] == "" { - all = all[:len(all)-1] - } - - if n <= 0 || n >= len(all) { - return all - } - - return all[len(all)-n:] -} - -// Match checks if any line in the buffer matches the given regex pattern. - -// Returns the first matching line, or empty string if no match. - -func (rb *ringBuffer) Match(pattern *regexp.Regexp) string { - rb.mu.Lock() - - defer rb.mu.Unlock() - - for _, line := range strings.Split(string(rb.buf), "\n") { - if pattern.MatchString(line) { - return line - } - } - - return "" -} - -// Len returns the current number of bytes in the buffer. - -func (rb *ringBuffer) Len() int { - rb.mu.Lock() - - defer rb.mu.Unlock() - - return len(rb.buf) -} - -// bgProcess represents a background process managed by ExecTool. - -type bgProcess struct { - id string - - command string - - cmd *exec.Cmd - - pid int - - startedAt time.Time - - output *ringBuffer - - done chan struct{} // closed when process exits - - exitErr error - - cancel context.CancelFunc // cancels the monitor goroutine -} - -// isRunning returns true if the process has not yet exited. - -func (bp *bgProcess) isRunning() bool { - select { - case <-bp.done: - - return false - - default: - - return true - } -} - type ExecTool struct { + execToolExt // fork-specific fields (see shell_ext.go) + workingDir string timeout time.Duration denyPatterns []*regexp.Regexp - allowRules [][]string // pre-split command prefix allowlist - customAllowPatterns []*regexp.Regexp restrictToWorkspace bool - localNetOnly bool // restrict curl/wget to localhost + RFC 1918 - allowRemote bool - - // Background process management - - bgMu sync.Mutex - - bgProcesses map[string]*bgProcess - - bgNextID int - - bgShutdown context.CancelFunc // cancels all bg monitor goroutines - - bgCtx context.Context } var ( @@ -351,25 +192,23 @@ func NewExecToolWithConfig(workingDir string, restrict bool, config *config.Conf bgCtx, bgCancel := context.WithCancel(context.Background()) return &ExecTool{ + execToolExt: execToolExt{ + bgProcesses: make(map[string]*bgProcess), + bgCtx: bgCtx, + bgShutdown: bgCancel, + }, + workingDir: workingDir, timeout: timeout, denyPatterns: denyPatterns, - allowRules: nil, - customAllowPatterns: customAllowPatterns, restrictToWorkspace: restrict, allowRemote: allowRemote, - - bgProcesses: make(map[string]*bgProcess), - - bgCtx: bgCtx, - - bgShutdown: bgCancel, }, nil } @@ -613,403 +452,6 @@ func (t *ExecTool) executeSync(ctx context.Context, command, cwd string) *ToolRe } } -// executeBg starts a background process and returns immediately. - -func (t *ExecTool) executeBg(command, cwd string) *ToolResult { - t.bgMu.Lock() - - // Check max processes limit - - running := 0 - - for _, bp := range t.bgProcesses { - if bp.isRunning() { - running++ - } - } - - if running >= bgMaxProcesses { - t.bgMu.Unlock() - - return ErrorResult( - - fmt.Sprintf("maximum background processes reached (%d). Kill an existing one first.", bgMaxProcesses), - ) - } - - t.bgNextID++ - - id := fmt.Sprintf("bg-%d", t.bgNextID) - - t.bgMu.Unlock() - - var cmd *exec.Cmd - - if runtime.GOOS == "windows" { - cmd = exec.Command("powershell", "-NoProfile", "-NonInteractive", "-Command", command) - } else { - cmd = exec.Command("sh", "-c", command) - } - - if cwd != "" { - cmd.Dir = cwd - } - - prepareCommandForTermination(cmd) - - output := newRingBuffer(bgRingBufSize) - - // Use pipes to capture output - - stdoutPipe, err := cmd.StdoutPipe() - if err != nil { - return ErrorResult(fmt.Sprintf("failed to create stdout pipe: %v", err)) - } - - stderrPipe, err := cmd.StderrPipe() - if err != nil { - return ErrorResult(fmt.Sprintf("failed to create stderr pipe: %v", err)) - } - - if err := cmd.Start(); err != nil { - return ErrorResult(fmt.Sprintf("failed to start background command: %v", err)) - } - - monitorCtx, monitorCancel := context.WithCancel(t.bgCtx) - - bp := &bgProcess{ - id: id, - - command: command, - - cmd: cmd, - - pid: cmd.Process.Pid, - - startedAt: time.Now(), - - output: output, - - done: make(chan struct{}), - - cancel: monitorCancel, - } - - t.bgMu.Lock() - - t.bgProcesses[id] = bp - - t.bgMu.Unlock() - - // io.Copy goroutines: pipe stdout/stderr into ring buffer - - go io.Copy(output, stdoutPipe) - - go io.Copy(output, stderrPipe) - - // cmd.Wait goroutine - - waitDone := make(chan error, 1) - - go func() { - waitDone <- cmd.Wait() - }() - - // Monitor goroutine: handles lifetime timer, process exit, and shutdown - - go func() { - lifetime := time.NewTimer(getBgMaxLifetime()) - - defer lifetime.Stop() - - select { - case err := <-waitDone: - - // Process exited naturally - - bp.exitErr = err - - close(bp.done) - - case <-lifetime.C: - - // Max lifetime exceeded — kill - - _ = terminateProcessTree(cmd) - - select { - case err := <-waitDone: - - bp.exitErr = err - - case <-time.After(2 * time.Second): - - if cmd.Process != nil { - _ = cmd.Process.Kill() - } - - bp.exitErr = <-waitDone - } - - close(bp.done) - - case <-monitorCtx.Done(): - - // Shutdown or explicit kill via cancel - - _ = terminateProcessTree(cmd) - - select { - case err := <-waitDone: - - bp.exitErr = err - - case <-time.After(2 * time.Second): - - if cmd.Process != nil { - _ = cmd.Process.Kill() - } - - bp.exitErr = <-waitDone - } - - select { - case <-bp.done: - - default: - - close(bp.done) - } - } - }() - - // Capture initial output (wait up to bgInitCapture) - - time.Sleep(bgInitCapture) - - initialOutput := output.String() - - var sb strings.Builder - - fmt.Fprintf(&sb, "Background process started.\n") - - fmt.Fprintf(&sb, " id: %s\n", id) - - fmt.Fprintf(&sb, " pid: %d\n", bp.pid) - - fmt.Fprintf(&sb, " cmd: %s\n", command) - - fmt.Fprintf(&sb, " max lifetime: %s\n", getBgMaxLifetime()) - - if initialOutput != "" { - fmt.Fprintf(&sb, "\nInitial output:\n%s", initialOutput) - } - - return &ToolResult{ - ForLLM: sb.String(), - - ForUser: fmt.Sprintf("Background process %s (pid=%d) started: %s", id, bp.pid, command), - } -} - -// handleBgAction handles bg_action=output and bg_action=kill. - -func (t *ExecTool) handleBgAction(action, bgID string) *ToolResult { - if bgID == "" { - return ErrorResult("bg_id is required for bg_action") - } - - t.bgMu.Lock() - - bp, ok := t.bgProcesses[bgID] - - t.bgMu.Unlock() - - if !ok { - return ErrorResult(fmt.Sprintf("background process %q not found", bgID)) - } - - switch action { - case "output": - - return t.bgOutput(bp) - - case "kill": - - return t.bgKill(bp) - - default: - - return ErrorResult(fmt.Sprintf("unknown bg_action %q (use 'output' or 'kill')", action)) - } -} - -func (t *ExecTool) bgOutput(bp *bgProcess) *ToolResult { - var sb strings.Builder - - fmt.Fprintf(&sb, "[%s] pid=%d %s\n", bp.id, bp.pid, bp.command) - - if bp.isRunning() { - uptime := time.Since(bp.startedAt).Truncate(time.Second) - - fmt.Fprintf(&sb, "Status: running (uptime: %s, max: %s)\n", uptime, getBgMaxLifetime()) - } else { - ran := time.Since(bp.startedAt).Truncate(time.Second) - - if bp.exitErr != nil { - fmt.Fprintf(&sb, "Status: exited with error (ran: %s): %v\n", ran, bp.exitErr) - } else { - fmt.Fprintf(&sb, "Status: exited=0 (ran: %s)\n", ran) - } - } - - output := bp.output.String() - - if output == "" { - fmt.Fprintf(&sb, "\n(no output)") - } else { - fmt.Fprintf(&sb, "\nOutput:\n%s", output) - } - - return &ToolResult{ - ForLLM: sb.String(), - - ForUser: sb.String(), - } -} - -func (t *ExecTool) bgKill(bp *bgProcess) *ToolResult { - if bp.isRunning() { - bp.cancel() // triggers monitor goroutine cleanup - - // Wait for process to actually exit - - select { - case <-bp.done: - - case <-time.After(5 * time.Second): - } - } - - t.bgMu.Lock() - - delete(t.bgProcesses, bp.id) - - t.bgMu.Unlock() - - msg := fmt.Sprintf("Background process %s (pid=%d) terminated: %s", bp.id, bp.pid, bp.command) - - return &ToolResult{ - ForLLM: msg, - - ForUser: msg, - } -} - -// BgProcesses returns a snapshot of background processes for use by bg_monitor. - -func (t *ExecTool) BgProcesses() map[string]*bgProcess { - t.bgMu.Lock() - - defer t.bgMu.Unlock() - - snapshot := make(map[string]*bgProcess, len(t.bgProcesses)) - - for k, v := range t.bgProcesses { - snapshot[k] = v - } - - return snapshot -} - -// RuntimeStatus implements StatusProvider for system prompt injection. - -func (t *ExecTool) RuntimeStatus() string { - t.bgMu.Lock() - - defer t.bgMu.Unlock() - - if len(t.bgProcesses) == 0 { - return "" - } - - // Sort by ID for stable output - - ids := make([]string, 0, len(t.bgProcesses)) - - for id := range t.bgProcesses { - ids = append(ids, id) - } - - sort.Strings(ids) - - var sb strings.Builder - - sb.WriteString("## Background Processes\n\n") - - for _, id := range ids { - bp := t.bgProcesses[id] - - if bp.isRunning() { - uptime := time.Since(bp.startedAt).Truncate(time.Second) - - fmt.Fprintf(&sb, " [%s] pid=%d running (uptime: %s, max: %s) %s\n", - - id, bp.pid, uptime, getBgMaxLifetime(), bp.command) - } else { - ran := time.Since(bp.startedAt).Truncate(time.Second) - - if bp.exitErr != nil { - fmt.Fprintf(&sb, " [%s] pid=%d exited=err (ran: %s) %s\n", - - id, bp.pid, ran, bp.command) - } else { - fmt.Fprintf(&sb, " [%s] pid=%d exited=0 (ran: %s) %s\n", - - id, bp.pid, ran, bp.command) - } - } - } - - sb.WriteString("\nUse exec with bg_action=\"output\" / \"kill\" and bg_id to manage.\n") - - sb.WriteString("Use bg_monitor for list/watch/tail operations.") - - return sb.String() -} - -// Shutdown terminates all background processes. Call on application exit. - -func (t *ExecTool) Shutdown() { - t.bgShutdown() // cancel all monitor goroutines - - t.bgMu.Lock() - - procs := make([]*bgProcess, 0, len(t.bgProcesses)) - - for _, bp := range t.bgProcesses { - procs = append(procs, bp) - } - - t.bgMu.Unlock() - - // Wait for all processes to exit - - for _, bp := range procs { - select { - case <-bp.done: - - case <-time.After(5 * time.Second): - - // Force kill if still running - - if bp.cmd.Process != nil { - _ = bp.cmd.Process.Kill() - } - } - } -} - func (t *ExecTool) guardCommand(command, cwd string) string { cmd := strings.TrimSpace(command) lower := strings.ToLower(cmd) @@ -1256,97 +698,3 @@ func matchAllowRules(cmd string, rules [][]string) bool { return false } - -func (t *ExecTool) SetLocalNetOnly(v bool) { - t.localNetOnly = v -} - -// isCurlOrWget reports whether command is a curl or wget invocation. - -func isCurlOrWget(command string) bool { - fields := strings.Fields(command) - - if len(fields) == 0 { - return false - } - - base := filepath.Base(fields[0]) - - return base == "curl" || base == "wget" -} - -// checkCurlLocalNet validates that all http/https URLs in a curl/wget command - -// target localhost or RFC 1918 private addresses. - -// Returns an error message string, or empty string if the command is allowed. - -func checkCurlLocalNet(command string) string { - for _, token := range strings.Fields(command) { - token = strings.Trim(token, "\"'") - - if !strings.HasPrefix(token, "http://") && !strings.HasPrefix(token, "https://") { - continue - } - - u, err := url.Parse(token) - if err != nil { - continue - } - - host := u.Hostname() - - if !isLocalHost(host) { - return fmt.Sprintf( - - "Command blocked by safety guard "+ - - "(curl/wget is restricted to localhost and private network; %q is a public address)", - - host, - ) - } - } - - return "" -} - -// isLocalHost reports whether host is localhost or a loopback/RFC 1918 private IP. - -// DNS resolution is intentionally avoided to prevent DNS rebinding attacks. - -func isLocalHost(host string) bool { - if strings.EqualFold(host, "localhost") { - return true - } - - ip := net.ParseIP(host) - - if ip == nil { - return false - } - - return ip.IsLoopback() || ip.IsPrivate() -} - -// SetBgMaxLifetimeForTest overrides bgMaxLifetime for testing purposes. - -// This is exposed only for tests; the returned function restores the original value. - -var bgMaxLifetimeOverride time.Duration - -func SetBgMaxLifetimeForTest(d time.Duration) func() { - old := bgMaxLifetimeOverride - - bgMaxLifetimeOverride = d - - return func() { bgMaxLifetimeOverride = old } -} - -func getBgMaxLifetime() time.Duration { - if bgMaxLifetimeOverride > 0 { - return bgMaxLifetimeOverride - } - - return bgMaxLifetime -} diff --git a/pkg/tools/shell_bg.go b/pkg/tools/shell_bg.go new file mode 100644 index 000000000..19ec015fc --- /dev/null +++ b/pkg/tools/shell_bg.go @@ -0,0 +1,573 @@ +package tools + +import ( + "context" + "fmt" + "io" + "os/exec" + "regexp" + "runtime" + "sort" + "strings" + "sync" + "time" +) + +const ( + bgMaxLifetime = 45 * time.Minute + + bgRingBufSize = 32 * 1024 // 32KB + + bgInitCapture = 3 * time.Second + + bgMaxProcesses = 10 +) + +// ringBuffer is a thread-safe circular buffer that retains the most recent bytes. + +type ringBuffer struct { + mu sync.Mutex + + buf []byte + + size int +} + +func newRingBuffer(size int) *ringBuffer { + return &ringBuffer{size: size} +} + +// Write appends data to the ring buffer, dropping oldest bytes if capacity is exceeded. + +func (rb *ringBuffer) Write(p []byte) (int, error) { + rb.mu.Lock() + + defer rb.mu.Unlock() + + rb.buf = append(rb.buf, p...) + + if len(rb.buf) > rb.size { + rb.buf = rb.buf[len(rb.buf)-rb.size:] + } + + return len(p), nil +} + +// String returns the current buffer contents. + +func (rb *ringBuffer) String() string { + rb.mu.Lock() + + defer rb.mu.Unlock() + + return string(rb.buf) +} + +// Lines returns the last n lines from the buffer. + +func (rb *ringBuffer) Lines(n int) []string { + rb.mu.Lock() + + defer rb.mu.Unlock() + + if len(rb.buf) == 0 { + return nil + } + + all := strings.Split(string(rb.buf), "\n") + + // Remove trailing empty element from final newline + + if len(all) > 0 && all[len(all)-1] == "" { + all = all[:len(all)-1] + } + + if n <= 0 || n >= len(all) { + return all + } + + return all[len(all)-n:] +} + +// Match checks if any line in the buffer matches the given regex pattern. + +// Returns the first matching line, or empty string if no match. + +func (rb *ringBuffer) Match(pattern *regexp.Regexp) string { + rb.mu.Lock() + + defer rb.mu.Unlock() + + for _, line := range strings.Split(string(rb.buf), "\n") { + if pattern.MatchString(line) { + return line + } + } + + return "" +} + +// Len returns the current number of bytes in the buffer. + +func (rb *ringBuffer) Len() int { + rb.mu.Lock() + + defer rb.mu.Unlock() + + return len(rb.buf) +} + +// bgProcess represents a background process managed by ExecTool. + +type bgProcess struct { + id string + + command string + + cmd *exec.Cmd + + pid int + + startedAt time.Time + + output *ringBuffer + + done chan struct{} // closed when process exits + + exitErr error + + cancel context.CancelFunc // cancels the monitor goroutine +} + +// isRunning returns true if the process has not yet exited. + +func (bp *bgProcess) isRunning() bool { + select { + case <-bp.done: + + return false + + default: + + return true + } +} + +// executeBg starts a background process and returns immediately. + +func (t *ExecTool) executeBg(command, cwd string) *ToolResult { + t.bgMu.Lock() + + // Check max processes limit + + running := 0 + + for _, bp := range t.bgProcesses { + if bp.isRunning() { + running++ + } + } + + if running >= bgMaxProcesses { + t.bgMu.Unlock() + + return ErrorResult( + + fmt.Sprintf("maximum background processes reached (%d). Kill an existing one first.", bgMaxProcesses), + ) + } + + t.bgNextID++ + + id := fmt.Sprintf("bg-%d", t.bgNextID) + + t.bgMu.Unlock() + + var cmd *exec.Cmd + + if runtime.GOOS == "windows" { + cmd = exec.Command("powershell", "-NoProfile", "-NonInteractive", "-Command", command) + } else { + cmd = exec.Command("sh", "-c", command) + } + + if cwd != "" { + cmd.Dir = cwd + } + + prepareCommandForTermination(cmd) + + output := newRingBuffer(bgRingBufSize) + + // Use pipes to capture output + + stdoutPipe, err := cmd.StdoutPipe() + if err != nil { + return ErrorResult(fmt.Sprintf("failed to create stdout pipe: %v", err)) + } + + stderrPipe, err := cmd.StderrPipe() + if err != nil { + return ErrorResult(fmt.Sprintf("failed to create stderr pipe: %v", err)) + } + + if err := cmd.Start(); err != nil { + return ErrorResult(fmt.Sprintf("failed to start background command: %v", err)) + } + + monitorCtx, monitorCancel := context.WithCancel(t.bgCtx) + + bp := &bgProcess{ + id: id, + + command: command, + + cmd: cmd, + + pid: cmd.Process.Pid, + + startedAt: time.Now(), + + output: output, + + done: make(chan struct{}), + + cancel: monitorCancel, + } + + t.bgMu.Lock() + + t.bgProcesses[id] = bp + + t.bgMu.Unlock() + + // io.Copy goroutines: pipe stdout/stderr into ring buffer + + go io.Copy(output, stdoutPipe) + + go io.Copy(output, stderrPipe) + + // cmd.Wait goroutine + + waitDone := make(chan error, 1) + + go func() { + waitDone <- cmd.Wait() + }() + + // Monitor goroutine: handles lifetime timer, process exit, and shutdown + + go func() { + lifetime := time.NewTimer(getBgMaxLifetime()) + + defer lifetime.Stop() + + select { + case err := <-waitDone: + + // Process exited naturally + + bp.exitErr = err + + close(bp.done) + + case <-lifetime.C: + + // Max lifetime exceeded — kill + + _ = terminateProcessTree(cmd) + + select { + case err := <-waitDone: + + bp.exitErr = err + + case <-time.After(2 * time.Second): + + if cmd.Process != nil { + _ = cmd.Process.Kill() + } + + bp.exitErr = <-waitDone + } + + close(bp.done) + + case <-monitorCtx.Done(): + + // Shutdown or explicit kill via cancel + + _ = terminateProcessTree(cmd) + + select { + case err := <-waitDone: + + bp.exitErr = err + + case <-time.After(2 * time.Second): + + if cmd.Process != nil { + _ = cmd.Process.Kill() + } + + bp.exitErr = <-waitDone + } + + select { + case <-bp.done: + + default: + + close(bp.done) + } + } + }() + + // Capture initial output (wait up to bgInitCapture) + + time.Sleep(bgInitCapture) + + initialOutput := output.String() + + var sb strings.Builder + + fmt.Fprintf(&sb, "Background process started.\n") + + fmt.Fprintf(&sb, " id: %s\n", id) + + fmt.Fprintf(&sb, " pid: %d\n", bp.pid) + + fmt.Fprintf(&sb, " cmd: %s\n", command) + + fmt.Fprintf(&sb, " max lifetime: %s\n", getBgMaxLifetime()) + + if initialOutput != "" { + fmt.Fprintf(&sb, "\nInitial output:\n%s", initialOutput) + } + + return &ToolResult{ + ForLLM: sb.String(), + + ForUser: fmt.Sprintf("Background process %s (pid=%d) started: %s", id, bp.pid, command), + } +} + +// handleBgAction handles bg_action=output and bg_action=kill. + +func (t *ExecTool) handleBgAction(action, bgID string) *ToolResult { + if bgID == "" { + return ErrorResult("bg_id is required for bg_action") + } + + t.bgMu.Lock() + + bp, ok := t.bgProcesses[bgID] + + t.bgMu.Unlock() + + if !ok { + return ErrorResult(fmt.Sprintf("background process %q not found", bgID)) + } + + switch action { + case "output": + + return t.bgOutput(bp) + + case "kill": + + return t.bgKill(bp) + + default: + + return ErrorResult(fmt.Sprintf("unknown bg_action %q (use 'output' or 'kill')", action)) + } +} + +func (t *ExecTool) bgOutput(bp *bgProcess) *ToolResult { + var sb strings.Builder + + fmt.Fprintf(&sb, "[%s] pid=%d %s\n", bp.id, bp.pid, bp.command) + + if bp.isRunning() { + uptime := time.Since(bp.startedAt).Truncate(time.Second) + + fmt.Fprintf(&sb, "Status: running (uptime: %s, max: %s)\n", uptime, getBgMaxLifetime()) + } else { + ran := time.Since(bp.startedAt).Truncate(time.Second) + + if bp.exitErr != nil { + fmt.Fprintf(&sb, "Status: exited with error (ran: %s): %v\n", ran, bp.exitErr) + } else { + fmt.Fprintf(&sb, "Status: exited=0 (ran: %s)\n", ran) + } + } + + output := bp.output.String() + + if output == "" { + fmt.Fprintf(&sb, "\n(no output)") + } else { + fmt.Fprintf(&sb, "\nOutput:\n%s", output) + } + + return &ToolResult{ + ForLLM: sb.String(), + + ForUser: sb.String(), + } +} + +func (t *ExecTool) bgKill(bp *bgProcess) *ToolResult { + if bp.isRunning() { + bp.cancel() // triggers monitor goroutine cleanup + + // Wait for process to actually exit + + select { + case <-bp.done: + + case <-time.After(5 * time.Second): + } + } + + t.bgMu.Lock() + + delete(t.bgProcesses, bp.id) + + t.bgMu.Unlock() + + msg := fmt.Sprintf("Background process %s (pid=%d) terminated: %s", bp.id, bp.pid, bp.command) + + return &ToolResult{ + ForLLM: msg, + + ForUser: msg, + } +} + +// BgProcesses returns a snapshot of background processes for use by bg_monitor. + +func (t *ExecTool) BgProcesses() map[string]*bgProcess { + t.bgMu.Lock() + + defer t.bgMu.Unlock() + + snapshot := make(map[string]*bgProcess, len(t.bgProcesses)) + + for k, v := range t.bgProcesses { + snapshot[k] = v + } + + return snapshot +} + +// RuntimeStatus implements StatusProvider for system prompt injection. + +func (t *ExecTool) RuntimeStatus() string { + t.bgMu.Lock() + + defer t.bgMu.Unlock() + + if len(t.bgProcesses) == 0 { + return "" + } + + // Sort by ID for stable output + + ids := make([]string, 0, len(t.bgProcesses)) + + for id := range t.bgProcesses { + ids = append(ids, id) + } + + sort.Strings(ids) + + var sb strings.Builder + + sb.WriteString("## Background Processes\n\n") + + for _, id := range ids { + bp := t.bgProcesses[id] + + if bp.isRunning() { + uptime := time.Since(bp.startedAt).Truncate(time.Second) + + fmt.Fprintf(&sb, " [%s] pid=%d running (uptime: %s, max: %s) %s\n", + + id, bp.pid, uptime, getBgMaxLifetime(), bp.command) + } else { + ran := time.Since(bp.startedAt).Truncate(time.Second) + + if bp.exitErr != nil { + fmt.Fprintf(&sb, " [%s] pid=%d exited=err (ran: %s) %s\n", + + id, bp.pid, ran, bp.command) + } else { + fmt.Fprintf(&sb, " [%s] pid=%d exited=0 (ran: %s) %s\n", + + id, bp.pid, ran, bp.command) + } + } + } + + sb.WriteString("\nUse exec with bg_action=\"output\" / \"kill\" and bg_id to manage.\n") + + sb.WriteString("Use bg_monitor for list/watch/tail operations.") + + return sb.String() +} + +// Shutdown terminates all background processes. Call on application exit. + +func (t *ExecTool) Shutdown() { + t.bgShutdown() // cancel all monitor goroutines + + t.bgMu.Lock() + + procs := make([]*bgProcess, 0, len(t.bgProcesses)) + + for _, bp := range t.bgProcesses { + procs = append(procs, bp) + } + + t.bgMu.Unlock() + + // Wait for all processes to exit + + for _, bp := range procs { + select { + case <-bp.done: + + case <-time.After(5 * time.Second): + + // Force kill if still running + + if bp.cmd.Process != nil { + _ = bp.cmd.Process.Kill() + } + } + } +} + +// SetBgMaxLifetimeForTest overrides bgMaxLifetime for testing purposes. + +// This is exposed only for tests; the returned function restores the original value. + +var bgMaxLifetimeOverride time.Duration + +func SetBgMaxLifetimeForTest(d time.Duration) func() { + old := bgMaxLifetimeOverride + + bgMaxLifetimeOverride = d + + return func() { bgMaxLifetimeOverride = old } +} + +func getBgMaxLifetime() time.Duration { + if bgMaxLifetimeOverride > 0 { + return bgMaxLifetimeOverride + } + + return bgMaxLifetime +} diff --git a/pkg/tools/shell_ext.go b/pkg/tools/shell_ext.go new file mode 100644 index 000000000..e3b034148 --- /dev/null +++ b/pkg/tools/shell_ext.go @@ -0,0 +1,26 @@ +package tools + +import ( + "context" + "sync" +) + +// execToolExt holds fork-specific fields for ExecTool. +// Embedded in ExecTool so existing field access (t.bgProcesses, etc.) continues to work. +type execToolExt struct { + allowRules [][]string // pre-split command prefix allowlist + + localNetOnly bool // restrict curl/wget to localhost + RFC 1918 + + // Background process management + + bgMu sync.Mutex + + bgProcesses map[string]*bgProcess + + bgNextID int + + bgShutdown context.CancelFunc // cancels all bg monitor goroutines + + bgCtx context.Context +} diff --git a/pkg/tools/shell_net.go b/pkg/tools/shell_net.go new file mode 100644 index 000000000..34f3397a6 --- /dev/null +++ b/pkg/tools/shell_net.go @@ -0,0 +1,82 @@ +package tools + +import ( + "fmt" + "net" + "net/url" + "path/filepath" + "strings" +) + +// SetLocalNetOnly restricts curl/wget to localhost and RFC 1918 private addresses. +func (t *ExecTool) SetLocalNetOnly(v bool) { + t.localNetOnly = v +} + +// isCurlOrWget reports whether command is a curl or wget invocation. + +func isCurlOrWget(command string) bool { + fields := strings.Fields(command) + + if len(fields) == 0 { + return false + } + + base := filepath.Base(fields[0]) + + return base == "curl" || base == "wget" +} + +// checkCurlLocalNet validates that all http/https URLs in a curl/wget command + +// target localhost or RFC 1918 private addresses. + +// Returns an error message string, or empty string if the command is allowed. + +func checkCurlLocalNet(command string) string { + for _, token := range strings.Fields(command) { + token = strings.Trim(token, "\"'") + + if !strings.HasPrefix(token, "http://") && !strings.HasPrefix(token, "https://") { + continue + } + + u, err := url.Parse(token) + if err != nil { + continue + } + + host := u.Hostname() + + if !isLocalHost(host) { + return fmt.Sprintf( + + "Command blocked by safety guard "+ + + "(curl/wget is restricted to localhost and private network; %q is a public address)", + + host, + ) + } + } + + return "" +} + +// isLocalHost reports whether host is localhost or a loopback/RFC 1918 private IP. + +// DNS resolution is intentionally avoided to prevent DNS rebinding attacks. + +func isLocalHost(host string) bool { + if strings.EqualFold(host, "localhost") { + return true + } + + ip := net.ParseIP(host) + + if ip == nil { + return false + } + + return ip.IsLoopback() || ip.IsPrivate() +}