refactor: reimplement filesystem tools with os.OpenRoot for enhanced security and simplified path validation.

This commit is contained in:
0x5487 2026-02-19 13:31:43 +08:00
parent 4e8b6ad764
commit d595d414d3
4 changed files with 257 additions and 120 deletions

View file

@ -3,6 +3,7 @@ package tools
import ( import (
"context" "context"
"fmt" "fmt"
"io"
"os" "os"
"strings" "strings"
) )
@ -67,38 +68,77 @@ func (t *EditFileTool) Execute(ctx context.Context, args map[string]interface{})
return ErrorResult("new_text is required") return ErrorResult("new_text is required")
} }
resolvedPath, err := validatePath(path, t.allowedDir, t.restrict) // If not restricted, perform operations directly
if err != nil { if !t.restrict {
return ErrorResult(err.Error()) content, err := os.ReadFile(path)
if err != nil {
if os.IsNotExist(err) {
return ErrorResult(fmt.Sprintf("file not found: %s", path))
}
return ErrorResult(fmt.Sprintf("failed to read file: %v", err))
}
contentStr := string(content)
if !strings.Contains(contentStr, oldText) {
return ErrorResult("old_text not found in file. Make sure it matches exactly")
}
count := strings.Count(contentStr, oldText)
if count > 1 {
return ErrorResult(fmt.Sprintf("old_text appears %d times. Please provide more context to make it unique", count))
}
newContent := strings.Replace(contentStr, oldText, newText, 1)
if err := os.WriteFile(path, []byte(newContent), 0644); err != nil {
return ErrorResult(fmt.Sprintf("failed to write file: %v", err))
}
return SilentResult(fmt.Sprintf("File edited: %s", path))
} }
if _, err := os.Stat(resolvedPath); os.IsNotExist(err) { // Use executeInRoot to safely access the file
return ErrorResult(fmt.Sprintf("file not found: %s", path)) return executeInRoot(t.allowedDir, path, func(root *os.Root, relPath string) (*ToolResult, error) {
} f, err := root.Open(relPath)
if err != nil {
if os.IsNotExist(err) {
return nil, fmt.Errorf("file not found: %s", path)
}
return nil, fmt.Errorf("failed to open file: %w", err)
}
content, err := os.ReadFile(resolvedPath) content, err := io.ReadAll(f)
if err != nil { f.Close()
return ErrorResult(fmt.Sprintf("failed to read file: %v", err))
}
contentStr := string(content) if err != nil {
return nil, fmt.Errorf("failed to read file: %v", err)
}
if !strings.Contains(contentStr, oldText) { contentStr := string(content)
return ErrorResult("old_text not found in file. Make sure it matches exactly")
}
count := strings.Count(contentStr, oldText) if !strings.Contains(contentStr, oldText) {
if count > 1 { return nil, fmt.Errorf("old_text not found in file. Make sure it matches exactly")
return ErrorResult(fmt.Sprintf("old_text appears %d times. Please provide more context to make it unique", count)) }
}
newContent := strings.Replace(contentStr, oldText, newText, 1) count := strings.Count(contentStr, oldText)
if count > 1 {
return nil, fmt.Errorf("old_text appears %d times. Please provide more context to make it unique", count)
}
if err := os.WriteFile(resolvedPath, []byte(newContent), 0644); err != nil { newContent := strings.Replace(contentStr, oldText, newText, 1)
return ErrorResult(fmt.Sprintf("failed to write file: %v", err))
}
return SilentResult(fmt.Sprintf("File edited: %s", path)) fw, err := root.Create(relPath)
if err != nil {
return nil, fmt.Errorf("failed to create file for writing: %w", err)
}
defer fw.Close()
if _, err := fw.Write([]byte(newContent)); err != nil {
return nil, fmt.Errorf("failed to write file: %v", err)
}
return SilentResult(fmt.Sprintf("File edited: %s", path)), nil
})
} }
type AppendFileTool struct { type AppendFileTool struct {
@ -146,20 +186,33 @@ func (t *AppendFileTool) Execute(ctx context.Context, args map[string]interface{
return ErrorResult("content is required") return ErrorResult("content is required")
} }
resolvedPath, err := validatePath(path, t.workspace, t.restrict) // If not restricted, perform operations directly
if err != nil { if !t.restrict {
return ErrorResult(err.Error()) f, err := os.OpenFile(path, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
if err != nil {
return ErrorResult(fmt.Sprintf("failed to open file: %v", err))
}
defer f.Close()
if _, err := f.WriteString(content); err != nil {
return ErrorResult(fmt.Sprintf("failed to append to file: %v", err))
}
return SilentResult(fmt.Sprintf("Appended to %s", path))
} }
f, err := os.OpenFile(resolvedPath, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644) // Use executeInRoot to safely access the file
if err != nil { return executeInRoot(t.workspace, path, func(root *os.Root, relPath string) (*ToolResult, error) {
return ErrorResult(fmt.Sprintf("failed to open file: %v", err)) f, err := root.OpenFile(relPath, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
} if err != nil {
defer f.Close() return nil, fmt.Errorf("failed to open file: %w", err)
}
defer f.Close()
if _, err := f.WriteString(content); err != nil { if _, err := f.WriteString(content); err != nil {
return ErrorResult(fmt.Sprintf("failed to append to file: %v", err)) return nil, fmt.Errorf("failed to append to file: %w", err)
} }
return SilentResult(fmt.Sprintf("Appended to %s", path)) return SilentResult(fmt.Sprintf("Appended to %s", path)), nil
})
} }

View file

@ -6,6 +6,8 @@ import (
"path/filepath" "path/filepath"
"strings" "strings"
"testing" "testing"
"github.com/stretchr/testify/assert"
) )
// TestEditTool_EditFile_Success verifies successful file editing // TestEditTool_EditFile_Success verifies successful file editing
@ -151,14 +153,13 @@ func TestEditTool_EditFile_OutsideAllowedDir(t *testing.T) {
result := tool.Execute(ctx, args) result := tool.Execute(ctx, args)
// Should return error result // Should return error result
if !result.IsError { assert.True(t, result.IsError, "Expected error when path is outside allowed directory")
t.Errorf("Expected error when path is outside allowed directory")
}
// Should mention outside allowed directory // Should mention outside allowed directory
if !strings.Contains(result.ForLLM, "outside") && !strings.Contains(result.ForUser, "outside") { // Note: ErrorResult only sets ForLLM by default, so ForUser might be empty.
t.Errorf("Expected 'outside allowed' message, got ForLLM: %s", result.ForLLM) // We check ForLLM as it's the primary error channel.
} assert.True(t, strings.Contains(result.ForLLM, "outside") || strings.Contains(result.ForLLM, "access denied") || strings.Contains(result.ForLLM, "escapes"),
"Expected 'outside allowed' or 'access denied' message, got ForLLM: %s", result.ForLLM)
} }
// TestEditTool_EditFile_MissingPath verifies error handling for missing path // TestEditTool_EditFile_MissingPath verifies error handling for missing path

