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

This commit is contained in:
0x5487 2026-02-22 14:56:23 +08:00
parent 4e69efa216
commit cb09447501
3 changed files with 46 additions and 94 deletions

View file

@ -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)

View file

@ -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()

View file

@ -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)