Skip to content

Commit 52e9812

Browse files
authored
fix(security): workspace sandbox avoid time-of-check/time-of-use (TOCTOU) races (sipeed#464)
* chore: Update default host bindings from 0.0.0.0 to 127.0.0.1 for various services and examples. * config: Update default host bindings to 0.0.0.0 for improved Docker accessibility and add related documentation. * refactor: reimplement filesystem tools with `os.OpenRoot` for enhanced security and simplified path validation. * chore: revert other PR content from this branch * docs: Update Chinese README. * docs: Update Chinese README. * docs: Update Chinese README. * refactor: Reorder filesystem helper functions, extract directory entry formatting logic, and enhance `WriteFileTool`'s result message. * feat: Enhance `mkdirAllInRoot` to prevent creating directories over existing files and add tests for directory creation functionality. * Refactor filesystem tools to use a `fileReadWriter` interface for both host and sandboxed I/O, improving atomic writes and error handling. * refactor: unify filesystem read/write operations with atomic write guarantees and clearer naming. * refactor: rename `appendFileWithRW` function to `appendFile` * refactor: unify filesystem access by introducing a `fileSystem` interface and updating tools to use it directly, removing `os.Root` dependency from `sandboxFs`. * chore: run make fmt * fix: `validatePath` now returns an error when the workspace is empty.
1 parent d1c461a commit 52e9812

File tree

4 files changed

+626
-105
lines changed

4 files changed

+626
-105
lines changed

pkg/tools/edit.go

Lines changed: 60 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -2,24 +2,27 @@ package tools
22

33
import (
44
"context"
5+
"errors"
56
"fmt"
6-
"os"
7+
"io/fs"
78
"strings"
89
)
910

1011
// EditFileTool edits a file by replacing old_text with new_text.
1112
// The old_text must exist exactly in the file.
1213
type EditFileTool struct {
13-
allowedDir string
14-
restrict bool
14+
fs fileSystem
1515
}
1616

1717
// NewEditFileTool creates a new EditFileTool with optional directory restriction.
18-
func NewEditFileTool(allowedDir string, restrict bool) *EditFileTool {
19-
return &EditFileTool{
20-
allowedDir: allowedDir,
21-
restrict: restrict,
22-
}
18+
func NewEditFileTool(workspace string, restrict bool) *EditFileTool {
19+
var fs fileSystem
20+
if restrict {
21+
fs = &sandboxFs{workspace: workspace}
22+
} else {
23+
fs = &hostFs{}
24+
}
25+
return &EditFileTool{fs: fs}
2326
}
2427

2528
func (t *EditFileTool) Name() string {
@@ -67,49 +70,24 @@ func (t *EditFileTool) Execute(ctx context.Context, args map[string]any) *ToolRe
6770
return ErrorResult("new_text is required")
6871
}
6972

70-
resolvedPath, err := validatePath(path, t.allowedDir, t.restrict)
71-
if err != nil {
73+
if err := editFile(t.fs, path, oldText, newText); err != nil {
7274
return ErrorResult(err.Error())
7375
}
74-
75-
if _, err = os.Stat(resolvedPath); os.IsNotExist(err) {
76-
return ErrorResult(fmt.Sprintf("file not found: %s", path))
77-
}
78-
79-
content, err := os.ReadFile(resolvedPath)
80-
if err != nil {
81-
return ErrorResult(fmt.Sprintf("failed to read file: %v", err))
82-
}
83-
84-
contentStr := string(content)
85-
86-
if !strings.Contains(contentStr, oldText) {
87-
return ErrorResult("old_text not found in file. Make sure it matches exactly")
88-
}
89-
90-
count := strings.Count(contentStr, oldText)
91-
if count > 1 {
92-
return ErrorResult(
93-
fmt.Sprintf("old_text appears %d times. Please provide more context to make it unique", count),
94-
)
95-
}
96-
97-
newContent := strings.Replace(contentStr, oldText, newText, 1)
98-
99-
if err := os.WriteFile(resolvedPath, []byte(newContent), 0o644); err != nil {
100-
return ErrorResult(fmt.Sprintf("failed to write file: %v", err))
101-
}
102-
10376
return SilentResult(fmt.Sprintf("File edited: %s", path))
10477
}
10578

10679
type AppendFileTool struct {
107-
workspace string
108-
restrict bool
80+
fs fileSystem
10981
}
11082

11183
func NewAppendFileTool(workspace string, restrict bool) *AppendFileTool {
112-
return &AppendFileTool{workspace: workspace, restrict: restrict}
84+
var fs fileSystem
85+
if restrict {
86+
fs = &sandboxFs{workspace: workspace}
87+
} else {
88+
fs = &hostFs{}
89+
}
90+
return &AppendFileTool{fs: fs}
11391
}
11492

11593
func (t *AppendFileTool) Name() string {
@@ -148,20 +126,52 @@ func (t *AppendFileTool) Execute(ctx context.Context, args map[string]any) *Tool
148126
return ErrorResult("content is required")
149127
}
150128

151-
resolvedPath, err := validatePath(path, t.workspace, t.restrict)
152-
if err != nil {
129+
if err := appendFile(t.fs, path, content); err != nil {
153130
return ErrorResult(err.Error())
154131
}
132+
return SilentResult(fmt.Sprintf("Appended to %s", path))
133+
}
155134

156-
f, err := os.OpenFile(resolvedPath, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0o644)
135+
// editFile reads the file via sysFs, performs the replacement, and writes back.
136+
// It uses a fileSystem interface, allowing the same logic for both restricted and unrestricted modes.
137+
func editFile(sysFs fileSystem, path, oldText, newText string) error {
138+
content, err := sysFs.ReadFile(path)
157139
if err != nil {
158-
return ErrorResult(fmt.Sprintf("failed to open file: %v", err))
140+
return err
159141
}
160-
defer f.Close()
161142

162-
if _, err := f.WriteString(content); err != nil {
163-
return ErrorResult(fmt.Sprintf("failed to append to file: %v", err))
143+
newContent, err := replaceEditContent(content, oldText, newText)
144+
if err != nil {
145+
return err
164146
}
165147

166-
return SilentResult(fmt.Sprintf("Appended to %s", path))
148+
return sysFs.WriteFile(path, newContent)
149+
}
150+
151+
// appendFile reads the existing content (if any) via sysFs, appends new content, and writes back.
152+
func appendFile(sysFs fileSystem, path, appendContent string) error {
153+
content, err := sysFs.ReadFile(path)
154+
if err != nil && !errors.Is(err, fs.ErrNotExist) {
155+
return err
156+
}
157+
158+
newContent := append(content, []byte(appendContent)...)
159+
return sysFs.WriteFile(path, newContent)
160+
}
161+
162+
// replaceEditContent handles the core logic of finding and replacing a single occurrence of oldText.
163+
func replaceEditContent(content []byte, oldText, newText string) ([]byte, error) {
164+
contentStr := string(content)
165+
166+
if !strings.Contains(contentStr, oldText) {
167+
return nil, fmt.Errorf("old_text not found in file. Make sure it matches exactly")
168+
}
169+
170+
count := strings.Count(contentStr, oldText)
171+
if count > 1 {
172+
return nil, fmt.Errorf("old_text appears %d times. Please provide more context to make it unique", count)
173+
}
174+
175+
newContent := strings.Replace(contentStr, oldText, newText, 1)
176+
return []byte(newContent), nil
167177
}

pkg/tools/edit_test.go

Lines changed: 154 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,8 @@ import (
66
"path/filepath"
77
"strings"
88
"testing"
9+
10+
"github.com/stretchr/testify/assert"
911
)
1012

1113
// TestEditTool_EditFile_Success verifies successful file editing
@@ -151,14 +153,18 @@ func TestEditTool_EditFile_OutsideAllowedDir(t *testing.T) {
151153
result := tool.Execute(ctx, args)
152154

153155
// Should return error result
154-
if !result.IsError {
155-
t.Errorf("Expected error when path is outside allowed directory")
156-
}
156+
assert.True(t, result.IsError, "Expected error when path is outside allowed directory")
157157

158158
// Should mention outside allowed directory
159-
if !strings.Contains(result.ForLLM, "outside") && !strings.Contains(result.ForUser, "outside") {
160-
t.Errorf("Expected 'outside allowed' message, got ForLLM: %s", result.ForLLM)
161-
}
159+
// Note: ErrorResult only sets ForLLM by default, so ForUser might be empty.
160+
// We check ForLLM as it's the primary error channel.
161+
assert.True(
162+
t,
163+
strings.Contains(result.ForLLM, "outside") || strings.Contains(result.ForLLM, "access denied") ||
164+
strings.Contains(result.ForLLM, "escapes"),
165+
"Expected 'outside allowed' or 'access denied' message, got ForLLM: %s",
166+
result.ForLLM,
167+
)
162168
}
163169

164170
// TestEditTool_EditFile_MissingPath verifies error handling for missing path
@@ -287,3 +293,145 @@ func TestEditTool_AppendFile_MissingContent(t *testing.T) {
287293
t.Errorf("Expected error when content is missing")
288294
}
289295
}
296+
297+
// TestReplaceEditContent verifies the helper function replaceEditContent
298+
func TestReplaceEditContent(t *testing.T) {
299+
tests := []struct {
300+
name string
301+
content []byte
302+
oldText string
303+
newText string
304+
expected []byte
305+
expectError bool
306+
}{
307+
{
308+
name: "successful replacement",
309+
content: []byte("hello world"),
310+
oldText: "world",
311+
newText: "universe",
312+
expected: []byte("hello universe"),
313+
expectError: false,
314+
},
315+
{
316+
name: "old text not found",
317+
content: []byte("hello world"),
318+
oldText: "golang",
319+
newText: "rust",
320+
expected: nil,
321+
expectError: true,
322+
},
323+
{
324+
name: "multiple matches found",
325+
content: []byte("test text test"),
326+
oldText: "test",
327+
newText: "done",
328+
expected: nil,
329+
expectError: true,
330+
},
331+
}
332+
333+
for _, tt := range tests {
334+
t.Run(tt.name, func(t *testing.T) {
335+
result, err := replaceEditContent(tt.content, tt.oldText, tt.newText)
336+
if tt.expectError {
337+
assert.Error(t, err)
338+
} else {
339+
assert.NoError(t, err)
340+
assert.Equal(t, tt.expected, result)
341+
}
342+
})
343+
}
344+
}
345+
346+
// TestAppendFileTool_AppendToNonExistent_Restricted verifies that AppendFileTool in restricted mode
347+
// can append to a file that does not yet exist — it should silently create the file.
348+
// This exercises the errors.Is(err, fs.ErrNotExist) path in appendFileWithRW + rootRW.
349+
func TestAppendFileTool_AppendToNonExistent_Restricted(t *testing.T) {
350+
workspace := t.TempDir()
351+
tool := NewAppendFileTool(workspace, true)
352+
ctx := context.Background()
353+
354+
args := map[string]any{
355+
"path": "brand_new_file.txt",
356+
"content": "first content",
357+
}
358+
359+
result := tool.Execute(ctx, args)
360+
assert.False(
361+
t,
362+
result.IsError,
363+
"Expected success when appending to non-existent file in restricted mode, got: %s",
364+
result.ForLLM,
365+
)
366+
367+
// Verify the file was created with correct content
368+
data, err := os.ReadFile(filepath.Join(workspace, "brand_new_file.txt"))
369+
assert.NoError(t, err)
370+
assert.Equal(t, "first content", string(data))
371+
}
372+
373+
// TestAppendFileTool_Restricted_Success verifies that AppendFileTool in restricted mode
374+
// correctly appends to an existing file within the sandbox.
375+
func TestAppendFileTool_Restricted_Success(t *testing.T) {
376+
workspace := t.TempDir()
377+
testFile := "existing.txt"
378+
err := os.WriteFile(filepath.Join(workspace, testFile), []byte("initial"), 0o644)
379+
assert.NoError(t, err)
380+
381+
tool := NewAppendFileTool(workspace, true)
382+
ctx := context.Background()
383+
args := map[string]any{
384+
"path": testFile,
385+
"content": " appended",
386+
}
387+
388+
result := tool.Execute(ctx, args)
389+
assert.False(t, result.IsError, "Expected success, got: %s", result.ForLLM)
390+
assert.True(t, result.Silent)
391+
392+
data, err := os.ReadFile(filepath.Join(workspace, testFile))
393+
assert.NoError(t, err)
394+
assert.Equal(t, "initial appended", string(data))
395+
}
396+
397+
// TestEditFileTool_Restricted_InPlaceEdit verifies that EditFileTool in restricted mode
398+
// correctly edits a file using the single-open editFileInRoot path.
399+
func TestEditFileTool_Restricted_InPlaceEdit(t *testing.T) {
400+
workspace := t.TempDir()
401+
testFile := "edit_target.txt"
402+
err := os.WriteFile(filepath.Join(workspace, testFile), []byte("Hello World"), 0o644)
403+
assert.NoError(t, err)
404+
405+
tool := NewEditFileTool(workspace, true)
406+
ctx := context.Background()
407+
args := map[string]any{
408+
"path": testFile,
409+
"old_text": "World",
410+
"new_text": "Go",
411+
}
412+
413+
result := tool.Execute(ctx, args)
414+
assert.False(t, result.IsError, "Expected success, got: %s", result.ForLLM)
415+
assert.True(t, result.Silent)
416+
417+
data, err := os.ReadFile(filepath.Join(workspace, testFile))
418+
assert.NoError(t, err)
419+
assert.Equal(t, "Hello Go", string(data))
420+
}
421+
422+
// TestEditFileTool_Restricted_FileNotFound verifies that editFileInRoot returns a proper
423+
// error message when the target file does not exist.
424+
func TestEditFileTool_Restricted_FileNotFound(t *testing.T) {
425+
workspace := t.TempDir()
426+
tool := NewEditFileTool(workspace, true)
427+
ctx := context.Background()
428+
args := map[string]any{
429+
"path": "no_such_file.txt",
430+
"old_text": "old",
431+
"new_text": "new",
432+
}
433+
434+
result := tool.Execute(ctx, args)
435+
assert.True(t, result.IsError)
436+
assert.Contains(t, result.ForLLM, "not found")
437+
}

0 commit comments

Comments
 (0)