View file

@ -3,78 +3,85 @@ package tools
import ( import (
"context" "context"
"fmt" "fmt"
"io"
"os" "os"
"path/filepath" "path/filepath"
"strings" "strings"
) )
// validatePath ensures the given path is within the workspace if restrict is true. // Helper to get a safe relative path for os.Root usage
func validatePath(path, workspace string, restrict bool) (string, error) { func getSafeRelPath(workspace, path string) (string, error) {
if workspace == "" { if workspace == "" {
return path, nil return "", fmt.Errorf("workspace is not defined")
} }
absWorkspace, err := filepath.Abs(workspace) // Clean the path first
if err != nil { path = filepath.Clean(path)
return "", fmt.Errorf("failed to resolve workspace path: %w", err)
}
var absPath string // If absolute, make it relative to workspace
if filepath.IsAbs(path) { if filepath.IsAbs(path) {
absPath = filepath.Clean(path) rel, err := filepath.Rel(workspace, path)
} else {
absPath, err = filepath.Abs(filepath.Join(absWorkspace, path))
if err != nil { if err != nil {
return "", fmt.Errorf("failed to resolve file path: %w", err) return "", fmt.Errorf("failed to calculate relative path: %w", err)
} }
path = rel
} }
if restrict { // Check for escape
if !isWithinWorkspace(absPath, absWorkspace) { if path == ".." || strings.HasPrefix(path, "../") {
return "", fmt.Errorf("access denied: path is outside the workspace") return "", fmt.Errorf("path escapes workspace: %s", path)
}
workspaceReal := absWorkspace
if resolved, err := filepath.EvalSymlinks(absWorkspace); err == nil {
workspaceReal = resolved
}
if resolved, err := filepath.EvalSymlinks(absPath); err == nil {
if !isWithinWorkspace(resolved, workspaceReal) {
return "", fmt.Errorf("access denied: symlink resolves outside workspace")
}
} else if os.IsNotExist(err) {
if parentResolved, err := resolveExistingAncestor(filepath.Dir(absPath)); err == nil {
if !isWithinWorkspace(parentResolved, workspaceReal) {
return "", fmt.Errorf("access denied: symlink resolves outside workspace")
}
} else if !os.IsNotExist(err) {
return "", fmt.Errorf("failed to resolve path: %w", err)
}
} else {
return "", fmt.Errorf("failed to resolve path: %w", err)
}
} }
return absPath, nil return path, nil
} }
func resolveExistingAncestor(path string) (string, error) { // executeInRoot executes a function within the safety of os.Root
for current := filepath.Clean(path); ; current = filepath.Dir(current) { func executeInRoot(workspace string, path string, fn func(root *os.Root, relPath string) (*ToolResult, error)) *ToolResult {
if resolved, err := filepath.EvalSymlinks(current); err == nil { if workspace == "" {
return resolved, nil return ErrorResult("workspace is not defined")
} else if !os.IsNotExist(err) {
return "", err
}
if filepath.Dir(current) == current {
return "", os.ErrNotExist
}
} }
// 1. Open the Root
root, err := os.OpenRoot(workspace)
if err != nil {
return ErrorResult(fmt.Sprintf("failed to open workspace root: %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
} }
func isWithinWorkspace(candidate, workspace string) bool { // mkdirAllInRoot mimics os.MkdirAll but within os.Root
rel, err := filepath.Rel(filepath.Clean(workspace), filepath.Clean(candidate)) func mkdirAllInRoot(root *os.Root, relPath string) error {
return err == nil && rel != ".." && !strings.HasPrefix(rel, ".."+string(os.PathSeparator)) relPath = filepath.Clean(relPath)
if relPath == "." || relPath == "/" {
return nil
}
dir := filepath.Dir(relPath)
if dir != "." && dir != "/" {
if err := mkdirAllInRoot(root, dir); err != nil {
return err
}
}
err := root.Mkdir(relPath, 0755)
if err != nil && !os.IsExist(err) {
return err
}
return nil
} }
type ReadFileTool struct { type ReadFileTool struct {
@ -113,17 +120,31 @@ func (t *ReadFileTool) Execute(ctx context.Context, args map[string]interface{})
return ErrorResult("path is required") return ErrorResult("path is required")
} }
resolvedPath, err := validatePath(path, t.workspace, t.restrict) // If restriction is disabled, fall back to standard os interactions (insecure but intended)
if err != nil { if !t.restrict {
return ErrorResult(err.Error()) content, err := os.ReadFile(path)
if err != nil {
return ErrorResult(fmt.Sprintf("failed to read file: %v", err))
}
return NewToolResult(string(content))
} }
content, err := os.ReadFile(resolvedPath) return executeInRoot(t.workspace, path, func(root *os.Root, relPath string) (*ToolResult, error) {
if err != nil { f, err := root.Open(relPath)
return ErrorResult(fmt.Sprintf("failed to read file: %v", err)) if err != nil {
} if os.IsNotExist(err) {
return nil, fmt.Errorf("failed to read file:file not found: %s", path)
}
return nil, fmt.Errorf("access denied or failed to open: %w", err)
}
defer f.Close()
return NewToolResult(string(content)) content, err := io.ReadAll(f)
if err != nil {
return nil, fmt.Errorf("failed to read file: %v", err)
}
return NewToolResult(string(content)), nil
})
} }
type WriteFileTool struct { type WriteFileTool struct {
@ -171,21 +192,37 @@ func (t *WriteFileTool) Execute(ctx context.Context, args map[string]interface{}
return ErrorResult("content is required") return ErrorResult("content is required")
} }
resolvedPath, err := validatePath(path, t.workspace, t.restrict) if !t.restrict {
if err != nil { if err := os.MkdirAll(filepath.Dir(path), 0755); err != nil {
return ErrorResult(err.Error()) return ErrorResult(fmt.Sprintf("failed to create directory: %v", err))
}
if err := os.WriteFile(path, []byte(content), 0644); err != nil {
return ErrorResult(fmt.Sprintf("failed to write file: %v", err))
}
return SilentResult(fmt.Sprintf("File written: %s", path))
} }
dir := filepath.Dir(resolvedPath) return executeInRoot(t.workspace, path, func(root *os.Root, relPath string) (*ToolResult, error) {
if err := os.MkdirAll(dir, 0755); err != nil { // Ensure parent directory exists within root using recursive creation
return ErrorResult(fmt.Sprintf("failed to create directory: %v", err)) dir := filepath.Dir(relPath)
} if dir != "." && dir != "/" {
if err := mkdirAllInRoot(root, dir); err != nil {
return nil, fmt.Errorf("failed to create parent directories: %w", err)
}
}
if err := os.WriteFile(resolvedPath, []byte(content), 0644); err != nil { f, err := root.Create(relPath)
return ErrorResult(fmt.Sprintf("failed to write file: %v", err)) if err != nil {
} return nil, fmt.Errorf("failed to create file: %w", err)
}
defer f.Close()
return SilentResult(fmt.Sprintf("File written: %s", path)) _, err = f.WriteString(content)
if err != nil {
return nil, fmt.Errorf("failed to write file: %w", err)
}
return &ToolResult{Silent: true}, nil
})
} }
type ListDirTool struct { type ListDirTool struct {
@ -224,16 +261,39 @@ func (t *ListDirTool) Execute(ctx context.Context, args map[string]interface{})
path = "." path = "."
} }
resolvedPath, err := validatePath(path, t.workspace, t.restrict) if !t.restrict {
if err != nil { entries, err := os.ReadDir(path)
return ErrorResult(err.Error()) if err != nil {
return ErrorResult(fmt.Sprintf("failed to read directory: %v", err))
}
return formatDirEntries(entries)
} }
entries, err := os.ReadDir(resolvedPath) return executeInRoot(t.workspace, path, func(root *os.Root, relPath string) (*ToolResult, error) {
if err != nil { f, err := root.Open(relPath)
return ErrorResult(fmt.Sprintf("failed to read directory: %v", err)) 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)
}
var result strings.Builder
for _, entry := range entries {
if entry.IsDir() {
result.WriteString("DIR: " + entry.Name() + "\n")
} else {
result.WriteString("FILE: " + entry.Name() + "\n")
}
}
return NewToolResult(result.String()), nil
})
}
func formatDirEntries(entries []os.DirEntry) *ToolResult {
result := "" result := ""
for _, entry := range entries { for _, entry := range entries {
if entry.IsDir() { if entry.IsDir() {
@ -242,6 +302,5 @@ func (t *ListDirTool) Execute(ctx context.Context, args map[string]interface{})
result += "FILE: " + entry.Name() + "\n" result += "FILE: " + entry.Name() + "\n"
} }
} }
return NewToolResult(result) return NewToolResult(result)
} }

