From b821b5556ddda935e338bdbfdf6c7adb0092a541 Mon Sep 17 00:00:00 2001 From: 0x5487 Date: Sun, 22 Feb 2026 15:37:44 +0800 Subject: [PATCH] refactor: unify filesystem access by introducing a `fileSystem` interface and updating tools to use it directly, removing `os.Root` dependency from `sandboxFs`. --- pkg/tools/edit.go | 69 +++++------ pkg/tools/filesystem.go | 226 +++++++++++++++++------------------ pkg/tools/filesystem_test.go | 58 ++++----- 3 files changed, 169 insertions(+), 184 deletions(-) diff --git a/pkg/tools/edit.go b/pkg/tools/edit.go index 0b0c24742..d3ab267bf 100644 --- a/pkg/tools/edit.go +++ b/pkg/tools/edit.go @@ -5,23 +5,24 @@ import ( "errors" "fmt" "io/fs" - "os" "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 { @@ -69,29 +70,24 @@ func (t *EditFileTool) Execute(ctx context.Context, args map[string]any) *ToolRe return ErrorResult("new_text is required") } - if t.restrict { - 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(&hostFs{}, path, oldText, newText); err != nil { + if err := editFile(t.fs, path, oldText, newText); err != nil { return ErrorResult(err.Error()) } 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 { @@ -130,27 +126,16 @@ func (t *AppendFileTool) Execute(ctx context.Context, args map[string]any) *Tool return ErrorResult("content is required") } - var rw fileReadWriter - if t.restrict { - return executeInWorkspace(t.workspace, path, func(root *os.Root, relPath string) (*ToolResult, error) { - if err := appendFile(&sandboxFs{root: root}, relPath, content); err != nil { - return nil, err - } - return SilentResult(fmt.Sprintf("Appended to %s", path)), nil - }) - } - - rw = &hostFs{} - if err := appendFile(rw, path, content); err != nil { + if err := appendFile(t.fs, path, content); err != nil { return ErrorResult(err.Error()) } return SilentResult(fmt.Sprintf("Appended to %s", path)) } -// editFile reads the file via rw, performs the replacement, and writes back. -// It uses a fileReadWriter, allowing the same logic for both restricted and unrestricted modes. -func editFile(rw fileReadWriter, path, oldText, newText string) error { - content, err := rw.Read(path) +// 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 err } @@ -160,18 +145,18 @@ func editFile(rw fileReadWriter, path, oldText, newText string) error { return err } - return rw.Write(path, newContent) + return sysFs.WriteFile(path, newContent) } -// appendFile reads the existing content (if any) via rw, appends new content, and writes back. -func appendFile(rw fileReadWriter, path, appendContent string) error { - content, err := rw.Read(path) +// 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 rw.Write(path, newContent) + return sysFs.WriteFile(path, newContent) } // replaceEditContent handles the core logic of finding and replacing a single occurrence of oldText. diff --git a/pkg/tools/filesystem.go b/pkg/tools/filesystem.go index b9e2a0d22..d713aebf5 100644 --- a/pkg/tools/filesystem.go +++ b/pkg/tools/filesystem.go @@ -82,12 +82,17 @@ func isWithinWorkspace(candidate, workspace string) bool { } type ReadFileTool struct { - workspace string - restrict bool + fs fileSystem } func NewReadFileTool(workspace string, restrict bool) *ReadFileTool { - return &ReadFileTool{workspace: workspace, restrict: restrict} + var fs fileSystem + if restrict { + fs = &sandboxFs{workspace: workspace} + } else { + fs = &hostFs{} + } + return &ReadFileTool{fs: fs} } func (t *ReadFileTool) Name() string { @@ -117,17 +122,7 @@ func (t *ReadFileTool) Execute(ctx context.Context, args map[string]any) *ToolRe return ErrorResult("path is required") } - if t.restrict { - 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 - } - return NewToolResult(string(content)), nil - }) - } - - content, err := (&hostFs{}).Read(path) + content, err := t.fs.ReadFile(path) if err != nil { return ErrorResult(err.Error()) } @@ -135,12 +130,17 @@ func (t *ReadFileTool) Execute(ctx context.Context, args map[string]any) *ToolRe } type WriteFileTool struct { - workspace string - restrict bool + fs fileSystem } func NewWriteFileTool(workspace string, restrict bool) *WriteFileTool { - return &WriteFileTool{workspace: workspace, restrict: restrict} + var fs fileSystem + if restrict { + fs = &sandboxFs{workspace: workspace} + } else { + fs = &hostFs{} + } + return &WriteFileTool{fs: fs} } func (t *WriteFileTool) Name() string { @@ -179,16 +179,7 @@ func (t *WriteFileTool) Execute(ctx context.Context, args map[string]any) *ToolR return ErrorResult("content is required") } - if t.restrict { - 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 := (&hostFs{}).Write(path, []byte(content)); err != nil { + if err := t.fs.WriteFile(path, []byte(content)); err != nil { return ErrorResult(err.Error()) } @@ -196,12 +187,17 @@ func (t *WriteFileTool) Execute(ctx context.Context, args map[string]any) *ToolR } type ListDirTool struct { - workspace string - restrict bool + fs fileSystem } func NewListDirTool(workspace string, restrict bool) *ListDirTool { - return &ListDirTool{workspace: workspace, restrict: restrict} + var fs fileSystem + if restrict { + fs = &sandboxFs{workspace: workspace} + } else { + fs = &hostFs{} + } + return &ListDirTool{fs: fs} } func (t *ListDirTool) Name() string { @@ -231,22 +227,11 @@ func (t *ListDirTool) Execute(ctx context.Context, args map[string]any) *ToolRes path = "." } - if !t.restrict { - entries, err := os.ReadDir(path) - if err != nil { - return ErrorResult(fmt.Sprintf("failed to read directory: %v", err)) - } - return formatDirEntries(entries) + entries, err := t.fs.ReadDir(path) + if err != nil { + return ErrorResult(fmt.Sprintf("failed to read directory: %v", err)) } - - 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) - } - - return formatDirEntries(entries), nil - }) + return formatDirEntries(entries) } func formatDirEntries(entries []os.DirEntry) *ToolResult { @@ -261,17 +246,18 @@ func formatDirEntries(entries []os.DirEntry) *ToolResult { return NewToolResult(result.String()) } -// fileReadWriter abstracts reading and writing files, allowing both unrestricted -// (host filesystem) and sandbox (os.Root) implementations to share the same logic. -type fileReadWriter interface { - Read(path string) ([]byte, error) - Write(path string, data []byte) error +// fileSystem abstracts reading, writing, and listing files, allowing both +// unrestricted (host filesystem) and sandbox (os.Root) implementations to share the same polymorphic interface. +type fileSystem interface { + ReadFile(path string) ([]byte, error) + WriteFile(path string, data []byte) error + ReadDir(path string) ([]os.DirEntry, error) } // hostFs is an unrestricted fileReadWriter that operates directly on the host filesystem. type hostFs struct{} -func (h *hostFs) Read(path string) ([]byte, error) { +func (h *hostFs) ReadFile(path string) ([]byte, error) { content, err := os.ReadFile(path) if err != nil { if os.IsNotExist(err) { @@ -285,7 +271,11 @@ func (h *hostFs) Read(path string) ([]byte, error) { return content, nil } -func (h *hostFs) Write(path string, data []byte) error { +func (h *hostFs) ReadDir(path string) ([]os.DirEntry, error) { + return os.ReadDir(path) +} + +func (h *hostFs) WriteFile(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) @@ -307,50 +297,88 @@ func (h *hostFs) Write(path string, data []byte) error { return nil } -// sandboxFs is a sandboxed fileReadWriter that operates within an os.Root boundary. -// All paths passed to Read/Write must be relative to the root. +// sandboxFs is a sandboxed fileSystem that operates within a strictly defined workspace using os.Root. type sandboxFs struct { - root *os.Root + workspace string } -func (r *sandboxFs) Read(path string) ([]byte, error) { - content, err := r.root.ReadFile(path) +func (r *sandboxFs) execute(path string, fn func(root *os.Root, relPath string) error) error { + if r.workspace == "" { + return fmt.Errorf("workspace is not defined") + } + + root, err := os.OpenRoot(r.workspace) if err != nil { - if os.IsNotExist(err) { - return nil, fmt.Errorf("failed to read file: file not found: %w", err) - } - // os.Root returns "escapes from parent" for paths outside the root - if os.IsPermission(err) || strings.Contains(err.Error(), "escapes from parent") || strings.Contains(err.Error(), "permission denied") { - return nil, fmt.Errorf("failed to read file: access denied: %w", err) - } - return nil, fmt.Errorf("failed to read file: %w", err) + return fmt.Errorf("failed to open workspace: %w", err) } - return content, nil + defer root.Close() + + relPath, err := getSafeRelPath(r.workspace, path) + if err != nil { + return err + } + + return fn(root, relPath) } -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 { - return fmt.Errorf("failed to create parent directories: %w", err) +func (r *sandboxFs) ReadFile(path string) ([]byte, error) { + var content []byte + err := r.execute(path, func(root *os.Root, relPath string) error { + fileContent, err := root.ReadFile(relPath) + if err != nil { + if os.IsNotExist(err) { + return fmt.Errorf("failed to read file: file not found: %w", err) + } + // os.Root returns "escapes from parent" for paths outside the root + if os.IsPermission(err) || strings.Contains(err.Error(), "escapes from parent") || strings.Contains(err.Error(), "permission denied") { + return fmt.Errorf("failed to read file: access denied: %w", err) + } + return fmt.Errorf("failed to read file: %w", err) } - } + content = fileContent + return nil + }) + return content, 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. - tmpRelPath := fmt.Sprintf("%s.%d.tmp", path, time.Now().UnixNano()) +func (r *sandboxFs) WriteFile(path string, data []byte) error { + return r.execute(path, func(root *os.Root, relPath string) error { + dir := filepath.Dir(relPath) + if dir != "." && dir != "/" { + if err := root.MkdirAll(dir, 0755); err != nil { + return fmt.Errorf("failed to create parent directories: %w", err) + } + } - 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) - } + // 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", relPath, time.Now().UnixNano()) - 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) - } - return nil + if err := root.WriteFile(tmpRelPath, data, 0644); err != nil { + root.Remove(tmpRelPath) // Ensure cleanup of partial/empty temp file + return fmt.Errorf("failed to write to temp file: %w", err) + } + + if err := root.Rename(tmpRelPath, relPath); err != nil { + root.Remove(tmpRelPath) + return fmt.Errorf("failed to rename temp file over target: %w", err) + } + return nil + }) +} + +func (r *sandboxFs) ReadDir(path string) ([]os.DirEntry, error) { + var entries []os.DirEntry + err := r.execute(path, func(root *os.Root, relPath string) error { + dirEntries, err := fs.ReadDir(root.FS(), relPath) + if err != nil { + return err + } + entries = dirEntries + return nil + }) + return entries, err } // Helper to get a safe relative path for os.Root usage @@ -374,31 +402,3 @@ func getSafeRelPath(workspace, path string) (string, error) { return rel, nil } - -// 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") - } - - // 1. Open the Root - root, err := os.OpenRoot(workspace) - if err != nil { - return ErrorResult(fmt.Sprintf("failed to open workspace: %v", err)) - } - defer root.Close() - - // 2. Calculate relative path - relPath, err := getSafeRelPath(workspace, path) - if err != nil { - return ErrorResult(err.Error()) - } - - // 3. Execute the operation - result, err := fn(root, relPath) - if err != nil { - return ErrorResult(err.Error()) - } - - return result -} diff --git a/pkg/tools/filesystem_test.go b/pkg/tools/filesystem_test.go index a64904f3a..1687ca395 100644 --- a/pkg/tools/filesystem_test.go +++ b/pkg/tools/filesystem_test.go @@ -17,7 +17,7 @@ func TestFilesystemTool_ReadFile_Success(t *testing.T) { testFile := filepath.Join(tmpDir, "test.txt") os.WriteFile(testFile, []byte("test content"), 0o644) - tool := &ReadFileTool{} + tool := NewReadFileTool("", false) ctx := context.Background() args := map[string]any{ "path": testFile, @@ -44,7 +44,7 @@ func TestFilesystemTool_ReadFile_Success(t *testing.T) { // TestFilesystemTool_ReadFile_NotFound verifies error handling for missing file func TestFilesystemTool_ReadFile_NotFound(t *testing.T) { - tool := &ReadFileTool{} + tool := NewReadFileTool("", false) ctx := context.Background() args := map[string]any{ "path": "/nonexistent_file_12345.txt", @@ -87,7 +87,7 @@ func TestFilesystemTool_WriteFile_Success(t *testing.T) { tmpDir := t.TempDir() testFile := filepath.Join(tmpDir, "newfile.txt") - tool := &WriteFileTool{} + tool := NewWriteFileTool("", false) ctx := context.Background() args := map[string]any{ "path": testFile, @@ -126,7 +126,7 @@ func TestFilesystemTool_WriteFile_CreateDir(t *testing.T) { tmpDir := t.TempDir() testFile := filepath.Join(tmpDir, "subdir", "newfile.txt") - tool := &WriteFileTool{} + tool := NewWriteFileTool("", false) ctx := context.Background() args := map[string]any{ "path": testFile, @@ -152,7 +152,7 @@ func TestFilesystemTool_WriteFile_CreateDir(t *testing.T) { // TestFilesystemTool_WriteFile_MissingPath verifies error handling for missing path func TestFilesystemTool_WriteFile_MissingPath(t *testing.T) { - tool := &WriteFileTool{} + tool := NewWriteFileTool("", false) ctx := context.Background() args := map[string]any{ "content": "test", @@ -168,7 +168,7 @@ func TestFilesystemTool_WriteFile_MissingPath(t *testing.T) { // TestFilesystemTool_WriteFile_MissingContent verifies error handling for missing content func TestFilesystemTool_WriteFile_MissingContent(t *testing.T) { - tool := &WriteFileTool{} + tool := NewWriteFileTool("", false) ctx := context.Background() args := map[string]any{ "path": "/tmp/test.txt", @@ -195,7 +195,7 @@ func TestFilesystemTool_ListDir_Success(t *testing.T) { os.WriteFile(filepath.Join(tmpDir, "file2.txt"), []byte("content"), 0o644) os.Mkdir(filepath.Join(tmpDir, "subdir"), 0o755) - tool := &ListDirTool{} + tool := NewListDirTool("", false) ctx := context.Background() args := map[string]any{ "path": tmpDir, @@ -219,7 +219,7 @@ func TestFilesystemTool_ListDir_Success(t *testing.T) { // TestFilesystemTool_ListDir_NotFound verifies error handling for non-existent directory func TestFilesystemTool_ListDir_NotFound(t *testing.T) { - tool := &ListDirTool{} + tool := NewListDirTool("", false) ctx := context.Background() args := map[string]any{ "path": "/nonexistent_directory_12345", @@ -240,7 +240,7 @@ func TestFilesystemTool_ListDir_NotFound(t *testing.T) { // TestFilesystemTool_ListDir_DefaultPath verifies default to current directory func TestFilesystemTool_ListDir_DefaultPath(t *testing.T) { - tool := &ListDirTool{} + tool := NewListDirTool("", false) ctx := context.Background() args := map[string]any{} @@ -371,7 +371,7 @@ func TestHostRW_Read_PermissionDenied(t *testing.T) { assert.NoError(t, err) defer os.Chmod(protected, 0644) // ensure cleanup - _, err = (&hostFs{}).Read(protected) + _, err = (&hostFs{}).ReadFile(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 := (&hostFs{}).Read(tmpDir) + _, err := (&hostFs{}).ReadFile(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 = (&sandboxFs{root: root}).Read("subdir") + _, err = (&sandboxFs{workspace: workspace}).ReadFile("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 := (&hostFs{}).Write(target, []byte("hello")) + err := (&hostFs{}).WriteFile(target, []byte("hello")) assert.NoError(t, err) data, err := os.ReadFile(target) @@ -416,12 +416,9 @@ func TestHostRW_Write_ParentDirMissing(t *testing.T) { // nested parent directories automatically within the sandbox. func TestRootRW_Write_ParentDirMissing(t *testing.T) { workspace := t.TempDir() - root, err := os.OpenRoot(workspace) - assert.NoError(t, err) - defer root.Close() relPath := "x/y/z/file.txt" - err = (&sandboxFs{root: root}).Write(relPath, []byte("nested")) + err := (&sandboxFs{workspace: workspace}).WriteFile(relPath, []byte("nested")) assert.NoError(t, err) data, err := os.ReadFile(filepath.Join(workspace, relPath)) @@ -435,7 +432,7 @@ func TestHostRW_Write(t *testing.T) { testFile := filepath.Join(tmpDir, "atomic_test.txt") testData := []byte("atomic test content") - err := (&hostFs{}).Write(testFile, testData) + err := (&hostFs{}).WriteFile(testFile, testData) assert.NoError(t, err) content, err := os.ReadFile(testFile) @@ -444,7 +441,7 @@ func TestHostRW_Write(t *testing.T) { // Verify it overwrites correctly newData := []byte("new atomic content") - err = (&hostFs{}).Write(testFile, newData) + err = (&hostFs{}).WriteFile(testFile, newData) assert.NoError(t, err) content, err = os.ReadFile(testFile) @@ -455,33 +452,36 @@ func TestHostRW_Write(t *testing.T) { // TestRootRW_Write verifies the rootRW.Write helper function func TestRootRW_Write(t *testing.T) { tmpDir := t.TempDir() - root, err := os.OpenRoot(tmpDir) - assert.NoError(t, err) - defer root.Close() relPath := "atomic_root_test.txt" testData := []byte("atomic root test content") - erw := &sandboxFs{root: root} - err = erw.Write(relPath, testData) + erw := &sandboxFs{workspace: tmpDir} + err := erw.WriteFile(relPath, testData) assert.NoError(t, err) + root, err := os.OpenRoot(tmpDir) + assert.NoError(t, err) + defer root.Close() + f, err := root.Open(relPath) assert.NoError(t, err) + defer f.Close() + content, err := io.ReadAll(f) assert.NoError(t, err) - f.Close() assert.Equal(t, testData, content) // Verify it overwrites correctly newData := []byte("new root atomic content") - err = erw.Write(relPath, newData) + err = erw.WriteFile(relPath, newData) assert.NoError(t, err) - f, err = root.Open(relPath) + f2, err := root.Open(relPath) assert.NoError(t, err) - content, err = io.ReadAll(f) + defer f2.Close() + + content, err = io.ReadAll(f2) assert.NoError(t, err) - f.Close() assert.Equal(t, newData, content) }