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 <noreply@anthropic.com>
This commit is contained in:
dj-oyu
parent
52f3546cbb
commit
9cdd71420f
4 changed files with 689 additions and 660 deletions
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
573
pkg/tools/shell_bg.go
Normal file
573
pkg/tools/shell_bg.go
Normal 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
26
pkg/tools/shell_ext.go
Normal 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
82
pkg/tools/shell_net.go
Normal 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()
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue