opencode/picoclaw
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

refactor: extract shell.go fork additions into separate files

Browse source 
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 <noreply@anthropic.com>
This commit is contained in:
dj-oyu 2026-03-13 16:06:33 +09:00
parent ea2d48e907
commit 2674bc6f33
4 changed files with 689 additions and 660 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

668
pkg/tools/shell.go
View file

@ -5,191 +5,32 @@ import (
"context" "context"
"errors" "errors"
"fmt" "fmt"
"io"
"net"
"net/url"
"os" "os"
"os/exec" "os/exec"
"path/filepath" "path/filepath"
"regexp" "regexp"
"runtime" "runtime"
"sort"
"strings" "strings"
"sync"
"time" "time"
"github.com/sipeed/picoclaw/pkg/config" "github.com/sipeed/picoclaw/pkg/config"
"github.com/sipeed/picoclaw/pkg/constants" "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 { type ExecTool struct {
execToolExt // fork-specific fields (see shell_ext.go)
workingDir string workingDir string
timeout time.Duration timeout time.Duration
denyPatterns []*regexp.Regexp denyPatterns []*regexp.Regexp
allowRules [][]string // pre-split command prefix allowlist
customAllowPatterns []*regexp.Regexp customAllowPatterns []*regexp.Regexp
restrictToWorkspace bool restrictToWorkspace bool
localNetOnly bool // restrict curl/wget to localhost + RFC 1918
allowRemote bool 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 ( var (
@ -351,25 +192,23 @@ func NewExecToolWithConfig(workingDir string, restrict bool, config *config.Conf
bgCtx, bgCancel := context.WithCancel(context.Background()) bgCtx, bgCancel := context.WithCancel(context.Background())
return &ExecTool{ return &ExecTool{
execToolExt: execToolExt{
bgProcesses: make(map[string]*bgProcess),
bgCtx: bgCtx,
bgShutdown: bgCancel,
},
workingDir: workingDir, workingDir: workingDir,
timeout: timeout, timeout: timeout,
denyPatterns: denyPatterns, denyPatterns: denyPatterns,
allowRules: nil,
customAllowPatterns: customAllowPatterns, customAllowPatterns: customAllowPatterns,
restrictToWorkspace: restrict, restrictToWorkspace: restrict,
allowRemote: allowRemote, allowRemote: allowRemote,
bgProcesses: make(map[string]*bgProcess),
bgCtx: bgCtx,
bgShutdown: bgCancel,
}, nil }, 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 { func (t *ExecTool) guardCommand(command, cwd string) string {
cmd := strings.TrimSpace(command) cmd := strings.TrimSpace(command)
lower := strings.ToLower(cmd) lower := strings.ToLower(cmd)
@ -1256,97 +698,3 @@ func matchAllowRules(cmd string, rules [][]string) bool {
return false 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
}

573
pkg/tools/shell_bg.go Normal file
View file

@ -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
}

26
pkg/tools/shell_ext.go Normal file
View file

@ -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
}

82
pkg/tools/shell_net.go Normal file
View file

@ -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()
}