refactor: unify filesystem read/write operations with atomic write guarantees and clearer naming.

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) {
+		return executeInWorkspace(t.allowedDir, path, func(root *os.Root, relPath string) (*ToolResult, error) {
-			if err := editFileInRoot(root, relPath, oldText, newText); err != nil {
+			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) {
+		return executeInWorkspace(t.workspace, path, func(root *os.Root, relPath string) (*ToolResult, error) {
-			if err := appendFileWithRW(&rootRW{root: root}, relPath, appendContent); err != nil {
+			if err := appendFileWithRW(&sandboxFs{root: root}, relPath, content); err != nil {
 				return nil, err
 			}
 			return SilentResult(fmt.Sprintf("Appended to %s", path)), nil
 		})
 	}
 
-	rw = &hostRW{}
+	rw = &hostFs{}
-	if err := appendFileWithRW(rw, path, appendContent); err != nil {
+	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)

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) {
+		return executeInWorkspace(t.workspace, path, func(root *os.Root, relPath string) (*ToolResult, error) {
-			content, err := (&rootRW{root: root}).Read(relPath)
+			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) {
+		return executeInWorkspace(t.workspace, path, func(root *os.Root, relPath string) (*ToolResult, error) {
-			if err := (&rootRW{root: root}).Write(relPath, []byte(content)); err != nil {
+			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) {
+	return executeInWorkspace(t.workspace, path, func(root *os.Root, relPath string) (*ToolResult, error) {
-		f, err := root.Open(relPath)
+		entries, err := fs.ReadDir(root.FS(), relPath)
-		if err != nil {
-			return nil, fmt.Errorf("failed to open directory: %w", err)
-		}
-		defer f.Close()
-
-		entries, err := f.ReadDir(-1)
 		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.
+// hostFs is an unrestricted fileReadWriter that operates directly on the host filesystem.
-type hostRW struct{}
+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 {
+	if err := r.root.WriteFile(tmpRelPath, data, 0644); err != nil {
-		fw.Close()
+		r.root.Remove(tmpRelPath) // Ensure cleanup of partial/empty temp file
-		r.root.Remove(tmpRelPath)
 		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
+// executeInWorkspace 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 {
+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()
 

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)