harden tool security boundaries across web, filesystem, and exec
Phase 1 - Security regression baseline - add blocking tests for loopback and redirect-to-private web_fetch targets - add filesystem symlink escape and prefix-bypass restriction tests - add shell timeout test to verify child-process cleanup behavior Phase 2 - Web fetch egress guard - introduce stdlib-only fetch target validator for host/IP policy checks - enforce target validation before connect, during redirect, and in dial path - keep public constructor API and add internal test-aware constructor path Phase 3 - Canonical filesystem boundary enforcement - replace prefix checks with canonical workspace containment via filepath.Rel - canonicalize workspace and target paths with symlink-aware resolution - protect create/read/write/list/edit/append flows through shared validation Phase 4 - Shell timeout process-tree cleanup - switch command execution to explicit start/wait with timeout select - add OS-specific process tree helpers for unix process groups and taskkill on windows - preserve existing output contract and timeout messaging semantics Phase 5 - Documentation and contributor guidance - document web_fetch network boundary and complementary security model - add tool security checklist for future built-in tool additions Verification - go test ./pkg/tools -run TestWebTool_WebFetch_Blocks (via golang:1.25) - go test ./pkg/tools -run TestFilesystemTool_Restrict (via golang:1.25) - go test ./pkg/tools -run TestShellTool_Timeout_KillsChildProcesses (via golang:1.25) - go test ./pkg/tools -run TestWebTool_WebFetch_ (via golang:1.25) - go test ./pkg/tools -run TestFilesystemTool_ (via golang:1.25) - go test ./pkg/tools -run TestShellTool_ (via golang:1.25) - go test ./pkg/tools (via golang:1.25) - go generate ./... && go test ./... (via golang:1.25)
This commit is contained in:
Jared Mahotiere
parent
3b4b6bfc02
commit
ddc942542e
11 changed files with 648 additions and 28 deletions
27
README.md
27
README.md
|
|
@ -511,6 +511,21 @@ When `restrict_to_workspace: true`, the following tools are sandboxed:
|
|||
| `append_file` | Append to files | Only files within workspace |
|
||||
| `exec` | Execute commands | Command paths must be within workspace |
|
||||
|
||||
#### Web Fetch Network Boundary
|
||||
|
||||
`web_fetch` enforces an outbound network boundary independent of `restrict_to_workspace`.
|
||||
|
||||
Blocked destination classes include:
|
||||
|
||||
* loopback
|
||||
* private RFC1918 / unique-local ranges
|
||||
* link-local
|
||||
* multicast
|
||||
* unspecified / non-routable internal targets
|
||||
* redirect hops that resolve to blocked targets
|
||||
|
||||
This policy is applied both before connect and during redirect handling, so a public URL cannot bounce into private infrastructure through redirects.
|
||||
|
||||
#### Additional Exec Protection
|
||||
|
||||
Even with `restrict_to_workspace: false`, the `exec` tool blocks these dangerous commands:
|
||||
|
|
@ -534,6 +549,11 @@ Even with `restrict_to_workspace: false`, the `exec` tool blocks these dangerous
|
|||
{tool=exec, error=Command blocked by safety guard (dangerous pattern detected)}
|
||||
```
|
||||
|
||||
```
|
||||
[ERROR] tool: Tool execution failed
|
||||
{tool=web_fetch, error=blocked destination: host "127.0.0.1" resolves to non-public IP 127.0.0.1}
|
||||
```
|
||||
|
||||
#### Disabling Restrictions (Security Risk)
|
||||
|
||||
If you need the agent to access paths outside the workspace:
|
||||
|
|
@ -560,7 +580,12 @@ export PICOCLAW_AGENTS_DEFAULTS_RESTRICT_TO_WORKSPACE=false
|
|||
|
||||
#### Security Boundary Consistency
|
||||
|
||||
The `restrict_to_workspace` setting applies consistently across all execution paths:
|
||||
PicoClaw enforces complementary boundaries:
|
||||
|
||||
* Filesystem + shell boundary via `restrict_to_workspace`
|
||||
* Network egress boundary via `web_fetch` public-target validation
|
||||
|
||||
The workspace boundary (`restrict_to_workspace`) applies consistently across all execution paths:
|
||||
|
||||
| Execution Path | Security Boundary |
|
||||
|----------------|-------------------|
|
||||
|
|
|
|||
9
pkg/tools/README.md
Normal file
9
pkg/tools/README.md
Normal file
|
|
@ -0,0 +1,9 @@
|
|||
# Tool Security Checklist
|
||||
|
||||
When adding a new built-in tool, include these minimum safety checks:
|
||||
|
||||
1. Path boundary: if the tool reads/writes files or executes commands with paths, enforce canonical workspace membership when `restrict_to_workspace=true`.
|
||||
2. Network boundary: if the tool performs outbound network calls, reject loopback/private/link-local/multicast/unspecified/internal targets and validate redirect hops.
|
||||
3. Timeout behavior: long-running operations must use deterministic timeout/cancel handling and terminate child processes where process trees are possible.
|
||||
4. Regression tests: add explicit tests for blocked behavior (not just happy-path errors), including redirect/path traversal/process-leak scenarios where relevant.
|
||||
5. Error clarity: return explicit denial reasons (`blocked destination`, `outside workspace`, `timed out`) so behavior is auditable in logs.
|
||||
|
|
@ -19,21 +19,98 @@ func validatePath(path, workspace string, restrict bool) (string, error) {
|
|||
return "", fmt.Errorf("failed to resolve workspace path: %w", err)
|
||||
}
|
||||
|
||||
var absPath string
|
||||
absPath, err := resolveTargetPath(path, absWorkspace)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
if !restrict {
|
||||
return absPath, nil
|
||||
}
|
||||
|
||||
canonicalPath, err := pathWithinWorkspace(absPath, absWorkspace)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
return canonicalPath, nil
|
||||
}
|
||||
|
||||
func resolveTargetPath(path, absWorkspace string) (string, error) {
|
||||
if filepath.IsAbs(path) {
|
||||
absPath = filepath.Clean(path)
|
||||
} else {
|
||||
absPath, err = filepath.Abs(filepath.Join(absWorkspace, path))
|
||||
return filepath.Clean(path), nil
|
||||
}
|
||||
|
||||
absPath, err := filepath.Abs(filepath.Join(absWorkspace, path))
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("failed to resolve file path: %w", err)
|
||||
}
|
||||
return filepath.Clean(absPath), nil
|
||||
}
|
||||
|
||||
func pathWithinWorkspace(target, workspace string) (string, error) {
|
||||
canonicalWorkspace, err := canonicalizeExistingPath(workspace)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("failed to canonicalize workspace path: %w", err)
|
||||
}
|
||||
|
||||
if restrict && !strings.HasPrefix(absPath, absWorkspace) {
|
||||
canonicalTarget, err := canonicalizePathForBoundary(target)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("failed to canonicalize target path: %w", err)
|
||||
}
|
||||
|
||||
rel, err := filepath.Rel(canonicalWorkspace, canonicalTarget)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("failed to evaluate workspace boundary: %w", err)
|
||||
}
|
||||
|
||||
if rel == ".." || strings.HasPrefix(rel, ".."+string(filepath.Separator)) {
|
||||
return "", fmt.Errorf("access denied: path is outside the workspace")
|
||||
}
|
||||
|
||||
return absPath, nil
|
||||
return canonicalTarget, nil
|
||||
}
|
||||
|
||||
func canonicalizeExistingPath(path string) (string, error) {
|
||||
resolved, err := filepath.EvalSymlinks(path)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return filepath.Clean(resolved), nil
|
||||
}
|
||||
|
||||
func canonicalizePathForBoundary(path string) (string, error) {
|
||||
cleanPath := filepath.Clean(path)
|
||||
segments := make([]string, 0, 4)
|
||||
current := cleanPath
|
||||
|
||||
for {
|
||||
_, err := os.Lstat(current)
|
||||
if err == nil {
|
||||
resolved, evalErr := filepath.EvalSymlinks(current)
|
||||
if evalErr != nil {
|
||||
return "", evalErr
|
||||
}
|
||||
|
||||
for i := len(segments) - 1; i >= 0; i-- {
|
||||
resolved = filepath.Join(resolved, segments[i])
|
||||
}
|
||||
|
||||
return filepath.Clean(resolved), nil
|
||||
}
|
||||
|
||||
if !os.IsNotExist(err) {
|
||||
return "", err
|
||||
}
|
||||
|
||||
parent := filepath.Dir(current)
|
||||
if parent == current {
|
||||
return "", fmt.Errorf("could not resolve existing parent for %q", path)
|
||||
}
|
||||
|
||||
segments = append(segments, filepath.Base(current))
|
||||
current = parent
|
||||
}
|
||||
}
|
||||
|
||||
type ReadFileTool struct {
|
||||
|
|
|
|||
|
|
@ -247,3 +247,55 @@ func TestFilesystemTool_ListDir_DefaultPath(t *testing.T) {
|
|||
t.Errorf("Expected success with default path '.', got IsError=true: %s", result.ForLLM)
|
||||
}
|
||||
}
|
||||
|
||||
// TestFilesystemTool_Restrict_BlocksSymlinkEscape verifies symlink traversal outside workspace is blocked.
|
||||
func TestFilesystemTool_Restrict_BlocksSymlinkEscape(t *testing.T) {
|
||||
workspace := t.TempDir()
|
||||
outside := t.TempDir()
|
||||
secretFile := filepath.Join(outside, "secret.txt")
|
||||
if err := os.WriteFile(secretFile, []byte("do-not-read"), 0644); err != nil {
|
||||
t.Fatalf("failed to create outside file: %v", err)
|
||||
}
|
||||
|
||||
linkPath := filepath.Join(workspace, "leak.txt")
|
||||
if err := os.Symlink(secretFile, linkPath); err != nil {
|
||||
t.Skipf("symlink unavailable in this environment: %v", err)
|
||||
}
|
||||
|
||||
tool := NewReadFileTool(workspace, true)
|
||||
result := tool.Execute(context.Background(), map[string]interface{}{"path": "leak.txt"})
|
||||
if !result.IsError {
|
||||
t.Fatalf("Expected symlink escape to be blocked")
|
||||
}
|
||||
if !strings.Contains(strings.ToLower(result.ForLLM), "outside") {
|
||||
t.Fatalf("Expected workspace boundary error, got: %s", result.ForLLM)
|
||||
}
|
||||
}
|
||||
|
||||
// TestFilesystemTool_Restrict_BlocksPrefixBypass verifies prefix confusion paths are blocked.
|
||||
func TestFilesystemTool_Restrict_BlocksPrefixBypass(t *testing.T) {
|
||||
parent := t.TempDir()
|
||||
workspace := filepath.Join(parent, "workspace")
|
||||
prefixBypassDir := filepath.Join(parent, "workspace-evil")
|
||||
|
||||
if err := os.MkdirAll(workspace, 0755); err != nil {
|
||||
t.Fatalf("failed to create workspace: %v", err)
|
||||
}
|
||||
if err := os.MkdirAll(prefixBypassDir, 0755); err != nil {
|
||||
t.Fatalf("failed to create bypass dir: %v", err)
|
||||
}
|
||||
|
||||
bypassFile := filepath.Join(prefixBypassDir, "stolen.txt")
|
||||
if err := os.WriteFile(bypassFile, []byte("secret"), 0644); err != nil {
|
||||
t.Fatalf("failed to create bypass file: %v", err)
|
||||
}
|
||||
|
||||
tool := NewReadFileTool(workspace, true)
|
||||
result := tool.Execute(context.Background(), map[string]interface{}{"path": bypassFile})
|
||||
if !result.IsError {
|
||||
t.Fatalf("Expected prefix bypass path to be blocked")
|
||||
}
|
||||
if !strings.Contains(strings.ToLower(result.ForLLM), "outside") {
|
||||
t.Fatalf("Expected workspace boundary error, got: %s", result.ForLLM)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
178
pkg/tools/network_guard.go
Normal file
178
pkg/tools/network_guard.go
Normal file
|
|
@ -0,0 +1,178 @@
|
|||
package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net"
|
||||
"net/netip"
|
||||
"net/url"
|
||||
"strings"
|
||||
)
|
||||
|
||||
var (
|
||||
cgnatPrefix = netip.MustParsePrefix("100.64.0.0/10")
|
||||
benchmarkPrefix = netip.MustParsePrefix("198.18.0.0/15")
|
||||
reservedPrefix = netip.MustParsePrefix("240.0.0.0/4")
|
||||
)
|
||||
|
||||
type fetchTargetValidator struct {
|
||||
resolver *net.Resolver
|
||||
allowedHosts map[string]struct{}
|
||||
}
|
||||
|
||||
func newFetchTargetValidator(allowHosts []string, resolver *net.Resolver) *fetchTargetValidator {
|
||||
allowed := make(map[string]struct{}, len(allowHosts))
|
||||
for _, host := range allowHosts {
|
||||
normalized := normalizeHostToken(host)
|
||||
if normalized != "" {
|
||||
allowed[normalized] = struct{}{}
|
||||
}
|
||||
}
|
||||
|
||||
if resolver == nil {
|
||||
resolver = net.DefaultResolver
|
||||
}
|
||||
|
||||
return &fetchTargetValidator{
|
||||
resolver: resolver,
|
||||
allowedHosts: allowed,
|
||||
}
|
||||
}
|
||||
|
||||
func (v *fetchTargetValidator) validateURL(ctx context.Context, target *url.URL) error {
|
||||
host := normalizeHostToken(target.Hostname())
|
||||
if host == "" {
|
||||
return fmt.Errorf("missing domain in URL")
|
||||
}
|
||||
|
||||
port := target.Port()
|
||||
if v.isAllowed(host, port) {
|
||||
return nil
|
||||
}
|
||||
|
||||
if isBlockedHostname(host) {
|
||||
return fmt.Errorf("blocked destination: host %q is internal-only", target.Hostname())
|
||||
}
|
||||
|
||||
if ip, ok := parseIPLiteral(host); ok {
|
||||
if IsBlockedIP(ip) {
|
||||
return fmt.Errorf("blocked destination: IP %s is not publicly routable", ip)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
addrs, err := v.resolver.LookupNetIP(ctx, "ip", host)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to resolve host %q: %w", host, err)
|
||||
}
|
||||
if len(addrs) == 0 {
|
||||
return fmt.Errorf("failed to resolve host %q: no records", host)
|
||||
}
|
||||
|
||||
for _, addr := range addrs {
|
||||
if IsBlockedIP(addr) {
|
||||
return fmt.Errorf("blocked destination: host %q resolves to non-public IP %s", host, addr)
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (v *fetchTargetValidator) isAllowed(host, port string) bool {
|
||||
if len(v.allowedHosts) == 0 {
|
||||
return false
|
||||
}
|
||||
|
||||
if _, ok := v.allowedHosts[host]; ok {
|
||||
return true
|
||||
}
|
||||
if port != "" {
|
||||
if _, ok := v.allowedHosts[host+":"+port]; ok {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// ValidateFetchTarget applies the default web fetch target policy.
|
||||
func ValidateFetchTarget(target *url.URL) error {
|
||||
return newFetchTargetValidator(nil, net.DefaultResolver).validateURL(context.Background(), target)
|
||||
}
|
||||
|
||||
// IsBlockedIP returns true for IPs that should not be reachable from web_fetch.
|
||||
func IsBlockedIP(addr netip.Addr) bool {
|
||||
if !addr.IsValid() {
|
||||
return true
|
||||
}
|
||||
|
||||
if addr.IsLoopback() ||
|
||||
addr.IsPrivate() ||
|
||||
addr.IsLinkLocalUnicast() ||
|
||||
addr.IsLinkLocalMulticast() ||
|
||||
addr.IsMulticast() ||
|
||||
addr.IsUnspecified() ||
|
||||
addr.IsInterfaceLocalMulticast() {
|
||||
return true
|
||||
}
|
||||
|
||||
if addr.Is4() {
|
||||
if cgnatPrefix.Contains(addr) || benchmarkPrefix.Contains(addr) || reservedPrefix.Contains(addr) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
|
||||
return false
|
||||
}
|
||||
|
||||
func isBlockedHostname(host string) bool {
|
||||
if host == "localhost" || strings.HasSuffix(host, ".localhost") {
|
||||
return true
|
||||
}
|
||||
if strings.HasSuffix(host, ".local") || strings.HasSuffix(host, ".internal") {
|
||||
return true
|
||||
}
|
||||
if host == "metadata.google.internal" || host == "metadata" {
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func parseIPLiteral(host string) (netip.Addr, bool) {
|
||||
if i := strings.Index(host, "%"); i >= 0 {
|
||||
host = host[:i]
|
||||
}
|
||||
addr, err := netip.ParseAddr(host)
|
||||
if err != nil {
|
||||
return netip.Addr{}, false
|
||||
}
|
||||
return addr.Unmap(), true
|
||||
}
|
||||
|
||||
func normalizeHostToken(raw string) string {
|
||||
return strings.TrimSuffix(strings.ToLower(strings.TrimSpace(raw)), ".")
|
||||
}
|
||||
|
||||
func guardedDialContext(base *net.Dialer, validator *fetchTargetValidator) func(ctx context.Context, network, address string) (net.Conn, error) {
|
||||
if base == nil {
|
||||
base = &net.Dialer{}
|
||||
}
|
||||
|
||||
return func(ctx context.Context, network, address string) (net.Conn, error) {
|
||||
host, port, err := net.SplitHostPort(address)
|
||||
if err != nil {
|
||||
host = address
|
||||
port = ""
|
||||
}
|
||||
|
||||
target := &url.URL{Host: host}
|
||||
if port != "" {
|
||||
target.Host = net.JoinHostPort(host, port)
|
||||
}
|
||||
|
||||
if err := validator.validateURL(ctx, target); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return base.DialContext(ctx, network, address)
|
||||
}
|
||||
}
|
||||
|
|
@ -94,26 +94,53 @@ func (t *ExecTool) Execute(ctx context.Context, args map[string]interface{}) *To
|
|||
|
||||
var cmd *exec.Cmd
|
||||
if runtime.GOOS == "windows" {
|
||||
cmd = exec.CommandContext(cmdCtx, "powershell", "-NoProfile", "-NonInteractive", "-Command", command)
|
||||
cmd = exec.Command("powershell", "-NoProfile", "-NonInteractive", "-Command", command)
|
||||
} else {
|
||||
cmd = exec.CommandContext(cmdCtx, "sh", "-c", command)
|
||||
cmd = exec.Command("sh", "-c", command)
|
||||
}
|
||||
if cwd != "" {
|
||||
cmd.Dir = cwd
|
||||
}
|
||||
|
||||
prepareCommandForTreeControl(cmd)
|
||||
|
||||
var stdout, stderr bytes.Buffer
|
||||
cmd.Stdout = &stdout
|
||||
cmd.Stderr = &stderr
|
||||
|
||||
err := cmd.Run()
|
||||
if err := cmd.Start(); err != nil {
|
||||
return ErrorResult(fmt.Sprintf("failed to start command: %v", err))
|
||||
}
|
||||
|
||||
waitDone := make(chan error, 1)
|
||||
go func() {
|
||||
waitDone <- cmd.Wait()
|
||||
}()
|
||||
|
||||
var err error
|
||||
timedOut := false
|
||||
|
||||
select {
|
||||
case err = <-waitDone:
|
||||
case <-cmdCtx.Done():
|
||||
if killErr := killCommandTree(cmd); killErr != nil {
|
||||
return ErrorResult(fmt.Sprintf("failed to terminate command tree: %v", killErr))
|
||||
}
|
||||
|
||||
timedOut = cmdCtx.Err() == context.DeadlineExceeded
|
||||
select {
|
||||
case err = <-waitDone:
|
||||
case <-time.After(3 * time.Second):
|
||||
return ErrorResult("command termination timed out after deadline")
|
||||
}
|
||||
}
|
||||
|
||||
output := stdout.String()
|
||||
if stderr.Len() > 0 {
|
||||
output += "\nSTDERR:\n" + stderr.String()
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
if cmdCtx.Err() == context.DeadlineExceeded {
|
||||
if timedOut {
|
||||
msg := fmt.Sprintf("Command timed out after %v", t.timeout)
|
||||
return &ToolResult{
|
||||
ForLLM: msg,
|
||||
|
|
@ -121,6 +148,8 @@ func (t *ExecTool) Execute(ctx context.Context, args map[string]interface{}) *To
|
|||
IsError: true,
|
||||
}
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
output += fmt.Sprintf("\nExit code: %v", err)
|
||||
}
|
||||
|
||||
|
|
|
|||
37
pkg/tools/shell_process_unix.go
Normal file
37
pkg/tools/shell_process_unix.go
Normal file
|
|
@ -0,0 +1,37 @@
|
|||
//go:build !windows
|
||||
|
||||
package tools
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"os/exec"
|
||||
"syscall"
|
||||
)
|
||||
|
||||
func prepareCommandForTreeControl(cmd *exec.Cmd) {
|
||||
cmd.SysProcAttr = &syscall.SysProcAttr{
|
||||
Setpgid: true,
|
||||
}
|
||||
}
|
||||
|
||||
func killCommandTree(cmd *exec.Cmd) error {
|
||||
if cmd == nil || cmd.Process == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
pgid, err := syscall.Getpgid(cmd.Process.Pid)
|
||||
if err == nil {
|
||||
if killErr := syscall.Kill(-pgid, syscall.SIGKILL); killErr == nil || killErr == syscall.ESRCH {
|
||||
return nil
|
||||
} else {
|
||||
return fmt.Errorf("failed to kill process group %d: %w", pgid, killErr)
|
||||
}
|
||||
}
|
||||
|
||||
if killErr := cmd.Process.Kill(); killErr != nil && killErr != os.ErrProcessDone {
|
||||
return fmt.Errorf("failed to kill process %d: %w", cmd.Process.Pid, killErr)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
30
pkg/tools/shell_process_windows.go
Normal file
30
pkg/tools/shell_process_windows.go
Normal file
|
|
@ -0,0 +1,30 @@
|
|||
//go:build windows
|
||||
|
||||
package tools
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os/exec"
|
||||
"strconv"
|
||||
"syscall"
|
||||
)
|
||||
|
||||
func prepareCommandForTreeControl(cmd *exec.Cmd) {
|
||||
cmd.SysProcAttr = &syscall.SysProcAttr{
|
||||
CreationFlags: syscall.CREATE_NEW_PROCESS_GROUP,
|
||||
}
|
||||
}
|
||||
|
||||
func killCommandTree(cmd *exec.Cmd) error {
|
||||
if cmd == nil || cmd.Process == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
pid := strconv.Itoa(cmd.Process.Pid)
|
||||
killCmd := exec.Command("taskkill", "/T", "/F", "/PID", pid)
|
||||
if err := killCmd.Run(); err != nil {
|
||||
return fmt.Errorf("taskkill failed for pid %s: %w", pid, err)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
|
@ -2,8 +2,12 @@ package tools
|
|||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"runtime"
|
||||
"strconv"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
|
@ -208,3 +212,95 @@ func TestShellTool_RestrictToWorkspace(t *testing.T) {
|
|||
t.Errorf("Expected 'blocked' message for path traversal, got ForLLM: %s, ForUser: %s", result.ForLLM, result.ForUser)
|
||||
}
|
||||
}
|
||||
|
||||
// TestShellTool_Timeout_KillsChildProcesses verifies timeout cleanup includes child processes.
|
||||
func TestShellTool_Timeout_KillsChildProcesses(t *testing.T) {
|
||||
tool := NewExecTool("", false)
|
||||
tool.SetTimeout(1200 * time.Millisecond)
|
||||
|
||||
pidFile := filepath.Join(t.TempDir(), "child.pid")
|
||||
var cmd string
|
||||
if runtime.GOOS == "windows" {
|
||||
escapedPidFile := strings.ReplaceAll(pidFile, "'", "''")
|
||||
cmd = fmt.Sprintf(
|
||||
"$p = Start-Process -FilePath powershell -ArgumentList '-NoProfile','-NonInteractive','-Command','Start-Sleep -Seconds 30' -WindowStyle Hidden -PassThru; Set-Content -Path '%s' -Value $p.Id; Start-Sleep -Seconds 30",
|
||||
escapedPidFile,
|
||||
)
|
||||
} else {
|
||||
cmd = fmt.Sprintf("sleep 30 & echo $! > %q; sleep 30", pidFile)
|
||||
}
|
||||
|
||||
result := tool.Execute(context.Background(), map[string]interface{}{"command": cmd})
|
||||
if !result.IsError {
|
||||
t.Fatalf("Expected timeout error for long-running command tree")
|
||||
}
|
||||
if !strings.Contains(strings.ToLower(result.ForLLM), "timed out") {
|
||||
t.Fatalf("Expected timeout message, got: %s", result.ForLLM)
|
||||
}
|
||||
|
||||
childPID, err := waitForPID(pidFile, 4*time.Second)
|
||||
if err != nil {
|
||||
t.Fatalf("failed to obtain child pid file before timeout: %v", err)
|
||||
}
|
||||
|
||||
// Give timeout cleanup a short grace period before checking process liveness.
|
||||
time.Sleep(400 * time.Millisecond)
|
||||
alive := processAlive(childPID)
|
||||
if alive {
|
||||
status := processStatus(childPID)
|
||||
killProcess(childPID)
|
||||
t.Fatalf("expected child process %d to be terminated after timeout; status: %s", childPID, status)
|
||||
}
|
||||
}
|
||||
|
||||
func waitForPID(path string, timeout time.Duration) (int, error) {
|
||||
deadline := time.Now().Add(timeout)
|
||||
for time.Now().Before(deadline) {
|
||||
raw, err := os.ReadFile(path)
|
||||
if err == nil {
|
||||
pid, parseErr := strconv.Atoi(strings.TrimSpace(string(raw)))
|
||||
if parseErr == nil && pid > 0 {
|
||||
return pid, nil
|
||||
}
|
||||
}
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
}
|
||||
return 0, fmt.Errorf("pid file %s not ready within %v", path, timeout)
|
||||
}
|
||||
|
||||
func processAlive(pid int) bool {
|
||||
if runtime.GOOS == "windows" {
|
||||
err := exec.Command("powershell", "-NoProfile", "-NonInteractive", "-Command", fmt.Sprintf("Get-Process -Id %d | Out-Null", pid)).Run()
|
||||
return err == nil
|
||||
}
|
||||
|
||||
output, err := exec.Command("sh", "-c", fmt.Sprintf("ps -o stat= -p %d", pid)).CombinedOutput()
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
|
||||
state := strings.TrimSpace(string(output))
|
||||
if state == "" || strings.HasPrefix(state, "Z") {
|
||||
return false
|
||||
}
|
||||
|
||||
return true
|
||||
}
|
||||
|
||||
func processStatus(pid int) string {
|
||||
if runtime.GOOS == "windows" {
|
||||
output, _ := exec.Command("tasklist", "/FI", fmt.Sprintf("PID eq %d", pid)).CombinedOutput()
|
||||
return strings.TrimSpace(string(output))
|
||||
}
|
||||
|
||||
output, _ := exec.Command("sh", "-c", fmt.Sprintf("ps -o pid=,ppid=,pgid=,stat=,cmd= -p %d", pid)).CombinedOutput()
|
||||
return strings.TrimSpace(string(output))
|
||||
}
|
||||
|
||||
func killProcess(pid int) {
|
||||
if runtime.GOOS == "windows" {
|
||||
_ = exec.Command("taskkill", "/PID", strconv.Itoa(pid), "/T", "/F").Run()
|
||||
return
|
||||
}
|
||||
_ = exec.Command("kill", "-KILL", strconv.Itoa(pid)).Run()
|
||||
}
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@ import (
|
|||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"regexp"
|
||||
|
|
@ -267,14 +268,20 @@ func (t *WebSearchTool) Execute(ctx context.Context, args map[string]interface{}
|
|||
|
||||
type WebFetchTool struct {
|
||||
maxChars int
|
||||
validator *fetchTargetValidator
|
||||
}
|
||||
|
||||
func NewWebFetchTool(maxChars int) *WebFetchTool {
|
||||
return newWebFetchTool(maxChars, nil)
|
||||
}
|
||||
|
||||
func newWebFetchTool(maxChars int, allowHosts []string) *WebFetchTool {
|
||||
if maxChars <= 0 {
|
||||
maxChars = 50000
|
||||
}
|
||||
return &WebFetchTool{
|
||||
maxChars: maxChars,
|
||||
validator: newFetchTargetValidator(allowHosts, net.DefaultResolver),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -330,6 +337,10 @@ func (t *WebFetchTool) Execute(ctx context.Context, args map[string]interface{})
|
|||
}
|
||||
}
|
||||
|
||||
if err := t.validator.validateURL(ctx, parsedURL); err != nil {
|
||||
return ErrorResult(err.Error())
|
||||
}
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, "GET", urlStr, nil)
|
||||
if err != nil {
|
||||
return ErrorResult(fmt.Sprintf("failed to create request: %v", err))
|
||||
|
|
@ -337,6 +348,11 @@ func (t *WebFetchTool) Execute(ctx context.Context, args map[string]interface{})
|
|||
|
||||
req.Header.Set("User-Agent", userAgent)
|
||||
|
||||
dialer := &net.Dialer{
|
||||
Timeout: 30 * time.Second,
|
||||
KeepAlive: 30 * time.Second,
|
||||
}
|
||||
|
||||
client := &http.Client{
|
||||
Timeout: 60 * time.Second,
|
||||
Transport: &http.Transport{
|
||||
|
|
@ -344,11 +360,15 @@ func (t *WebFetchTool) Execute(ctx context.Context, args map[string]interface{})
|
|||
IdleConnTimeout: 30 * time.Second,
|
||||
DisableCompression: false,
|
||||
TLSHandshakeTimeout: 15 * time.Second,
|
||||
DialContext: guardedDialContext(dialer, t.validator),
|
||||
},
|
||||
CheckRedirect: func(req *http.Request, via []*http.Request) error {
|
||||
if len(via) >= 5 {
|
||||
return fmt.Errorf("stopped after 5 redirects")
|
||||
}
|
||||
if err := t.validator.validateURL(req.Context(), req.URL); err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
},
|
||||
}
|
||||
|
|
|
|||
|
|
@ -5,10 +5,17 @@ import (
|
|||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/url"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// newWebFetchToolForTests centralizes test construction.
|
||||
// allowHosts is wired in Phase 2 once fetch target policy is introduced.
|
||||
func newWebFetchToolForTests(maxChars int, allowHosts ...string) *WebFetchTool {
|
||||
return newWebFetchTool(maxChars, allowHosts)
|
||||
}
|
||||
|
||||
// TestWebTool_WebFetch_Success verifies successful URL fetching
|
||||
func TestWebTool_WebFetch_Success(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
|
|
@ -18,7 +25,11 @@ func TestWebTool_WebFetch_Success(t *testing.T) {
|
|||
}))
|
||||
defer server.Close()
|
||||
|
||||
tool := NewWebFetchTool(50000)
|
||||
parsed, err := url.Parse(server.URL)
|
||||
if err != nil {
|
||||
t.Fatalf("failed to parse test server URL: %v", err)
|
||||
}
|
||||
tool := newWebFetchToolForTests(50000, parsed.Host)
|
||||
ctx := context.Background()
|
||||
args := map[string]interface{}{
|
||||
"url": server.URL,
|
||||
|
|
@ -54,7 +65,11 @@ func TestWebTool_WebFetch_JSON(t *testing.T) {
|
|||
}))
|
||||
defer server.Close()
|
||||
|
||||
tool := NewWebFetchTool(50000)
|
||||
parsed, err := url.Parse(server.URL)
|
||||
if err != nil {
|
||||
t.Fatalf("failed to parse test server URL: %v", err)
|
||||
}
|
||||
tool := newWebFetchToolForTests(50000, parsed.Host)
|
||||
ctx := context.Background()
|
||||
args := map[string]interface{}{
|
||||
"url": server.URL,
|
||||
|
|
@ -145,7 +160,11 @@ func TestWebTool_WebFetch_Truncation(t *testing.T) {
|
|||
}))
|
||||
defer server.Close()
|
||||
|
||||
tool := NewWebFetchTool(1000) // Limit to 1000 chars
|
||||
parsed, err := url.Parse(server.URL)
|
||||
if err != nil {
|
||||
t.Fatalf("failed to parse test server URL: %v", err)
|
||||
}
|
||||
tool := newWebFetchToolForTests(1000, parsed.Host) // Limit to 1000 chars
|
||||
ctx := context.Background()
|
||||
args := map[string]interface{}{
|
||||
"url": server.URL,
|
||||
|
|
@ -206,7 +225,11 @@ func TestWebTool_WebFetch_HTMLExtraction(t *testing.T) {
|
|||
}))
|
||||
defer server.Close()
|
||||
|
||||
tool := NewWebFetchTool(50000)
|
||||
parsed, err := url.Parse(server.URL)
|
||||
if err != nil {
|
||||
t.Fatalf("failed to parse test server URL: %v", err)
|
||||
}
|
||||
tool := newWebFetchToolForTests(50000, parsed.Host)
|
||||
ctx := context.Background()
|
||||
args := map[string]interface{}{
|
||||
"url": server.URL,
|
||||
|
|
@ -250,3 +273,47 @@ func TestWebTool_WebFetch_MissingDomain(t *testing.T) {
|
|||
t.Errorf("Expected domain error message, got ForLLM: %s", result.ForLLM)
|
||||
}
|
||||
}
|
||||
|
||||
// TestWebTool_WebFetch_BlocksLoopback verifies loopback targets are blocked
|
||||
func TestWebTool_WebFetch_BlocksLoopback(t *testing.T) {
|
||||
tool := NewWebFetchTool(50000)
|
||||
ctx := context.Background()
|
||||
args := map[string]interface{}{
|
||||
"url": "http://127.0.0.1/",
|
||||
}
|
||||
|
||||
result := tool.Execute(ctx, args)
|
||||
if !result.IsError {
|
||||
t.Fatalf("Expected blocked destination error for loopback target")
|
||||
}
|
||||
if !strings.Contains(strings.ToLower(result.ForLLM), "blocked destination") {
|
||||
t.Fatalf("Expected blocked destination message, got: %s", result.ForLLM)
|
||||
}
|
||||
}
|
||||
|
||||
// TestWebTool_WebFetch_BlocksRedirectToPrivate verifies private redirect hops are blocked
|
||||
func TestWebTool_WebFetch_BlocksRedirectToPrivate(t *testing.T) {
|
||||
redirectServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
http.Redirect(w, r, "http://127.0.0.1:1/private", http.StatusFound)
|
||||
}))
|
||||
defer redirectServer.Close()
|
||||
|
||||
parsed, err := url.Parse(redirectServer.URL)
|
||||
if err != nil {
|
||||
t.Fatalf("failed to parse test server URL: %v", err)
|
||||
}
|
||||
|
||||
tool := newWebFetchToolForTests(50000, parsed.Host)
|
||||
ctx := context.Background()
|
||||
args := map[string]interface{}{
|
||||
"url": redirectServer.URL,
|
||||
}
|
||||
|
||||
result := tool.Execute(ctx, args)
|
||||
if !result.IsError {
|
||||
t.Fatalf("Expected blocked destination error for redirect to private target")
|
||||
}
|
||||
if !strings.Contains(strings.ToLower(result.ForLLM), "blocked destination") {
|
||||
t.Fatalf("Expected blocked destination message, got: %s", result.ForLLM)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue