fix(security): workspace sandbox avoid time-of-check/time-of-use (TOCTOU) races

pkg/tools/edit.go

@@ -2,24 +2,27 @@ package tools
 
 import (
 	"context"
+	"errors"
 	"fmt"
-	"os"
+	"io/fs"
 	"strings"
 )
 
 // EditFileTool edits a file by replacing old_text with new_text.
 // The old_text must exist exactly in the file.
 type EditFileTool struct {
-	allowedDir string
-	restrict   bool
+	fs fileSystem
 }
 
 // NewEditFileTool creates a new EditFileTool with optional directory restriction.
-func NewEditFileTool(allowedDir string, restrict bool) *EditFileTool {
-	return &EditFileTool{
-		allowedDir: allowedDir,
-		restrict:   restrict,
-	}
+func NewEditFileTool(workspace string, restrict bool) *EditFileTool {
+	var fs fileSystem
+	if restrict {
+		fs = &sandboxFs{workspace: workspace}
+	} else {
+		fs = &hostFs{}
+	}
+	return &EditFileTool{fs: fs}
 }
 
 func (t *EditFileTool) Name() string {
@@ -67,49 +70,24 @@ func (t *EditFileTool) Execute(ctx context.Context, args map[string]any) *ToolRe
 		return ErrorResult("new_text is required")
 	}
 
-	resolvedPath, err := validatePath(path, t.allowedDir, t.restrict)
-	if err != nil {
+	if err := editFile(t.fs, path, oldText, newText); err != nil {
 		return ErrorResult(err.Error())
 	}
-
-	if _, err = os.Stat(resolvedPath); os.IsNotExist(err) {
-		return ErrorResult(fmt.Sprintf("file not found: %s", path))
-	}
-
-	content, err := os.ReadFile(resolvedPath)
-	if err != nil {
-		return ErrorResult(fmt.Sprintf("failed to read file: %v", err))
-	}
-
-	contentStr := string(content)
-
-	if !strings.Contains(contentStr, oldText) {
-		return ErrorResult("old_text not found in file. Make sure it matches exactly")
-	}
-
-	count := strings.Count(contentStr, oldText)
-	if count > 1 {
-		return ErrorResult(
-			fmt.Sprintf("old_text appears %d times. Please provide more context to make it unique", count),
-		)
-	}
-
-	newContent := strings.Replace(contentStr, oldText, newText, 1)
-
-	if err := os.WriteFile(resolvedPath, []byte(newContent), 0o644); err != nil {
-		return ErrorResult(fmt.Sprintf("failed to write file: %v", err))
-	}
-
 	return SilentResult(fmt.Sprintf("File edited: %s", path))
 }
 
 type AppendFileTool struct {
-	workspace string
-	restrict  bool
+	fs fileSystem
 }
 
 func NewAppendFileTool(workspace string, restrict bool) *AppendFileTool {
-	return &AppendFileTool{workspace: workspace, restrict: restrict}
+	var fs fileSystem
+	if restrict {
+		fs = &sandboxFs{workspace: workspace}
+	} else {
+		fs = &hostFs{}
+	}
+	return &AppendFileTool{fs: fs}
 }
 
 func (t *AppendFileTool) Name() string {
@@ -148,20 +126,52 @@ func (t *AppendFileTool) Execute(ctx context.Context, args map[string]any) *Tool
 		return ErrorResult("content is required")
 	}
 
-	resolvedPath, err := validatePath(path, t.workspace, t.restrict)
-	if err != nil {
+	if err := appendFile(t.fs, path, content); err != nil {
 		return ErrorResult(err.Error())
 	}
+	return SilentResult(fmt.Sprintf("Appended to %s", path))
+}
 
-	f, err := os.OpenFile(resolvedPath, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0o644)
+// editFile reads the file via sysFs, performs the replacement, and writes back.
+// It uses a fileSystem interface, allowing the same logic for both restricted and unrestricted modes.
+func editFile(sysFs fileSystem, path, oldText, newText string) error {
+	content, err := sysFs.ReadFile(path)
 	if err != nil {
-		return ErrorResult(fmt.Sprintf("failed to open file: %v", err))
+		return err
 	}
-	defer f.Close()
 
-	if _, err := f.WriteString(content); err != nil {
-		return ErrorResult(fmt.Sprintf("failed to append to file: %v", err))
+	newContent, err := replaceEditContent(content, oldText, newText)
+	if err != nil {
+		return err
 	}
 
-	return SilentResult(fmt.Sprintf("Appended to %s", path))
+	return sysFs.WriteFile(path, newContent)
+}
+
+// appendFile reads the existing content (if any) via sysFs, appends new content, and writes back.
+func appendFile(sysFs fileSystem, path, appendContent string) error {
+	content, err := sysFs.ReadFile(path)
+	if err != nil && !errors.Is(err, fs.ErrNotExist) {
+		return err
+	}
+
+	newContent := append(content, []byte(appendContent)...)
+	return sysFs.WriteFile(path, newContent)
+}
+
+// replaceEditContent handles the core logic of finding and replacing a single occurrence of oldText.
+func replaceEditContent(content []byte, oldText, newText string) ([]byte, error) {
+	contentStr := string(content)
+
+	if !strings.Contains(contentStr, oldText) {
+		return nil, fmt.Errorf("old_text not found in file. Make sure it matches exactly")
+	}
+
+	count := strings.Count(contentStr, oldText)
+	if count > 1 {
+		return nil, fmt.Errorf("old_text appears %d times. Please provide more context to make it unique", count)
+	}
+
+	newContent := strings.Replace(contentStr, oldText, newText, 1)
+	return []byte(newContent), nil
 }

pkg/tools/edit_test.go

@@ -6,6 +6,8 @@ import (
 	"path/filepath"
 	"strings"
 	"testing"
+
+	"github.com/stretchr/testify/assert"
 )
 
 // TestEditTool_EditFile_Success verifies successful file editing
@@ -151,14 +153,18 @@ func TestEditTool_EditFile_OutsideAllowedDir(t *testing.T) {
 	result := tool.Execute(ctx, args)
 
 	// Should return error result
-	if !result.IsError {
-		t.Errorf("Expected error when path is outside allowed directory")
-	}
+	assert.True(t, result.IsError, "Expected error when path is outside allowed directory")
 
 	// Should mention outside allowed directory
-	if !strings.Contains(result.ForLLM, "outside") && !strings.Contains(result.ForUser, "outside") {
-		t.Errorf("Expected 'outside allowed' message, got ForLLM: %s", result.ForLLM)
-	}
+	// Note: ErrorResult only sets ForLLM by default, so ForUser might be empty.
+	// We check ForLLM as it's the primary error channel.
+	assert.True(
+		t,
+		strings.Contains(result.ForLLM, "outside") || strings.Contains(result.ForLLM, "access denied") ||
+			strings.Contains(result.ForLLM, "escapes"),
+		"Expected 'outside allowed' or 'access denied' message, got ForLLM: %s",
+		result.ForLLM,
+	)
 }
 
 // TestEditTool_EditFile_MissingPath verifies error handling for missing path
@@ -287,3 +293,145 @@ func TestEditTool_AppendFile_MissingContent(t *testing.T) {
 		t.Errorf("Expected error when content is missing")
 	}
 }
+
+// TestReplaceEditContent verifies the helper function replaceEditContent
+func TestReplaceEditContent(t *testing.T) {
+	tests := []struct {
+		name        string
+		content     []byte
+		oldText     string
+		newText     string
+		expected    []byte
+		expectError bool
+	}{
+		{
+			name:        "successful replacement",
+			content:     []byte("hello world"),
+			oldText:     "world",
+			newText:     "universe",
+			expected:    []byte("hello universe"),
+			expectError: false,
+		},
+		{
+			name:        "old text not found",
+			content:     []byte("hello world"),
+			oldText:     "golang",
+			newText:     "rust",
+			expected:    nil,
+			expectError: true,
+		},
+		{
+			name:        "multiple matches found",
+			content:     []byte("test text test"),
+			oldText:     "test",
+			newText:     "done",
+			expected:    nil,
+			expectError: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			result, err := replaceEditContent(tt.content, tt.oldText, tt.newText)
+			if tt.expectError {
+				assert.Error(t, err)
+			} else {
+				assert.NoError(t, err)
+				assert.Equal(t, tt.expected, result)
+			}
+		})
+	}
+}
+
+// TestAppendFileTool_AppendToNonExistent_Restricted verifies that AppendFileTool in restricted mode
+// can append to a file that does not yet exist — it should silently create the file.
+// This exercises the errors.Is(err, fs.ErrNotExist) path in appendFileWithRW + rootRW.
+func TestAppendFileTool_AppendToNonExistent_Restricted(t *testing.T) {
+	workspace := t.TempDir()
+	tool := NewAppendFileTool(workspace, true)
+	ctx := context.Background()
+
+	args := map[string]any{
+		"path":    "brand_new_file.txt",
+		"content": "first content",
+	}
+
+	result := tool.Execute(ctx, args)
+	assert.False(
+		t,
+		result.IsError,
+		"Expected success when appending to non-existent file in restricted mode, got: %s",
+		result.ForLLM,
+	)
+
+	// Verify the file was created with correct content
+	data, err := os.ReadFile(filepath.Join(workspace, "brand_new_file.txt"))
+	assert.NoError(t, err)
+	assert.Equal(t, "first content", string(data))
+}
+
+// TestAppendFileTool_Restricted_Success verifies that AppendFileTool in restricted mode
+// correctly appends to an existing file within the sandbox.
+func TestAppendFileTool_Restricted_Success(t *testing.T) {
+	workspace := t.TempDir()
+	testFile := "existing.txt"
+	err := os.WriteFile(filepath.Join(workspace, testFile), []byte("initial"), 0o644)
+	assert.NoError(t, err)
+
+	tool := NewAppendFileTool(workspace, true)
+	ctx := context.Background()
+	args := map[string]any{
+		"path":    testFile,
+		"content": " appended",
+	}
+
+	result := tool.Execute(ctx, args)
+	assert.False(t, result.IsError, "Expected success, got: %s", result.ForLLM)
+	assert.True(t, result.Silent)
+
+	data, err := os.ReadFile(filepath.Join(workspace, testFile))
+	assert.NoError(t, err)
+	assert.Equal(t, "initial appended", string(data))
+}
+
+// TestEditFileTool_Restricted_InPlaceEdit verifies that EditFileTool in restricted mode
+// correctly edits a file using the single-open editFileInRoot path.
+func TestEditFileTool_Restricted_InPlaceEdit(t *testing.T) {
+	workspace := t.TempDir()
+	testFile := "edit_target.txt"
+	err := os.WriteFile(filepath.Join(workspace, testFile), []byte("Hello World"), 0o644)
+	assert.NoError(t, err)
+
+	tool := NewEditFileTool(workspace, true)
+	ctx := context.Background()
+	args := map[string]any{
+		"path":     testFile,
+		"old_text": "World",
+		"new_text": "Go",
+	}
+
+	result := tool.Execute(ctx, args)
+	assert.False(t, result.IsError, "Expected success, got: %s", result.ForLLM)
+	assert.True(t, result.Silent)
+
+	data, err := os.ReadFile(filepath.Join(workspace, testFile))
+	assert.NoError(t, err)
+	assert.Equal(t, "Hello Go", string(data))
+}
+
+// TestEditFileTool_Restricted_FileNotFound verifies that editFileInRoot returns a proper
+// error message when the target file does not exist.
+func TestEditFileTool_Restricted_FileNotFound(t *testing.T) {
+	workspace := t.TempDir()
+	tool := NewEditFileTool(workspace, true)
+	ctx := context.Background()
+	args := map[string]any{
+		"path":     "no_such_file.txt",
+		"old_text": "old",
+		"new_text": "new",
+	}
+
+	result := tool.Execute(ctx, args)
+	assert.True(t, result.IsError)
+	assert.Contains(t, result.ForLLM, "not found")
+}