diff --git a/pkg/tools/edit.go b/pkg/tools/edit.go index dc5f9795e..7df909216 100644 --- a/pkg/tools/edit.go +++ b/pkg/tools/edit.go @@ -4,7 +4,6 @@ import ( "context" "errors" "fmt" - "io" "io/fs" "os" "strings" @@ -71,15 +70,16 @@ func (t *EditFileTool) Execute(ctx context.Context, args map[string]any) *ToolRe } if t.restrict { - return executeInRoot(t.allowedDir, path, func(root *os.Root, relPath string) (*ToolResult, error) { - if err := editFileInRoot(root, relPath, oldText, newText); err != nil { + return executeInWorkspace(t.allowedDir, path, func(root *os.Root, relPath string) (*ToolResult, error) { + rw := &sandboxFs{root: root} + if err := editFile(rw, relPath, oldText, newText); err != nil { return nil, err } return SilentResult(fmt.Sprintf("File edited: %s", path)), nil }) } - if err := editFile(&hostRW{}, path, oldText, newText); err != nil { + if err := editFile(&hostFs{}, path, oldText, newText); err != nil { return ErrorResult(err.Error()) } return SilentResult(fmt.Sprintf("File edited: %s", path)) @@ -125,23 +125,23 @@ func (t *AppendFileTool) Execute(ctx context.Context, args map[string]any) *Tool return ErrorResult("path is required") } - appendContent, ok := args["content"].(string) + content, ok := args["content"].(string) if !ok { return ErrorResult("content is required") } var rw fileReadWriter if t.restrict { - return executeInRoot(t.workspace, path, func(root *os.Root, relPath string) (*ToolResult, error) { - if err := appendFileWithRW(&rootRW{root: root}, relPath, appendContent); err != nil { + return executeInWorkspace(t.workspace, path, func(root *os.Root, relPath string) (*ToolResult, error) { + if err := appendFileWithRW(&sandboxFs{root: root}, relPath, content); err != nil { return nil, err } return SilentResult(fmt.Sprintf("Appended to %s", path)), nil }) } - rw = &hostRW{} - if err := appendFileWithRW(rw, path, appendContent); err != nil { + rw = &hostFs{} + if err := appendFileWithRW(rw, path, content); err != nil { return ErrorResult(err.Error()) } return SilentResult(fmt.Sprintf("Appended to %s", path)) @@ -163,46 +163,6 @@ func editFile(rw fileReadWriter, path, oldText, newText string) error { return rw.Write(path, newContent) } -// editFileInRoot performs an in-place edit within an os.Root using a single open call. -// By opening with O_RDWR and reusing the same file descriptor for both read and write, -// we narrow the TOCTOU window compared to two separate open calls. -func editFileInRoot(root *os.Root, relPath, oldText, newText string) error { - f, err := root.OpenFile(relPath, os.O_RDWR, 0) - if err != nil { - if os.IsNotExist(err) { - return fmt.Errorf("failed to read file: file not found: %w", err) - } - if os.IsPermission(err) || strings.Contains(err.Error(), "escapes from parent") { - return fmt.Errorf("failed to read file: access denied: %w", err) - } - return fmt.Errorf("failed to open file for editing: %w", err) - } - defer f.Close() - - content, err := io.ReadAll(f) - if err != nil { - return fmt.Errorf("failed to read file content: %w", err) - } - - newContent, err := replaceEditContent(content, oldText, newText) - if err != nil { - return err - } - - // Truncate the file and seek back to the beginning before writing. - if err := f.Truncate(0); err != nil { - return fmt.Errorf("failed to truncate file for in-place edit: %w", err) - } - if _, err := f.Seek(0, io.SeekStart); err != nil { - return fmt.Errorf("failed to seek to beginning of file: %w", err) - } - - if _, err := f.Write(newContent); err != nil { - return fmt.Errorf("failed to write edited content: %w", err) - } - return nil -} - // appendFileWithRW reads the existing content (if any) via rw, appends new content, and writes back. func appendFileWithRW(rw fileReadWriter, path, appendContent string) error { content, err := rw.Read(path) diff --git a/pkg/tools/filesystem.go b/pkg/tools/filesystem.go index 114969bf9..b9e2a0d22 100644 --- a/pkg/tools/filesystem.go +++ b/pkg/tools/filesystem.go @@ -3,6 +3,7 @@ package tools import ( "context" "fmt" + "io/fs" "os" "path/filepath" "strings" @@ -117,8 +118,8 @@ func (t *ReadFileTool) Execute(ctx context.Context, args map[string]any) *ToolRe } if t.restrict { - return executeInRoot(t.workspace, path, func(root *os.Root, relPath string) (*ToolResult, error) { - content, err := (&rootRW{root: root}).Read(relPath) + return executeInWorkspace(t.workspace, path, func(root *os.Root, relPath string) (*ToolResult, error) { + content, err := (&sandboxFs{root: root}).Read(relPath) if err != nil { return nil, err } @@ -126,7 +127,7 @@ func (t *ReadFileTool) Execute(ctx context.Context, args map[string]any) *ToolRe }) } - content, err := (&hostRW{}).Read(path) + content, err := (&hostFs{}).Read(path) if err != nil { return ErrorResult(err.Error()) } @@ -179,15 +180,15 @@ func (t *WriteFileTool) Execute(ctx context.Context, args map[string]any) *ToolR } if t.restrict { - return executeInRoot(t.workspace, path, func(root *os.Root, relPath string) (*ToolResult, error) { - if err := (&rootRW{root: root}).Write(relPath, []byte(content)); err != nil { + return executeInWorkspace(t.workspace, path, func(root *os.Root, relPath string) (*ToolResult, error) { + if err := (&sandboxFs{root: root}).Write(relPath, []byte(content)); err != nil { return nil, err } return SilentResult(fmt.Sprintf("File written: %s", path)), nil }) } - if err := (&hostRW{}).Write(path, []byte(content)); err != nil { + if err := (&hostFs{}).Write(path, []byte(content)); err != nil { return ErrorResult(err.Error()) } @@ -238,14 +239,8 @@ func (t *ListDirTool) Execute(ctx context.Context, args map[string]any) *ToolRes return formatDirEntries(entries) } - return executeInRoot(t.workspace, path, func(root *os.Root, relPath string) (*ToolResult, error) { - f, err := root.Open(relPath) - if err != nil { - return nil, fmt.Errorf("failed to open directory: %w", err) - } - defer f.Close() - - entries, err := f.ReadDir(-1) + return executeInWorkspace(t.workspace, path, func(root *os.Root, relPath string) (*ToolResult, error) { + entries, err := fs.ReadDir(root.FS(), relPath) if err != nil { return nil, fmt.Errorf("failed to read directory: %w", err) } @@ -273,10 +268,10 @@ type fileReadWriter interface { Write(path string, data []byte) error } -// hostRW is an unrestricted fileReadWriter that operates directly on the host filesystem. -type hostRW struct{} +// hostFs is an unrestricted fileReadWriter that operates directly on the host filesystem. +type hostFs struct{} -func (h *hostRW) Read(path string) ([]byte, error) { +func (h *hostFs) Read(path string) ([]byte, error) { content, err := os.ReadFile(path) if err != nil { if os.IsNotExist(err) { @@ -290,14 +285,18 @@ func (h *hostRW) Read(path string) ([]byte, error) { return content, nil } -func (h *hostRW) Write(path string, data []byte) error { +func (h *hostFs) Write(path string, data []byte) error { dir := filepath.Dir(path) if err := os.MkdirAll(dir, 0755); err != nil { return fmt.Errorf("failed to create parent directories: %w", err) } + // We use a "write-then-rename" pattern here to ensure an atomic write. + // This prevents the target file from being left in a truncated or partial state + // if the operation is interrupted, as the rename operation is atomic on Linux. tmpPath := fmt.Sprintf("%s.%d.tmp", path, time.Now().UnixNano()) if err := os.WriteFile(tmpPath, data, 0644); err != nil { + os.Remove(tmpPath) // Ensure cleanup of partial/empty temp file return fmt.Errorf("failed to write temp file: %w", err) } @@ -308,13 +307,13 @@ func (h *hostRW) Write(path string, data []byte) error { return nil } -// rootRW is a sandboxed fileReadWriter that operates within an os.Root boundary. +// sandboxFs is a sandboxed fileReadWriter that operates within an os.Root boundary. // All paths passed to Read/Write must be relative to the root. -type rootRW struct { +type sandboxFs struct { root *os.Root } -func (r *rootRW) Read(path string) ([]byte, error) { +func (r *sandboxFs) Read(path string) ([]byte, error) { content, err := r.root.ReadFile(path) if err != nil { if os.IsNotExist(err) { @@ -329,7 +328,7 @@ func (r *rootRW) Read(path string) ([]byte, error) { return content, nil } -func (r *rootRW) Write(path string, data []byte) error { +func (r *sandboxFs) Write(path string, data []byte) error { dir := filepath.Dir(path) if dir != "." && dir != "/" { if err := r.root.MkdirAll(dir, 0755); err != nil { @@ -337,23 +336,16 @@ func (r *rootRW) Write(path string, data []byte) error { } } + // We use a "write-then-rename" pattern here to ensure an atomic write. + // This prevents the target file from being left in a truncated or partial state + // if the operation is interrupted, as the rename operation is atomic on Linux. tmpRelPath := fmt.Sprintf("%s.%d.tmp", path, time.Now().UnixNano()) - fw, err := r.root.Create(tmpRelPath) - if err != nil { - return fmt.Errorf("failed to create temp file for writing: %w", err) - } - if _, err := fw.Write(data); err != nil { - fw.Close() - r.root.Remove(tmpRelPath) + if err := r.root.WriteFile(tmpRelPath, data, 0644); err != nil { + r.root.Remove(tmpRelPath) // Ensure cleanup of partial/empty temp file return fmt.Errorf("failed to write to temp file: %w", err) } - if err := fw.Close(); err != nil { - r.root.Remove(tmpRelPath) - return fmt.Errorf("failed to close temp file: %w", err) - } - if err := r.root.Rename(tmpRelPath, path); err != nil { r.root.Remove(tmpRelPath) return fmt.Errorf("failed to rename temp file over target: %w", err) @@ -383,8 +375,8 @@ func getSafeRelPath(workspace, path string) (string, error) { return rel, nil } -// executeInRoot executes a function within the safety of os.Root -func executeInRoot(workspace string, path string, fn func(root *os.Root, relPath string) (*ToolResult, error)) *ToolResult { +// executeInWorkspace executes a function within the safety of os.Root +func executeInWorkspace(workspace string, path string, fn func(root *os.Root, relPath string) (*ToolResult, error)) *ToolResult { if workspace == "" { return ErrorResult("workspace is not defined") } @@ -392,7 +384,7 @@ func executeInRoot(workspace string, path string, fn func(root *os.Root, relPath // 1. Open the Root root, err := os.OpenRoot(workspace) if err != nil { - return ErrorResult(fmt.Sprintf("failed to open workspace root: %v", err)) + return ErrorResult(fmt.Sprintf("failed to open workspace: %v", err)) } defer root.Close() diff --git a/pkg/tools/filesystem_test.go b/pkg/tools/filesystem_test.go index aa2d970e9..a64904f3a 100644 --- a/pkg/tools/filesystem_test.go +++ b/pkg/tools/filesystem_test.go @@ -371,7 +371,7 @@ func TestHostRW_Read_PermissionDenied(t *testing.T) { assert.NoError(t, err) defer os.Chmod(protected, 0644) // ensure cleanup - _, err = (&hostRW{}).Read(protected) + _, err = (&hostFs{}).Read(protected) assert.Error(t, err) assert.Contains(t, err.Error(), "access denied") } @@ -380,7 +380,7 @@ func TestHostRW_Read_PermissionDenied(t *testing.T) { func TestHostRW_Read_Directory(t *testing.T) { tmpDir := t.TempDir() - _, err := (&hostRW{}).Read(tmpDir) + _, err := (&hostFs{}).Read(tmpDir) assert.Error(t, err, "expected error when reading a directory as a file") } @@ -395,7 +395,7 @@ func TestRootRW_Read_Directory(t *testing.T) { err = root.Mkdir("subdir", 0755) assert.NoError(t, err) - _, err = (&rootRW{root: root}).Read("subdir") + _, err = (&sandboxFs{root: root}).Read("subdir") assert.Error(t, err, "expected error when reading a directory as a file") } @@ -404,7 +404,7 @@ func TestHostRW_Write_ParentDirMissing(t *testing.T) { tmpDir := t.TempDir() target := filepath.Join(tmpDir, "a", "b", "c", "file.txt") - err := (&hostRW{}).Write(target, []byte("hello")) + err := (&hostFs{}).Write(target, []byte("hello")) assert.NoError(t, err) data, err := os.ReadFile(target) @@ -421,7 +421,7 @@ func TestRootRW_Write_ParentDirMissing(t *testing.T) { defer root.Close() relPath := "x/y/z/file.txt" - err = (&rootRW{root: root}).Write(relPath, []byte("nested")) + err = (&sandboxFs{root: root}).Write(relPath, []byte("nested")) assert.NoError(t, err) data, err := os.ReadFile(filepath.Join(workspace, relPath)) @@ -435,7 +435,7 @@ func TestHostRW_Write(t *testing.T) { testFile := filepath.Join(tmpDir, "atomic_test.txt") testData := []byte("atomic test content") - err := (&hostRW{}).Write(testFile, testData) + err := (&hostFs{}).Write(testFile, testData) assert.NoError(t, err) content, err := os.ReadFile(testFile) @@ -444,7 +444,7 @@ func TestHostRW_Write(t *testing.T) { // Verify it overwrites correctly newData := []byte("new atomic content") - err = (&hostRW{}).Write(testFile, newData) + err = (&hostFs{}).Write(testFile, newData) assert.NoError(t, err) content, err = os.ReadFile(testFile) @@ -462,7 +462,7 @@ func TestRootRW_Write(t *testing.T) { relPath := "atomic_root_test.txt" testData := []byte("atomic root test content") - erw := &rootRW{root: root} + erw := &sandboxFs{root: root} err = erw.Write(relPath, testData) assert.NoError(t, err)