refactor: unify filesystem read/write operations with atomic write guarantees and clearer naming.
This commit is contained in:
parent
4e69efa216
commit
cb09447501
3 changed files with 46 additions and 94 deletions
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue