Skip to content

Commit 3058016

Browse files
authored
refactor(tool): extract safety helpers into internal/toolsafety (#315)
Move the self-contained safety layer (destructive-command detection, credential detection, binary detection, output truncation, SSRF guards, command segmentation) out of the 41k-LOC internal/tool package into a new low-level internal/toolsafety package. tool re-exports the symbols so existing callers keep compiling, and internal/engine/search now imports toolsafety directly instead of the god package. internal/tool loses ~500 lines of implementation and ~660 lines of tests.
1 parent 22e0273 commit 3058016

9 files changed

Lines changed: 341 additions & 268 deletions

File tree

‎internal/engine/search/url_scraper.go‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ import (
1313
"sync"
1414
"time"
1515

16-
"github.com/GrayCodeAI/rho/internal/tool"
16+
"github.com/GrayCodeAI/rho/internal/toolsafety"
1717
)
1818

1919
// URLScraper detects URLs in conversation text and fetches/extracts their content.
@@ -113,7 +113,7 @@ func (s *URLScraper) Fetch(ctx context.Context, rawURL string) (*ScrapeResult, e
113113
// SSRF guard: reject private/link-local targets and pin the resolved IP
114114
// to prevent DNS rebinding. Without this, a scraped URL could reach cloud
115115
// metadata endpoints or internal services.
116-
pinnedURL, origHost, err := tool.ValidateURLPublic(ctx, rawURL)
116+
pinnedURL, origHost, err := toolsafety.ValidateURLPublic(ctx, rawURL)
117117
if err != nil {
118118
return nil, err
119119
}
@@ -128,7 +128,7 @@ func (s *URLScraper) Fetch(ctx context.Context, rawURL string) (*ScrapeResult, e
128128
req.Header.Set("User-Agent", s.UserAgent)
129129
req.Header.Set("Accept", "text/html, application/json, text/plain, */*")
130130

131-
client := tool.SSRFSafeClient(ctx, s.Timeout)
131+
client := toolsafety.SSRFSafeClient(ctx, s.Timeout)
132132
resp, err := client.Do(req)
133133
if err != nil {
134134
return nil, fmt.Errorf("fetching URL: %w", err)

‎internal/engine/search/url_scraper_test.go‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ import (
1010
"time"
1111

1212
"github.com/GrayCodeAI/rho/internal/testutil"
13-
"github.com/GrayCodeAI/rho/internal/tool"
13+
"github.com/GrayCodeAI/rho/internal/toolsafety"
1414
)
1515

1616
func TestDetectURLs_HTTPAndHTTPS(t *testing.T) {
@@ -330,7 +330,7 @@ func TestCachePreventsRefetch(t *testing.T) {
330330
defer ts.Close()
331331

332332
scraper := NewURLScraper()
333-
ctx := tool.WithSSRFSkip(context.Background())
333+
ctx := toolsafety.WithSSRFSkip(context.Background())
334334

335335
// First fetch should hit server.
336336
_, err := scraper.Fetch(ctx, ts.URL)
@@ -359,7 +359,7 @@ func TestFetch_HTMLContent(t *testing.T) {
359359
defer ts.Close()
360360

361361
scraper := NewURLScraper()
362-
result, err := scraper.Fetch(tool.WithSSRFSkip(context.Background()), ts.URL)
362+
result, err := scraper.Fetch(toolsafety.WithSSRFSkip(context.Background()), ts.URL)
363363
if err != nil {
364364
t.Fatalf("fetch failed: %v", err)
365365
}
@@ -386,7 +386,7 @@ func TestFetch_JSONContent(t *testing.T) {
386386
defer ts.Close()
387387

388388
scraper := NewURLScraper()
389-
result, err := scraper.Fetch(tool.WithSSRFSkip(context.Background()), ts.URL)
389+
result, err := scraper.Fetch(toolsafety.WithSSRFSkip(context.Background()), ts.URL)
390390
if err != nil {
391391
t.Fatalf("fetch failed: %v", err)
392392
}
@@ -494,7 +494,7 @@ func TestTokenEstimate(t *testing.T) {
494494
defer ts.Close()
495495

496496
scraper := NewURLScraper()
497-
result, err := scraper.Fetch(tool.WithSSRFSkip(context.Background()), ts.URL)
497+
result, err := scraper.Fetch(toolsafety.WithSSRFSkip(context.Background()), ts.URL)
498498
if err != nil {
499499
t.Fatalf("fetch failed: %v", err)
500500
}

‎internal/tool/bash.go‎

Lines changed: 5 additions & 117 deletions
Original file line numberDiff line numberDiff line change
@@ -180,118 +180,6 @@ func (BashTool) Parameters() map[string]interface{} {
180180
// bashSchema is the single source of truth for Bash's input schema.
181181
var bashSchema = BashTool{}.Schema()
182182

183-
// SegmentCommand splits a command string on &&, ||, ;, and | (respecting quotes
184-
// and heredocs) into individual segments for independent analysis.
185-
func SegmentCommand(cmd string) []string {
186-
var segments []string
187-
var current strings.Builder
188-
inSingle, inDouble := false, false
189-
inHeredoc := false
190-
heredocDelim := ""
191-
runes := []rune(cmd)
192-
for i := 0; i < len(runes); i++ {
193-
ch := runes[i]
194-
195-
// If inside a heredoc body, consume until we find the delimiter on its own line
196-
if inHeredoc {
197-
current.WriteRune(ch)
198-
if ch == '\n' {
199-
lineStart := i + 1
200-
lineEnd := lineStart
201-
for lineEnd < len(runes) && runes[lineEnd] != '\n' {
202-
lineEnd++
203-
}
204-
line := strings.TrimSpace(string(runes[lineStart:lineEnd]))
205-
if line == heredocDelim {
206-
for j := lineStart; j <= lineEnd && j < len(runes); j++ {
207-
current.WriteRune(runes[j])
208-
}
209-
i = lineEnd
210-
inHeredoc = false
211-
heredocDelim = ""
212-
}
213-
}
214-
continue
215-
}
216-
217-
if ch == '\'' && !inDouble {
218-
inSingle = !inSingle
219-
current.WriteRune(ch)
220-
continue
221-
}
222-
if ch == '"' && !inSingle {
223-
inDouble = !inDouble
224-
current.WriteRune(ch)
225-
continue
226-
}
227-
if inSingle || inDouble {
228-
current.WriteRune(ch)
229-
continue
230-
}
231-
232-
// Detect heredoc: <<EOF, << EOF, <<-EOF
233-
if ch == '<' && i+1 < len(runes) && runes[i+1] == '<' {
234-
j := i + 2
235-
if j < len(runes) && runes[j] == '-' {
236-
j++
237-
}
238-
for j < len(runes) && runes[j] == ' ' {
239-
j++
240-
}
241-
if j < len(runes) {
242-
delimStart := j
243-
delimQuote := rune(0)
244-
if runes[j] == '\'' || runes[j] == '"' {
245-
delimQuote = runes[j]
246-
j++
247-
delimStart = j
248-
for j < len(runes) && runes[j] != delimQuote {
249-
j++
250-
}
251-
} else {
252-
for j < len(runes) && runes[j] != ' ' && runes[j] != '\n' && runes[j] != '<' && runes[j] != '>' && runes[j] != '|' && runes[j] != '&' && runes[j] != ';' {
253-
j++
254-
}
255-
}
256-
if j > delimStart {
257-
heredocDelim = string(runes[delimStart:j])
258-
inHeredoc = true
259-
for k := i; k < j; k++ {
260-
current.WriteRune(runes[k])
261-
}
262-
i = j - 1
263-
continue
264-
}
265-
}
266-
}
267-
268-
// Check for &&, ||
269-
if i+1 < len(runes) && ((ch == '&' && runes[i+1] == '&') || (ch == '|' && runes[i+1] == '|')) {
270-
if s := strings.TrimSpace(current.String()); s != "" {
271-
segments = append(segments, s)
272-
}
273-
current.Reset()
274-
i++ // skip second char
275-
continue
276-
}
277-
// Check for ; or single |
278-
if ch == ';' || ch == '|' {
279-
if s := strings.TrimSpace(current.String()); s != "" {
280-
segments = append(segments, s)
281-
}
282-
current.Reset()
283-
continue
284-
}
285-
current.WriteRune(ch)
286-
}
287-
if s := strings.TrimSpace(current.String()); s != "" {
288-
segments = append(segments, s)
289-
}
290-
return segments
291-
}
292-
293-
// IsSuspicious returns true if the command needs a permission prompt.
294-
// This is fail-closed: anything we can't confidently classify as safe gets flagged.
295183
func IsSuspicious(command string) bool {
296184
// Whole-command checks that apply regardless of segmentation
297185
if strings.Contains(command, "\r") {
@@ -500,23 +388,23 @@ func (BashTool) Execute(ctx context.Context, input json.RawMessage) (string, err
500388
if tc := GetToolContext(ctx); tc != nil && tc.WorkingDir != "" {
501389
cmd.Dir = tc.WorkingDir
502390
}
503-
// Use a limitedWriter to cap output at maxOutputBytes instead of
391+
// Use a limitedWriter to cap output at MaxOutputBytes instead of
504392
// CombinedOutput, which buffers the entire output in memory. A command
505393
// like `yes` or `cat /dev/urandom` can produce GBs before the timeout
506394
// kills it; the limitedWriter keeps memory bounded while the command
507395
// continues to run (writes are silently discarded after the cap).
508396
var lw limitedWriter
509-
// Cap one byte above maxOutputBytes so that TruncateOutput's > branch
510-
// fires when the cap is reached. At exactly maxOutputBytes (no discard)
397+
// Cap one byte above MaxOutputBytes so that TruncateOutput's > branch
398+
// fires when the cap is reached. At exactly MaxOutputBytes (no discard)
511399
// TruncateOutput returns unchanged, which is correct.
512-
lw.maxBytes = maxOutputBytes + 1
400+
lw.maxBytes = MaxOutputBytes + 1
513401
cmd.Stdout = &lw
514402
cmd.Stderr = &lw
515403
err := cmd.Run()
516404
result := lw.buf.String()
517405

518406
// Apply safety output truncation (50KB) — the limitedWriter may have
519-
// captured up to maxOutputBytes (500KB), so we still truncate for the
407+
// captured up to MaxOutputBytes (500KB), so we still truncate for the
520408
// final result returned to the model.
521409
result = TruncateOutput(result)
522410
result = strings.TrimRight(result, "\n")

‎internal/tool/download.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -50,12 +50,12 @@ func (DownloadTool) Execute(ctx context.Context, input json.RawMessage) (string,
5050
if reason := IsSensitivePath(p.Destination); reason != "" {
5151
return "", fmt.Errorf("write blocked: %s", reason)
5252
}
53-
pinnedURL, origHost, err := validateURLPublic(ctx, p.URL)
53+
pinnedURL, origHost, err := ValidateURLPublic(ctx, p.URL)
5454
if err != nil {
5555
return "", err
5656
}
5757

58-
client := ssrfSafeClient(ctx, 2*time.Minute)
58+
client := SSRFSafeClient(ctx, 2*time.Minute)
5959
req, err := http.NewRequestWithContext(ctx, http.MethodGet, pinnedURL, nil)
6060
if err != nil {
6161
return "", fmt.Errorf("create request: %w", err)
Lines changed: 137 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,137 @@
1+
package tool
2+
3+
import (
4+
"encoding/json"
5+
"os"
6+
"path/filepath"
7+
"strings"
8+
"testing"
9+
)
10+
11+
// These tests exercise the file tools' use of the sensitive-path policy, so
12+
// they live in package tool (the policy itself is tested in toolsafety).
13+
14+
func TestFileToolsBlockFluxProviderConfig(t *testing.T) {
15+
fluxDir := filepath.Join(t.TempDir(), "flux")
16+
t.Setenv("FLUX_CONFIG_DIR", fluxDir)
17+
providerPath := filepath.Join(fluxDir, "provider.json")
18+
19+
readInput, _ := json.Marshal(map[string]string{"path": providerPath})
20+
editInput, _ := json.Marshal(map[string]string{
21+
"path": providerPath, "old_str": "old", "new_str": "new",
22+
})
23+
writeInput, _ := json.Marshal(map[string]string{
24+
"path": providerPath, "content": "safe routing metadata",
25+
})
26+
tests := []struct {
27+
name string
28+
run func() error
29+
}{
30+
{name: "Read", run: func() error {
31+
_, err := (FileReadTool{}).Execute(testCtx(), readInput)
32+
return err
33+
}},
34+
{name: "Edit", run: func() error {
35+
_, err := (FileEditTool{}).Execute(testCtx(), editInput)
36+
return err
37+
}},
38+
{name: "Write", run: func() error {
39+
_, err := (FileWriteTool{}).Execute(testCtx(), writeInput)
40+
return err
41+
}},
42+
}
43+
44+
for _, tt := range tests {
45+
t.Run(tt.name, func(t *testing.T) {
46+
err := tt.run()
47+
if err == nil || !strings.Contains(err.Error(), "blocked") {
48+
t.Fatalf("%s error = %v, want sensitive-path block", tt.name, err)
49+
}
50+
})
51+
}
52+
}
53+
54+
// TestFileRead_BlocksSymlinkToSensitiveFile verifies the read tool resolves
55+
// symlinks before opening (M13): reading through a symlink that points at a
56+
// sensitive target is blocked, while a symlink to an ordinary file works.
57+
func TestFileRead_BlocksSymlinkToSensitiveFile(t *testing.T) {
58+
fluxDir := filepath.Join(t.TempDir(), "flux")
59+
if err := os.MkdirAll(fluxDir, 0o755); err != nil {
60+
t.Fatal(err)
61+
}
62+
t.Setenv("FLUX_CONFIG_DIR", fluxDir)
63+
providerPath := filepath.Join(fluxDir, "provider.json")
64+
if err := os.WriteFile(providerPath, []byte(`{"key":"x"}`), 0o600); err != nil {
65+
t.Fatal(err)
66+
}
67+
68+
workDir := t.TempDir()
69+
link := filepath.Join(workDir, "readme.md")
70+
if err := os.Symlink(providerPath, link); err != nil {
71+
t.Fatal(err)
72+
}
73+
in, _ := json.Marshal(map[string]string{"path": link})
74+
_, err := (FileReadTool{}).Execute(testCtx(), in)
75+
if err == nil || !strings.Contains(err.Error(), "blocked") {
76+
t.Fatalf("expected sensitive-path block for symlinked provider config, got %v", err)
77+
}
78+
79+
// A symlink to an ordinary file must still read fine.
80+
plain := filepath.Join(workDir, "plain.txt")
81+
if err := os.WriteFile(plain, []byte("hello"), 0o600); err != nil {
82+
t.Fatal(err)
83+
}
84+
link2 := filepath.Join(workDir, "link2.txt")
85+
if err := os.Symlink(plain, link2); err != nil {
86+
t.Fatal(err)
87+
}
88+
in2, _ := json.Marshal(map[string]string{"path": link2})
89+
out, err := (FileReadTool{}).Execute(testCtx(), in2)
90+
if err != nil {
91+
t.Fatalf("expected symlinked plain file to read, got %v", err)
92+
}
93+
if !strings.Contains(out, "hello") {
94+
t.Fatalf("expected content through symlink, got %q", out)
95+
}
96+
}
97+
98+
// TestIsSensitivePath_SecretBasenames verifies the expanded basename blocklist
99+
// (secrets.txt, .git-credentials, private keys, …) applies anywhere, not just
100+
// under the home directory.
101+
102+
func TestBashSensitivePathIntegration(t *testing.T) {
103+
if !IsSuspicious("cat ~/.ssh/id_rsa") {
104+
t.Error("IsSuspicious should flag reads of SSH private keys")
105+
}
106+
if !isHardDeny("cat ~/.aws/credentials") {
107+
t.Error("isHardDeny should block credential reads when prompts are bypassed")
108+
}
109+
if isHardDeny("go build ./...") {
110+
t.Error("isHardDeny should not block ordinary commands")
111+
}
112+
}
113+
114+
func TestCommandReferencesSensitivePath_FluxConfigDir(t *testing.T) {
115+
fluxDir := filepath.Join(t.TempDir(), "flux config with spaces")
116+
t.Setenv("FLUX_CONFIG_DIR", fluxDir)
117+
118+
commands := []string{
119+
`cat "` + filepath.Join(fluxDir, "provider.json") + `"`,
120+
`cat "$FLUX_CONFIG_DIR/provider.json"`,
121+
`cat "${FLUX_CONFIG_DIR}/provider.json"`,
122+
`cat "${FLUX_CONFIG_DIR%/}/provider.json"`,
123+
`printf '%s\n' "$FLUX_CONFIG_DIR"`,
124+
"cat " + strings.ReplaceAll(filepath.Join(fluxDir, "provider.json"), " ", `\ `),
125+
}
126+
for _, command := range commands {
127+
if reason := CommandReferencesSensitivePath(command); reason == "" {
128+
t.Fatalf("CommandReferencesSensitivePath(%q) = empty, want blocked", command)
129+
}
130+
if !IsSuspicious(command) {
131+
t.Fatalf("IsSuspicious(%q) = false, want Bash prompt", command)
132+
}
133+
if !isHardDeny(command) {
134+
t.Fatalf("isHardDeny(%q) = false, want Bash block", command)
135+
}
136+
}
137+
}

0 commit comments

Comments
 (0)