View file

@ -6,6 +6,8 @@ import (
"path/filepath" "path/filepath"
"strings" "strings"
"testing" "testing"
"github.com/stretchr/testify/assert"
) )
// TestFilesystemTool_ReadFile_Success verifies successful file reading // TestFilesystemTool_ReadFile_Success verifies successful file reading
@ -275,7 +277,29 @@ func TestFilesystemTool_ReadFile_RejectsSymlinkEscape(t *testing.T) {
if !result.IsError { if !result.IsError {
t.Fatalf("expected symlink escape to be blocked") t.Fatalf("expected symlink escape to be blocked")
} }
if !strings.Contains(result.ForLLM, "symlink resolves outside workspace") { // os.Root might return different errors depending on platform/implementation
// but it definitely should error.
// Our wrapper returns "access denied or file not found"
if !strings.Contains(result.ForLLM, "access denied") && !strings.Contains(result.ForLLM, "file not found") && !strings.Contains(result.ForLLM, "no such file") {
t.Fatalf("expected symlink escape error, got: %s", result.ForLLM) t.Fatalf("expected symlink escape error, got: %s", result.ForLLM)
} }
} }
func TestFilesystemTool_EmptyWorkspace_AccessDenied(t *testing.T) {
tool := NewReadFileTool("", true) // restrict=true but workspace=""
// Try to read a sensitive file (simulated by a temp file outside workspace)
tmpDir := t.TempDir()
secretFile := filepath.Join(tmpDir, "shadow")
os.WriteFile(secretFile, []byte("secret data"), 0600)
result := tool.Execute(context.Background(), map[string]any{
"path": secretFile,
})
// We EXPECT IsError=true (access blocked due to empty workspace)
assert.True(t, result.IsError, "Security Regression: Empty workspace allowed access! content: %s", result.ForLLM)
// Verify it failed for the right reason
assert.Contains(t, result.ForLLM, "workspace is not defined", "Expected 'workspace is not defined' error")
}