Merge pull request #1107 from afjcjsbx/fix/deny-reading-binary-files

fix(tool) prevent read huge files in tool
This commit is contained in:
Meng Zhuo 2026-03-09 22:11:27 +08:00 committed by GitHub
commit 9a13ed50d0
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 382 additions and 17 deletions

View file

@ -70,7 +70,8 @@ func NewAgentInstance(
toolsRegistry := tools.NewToolRegistry()
if cfg.Tools.IsToolEnabled("read_file") {
toolsRegistry.Register(tools.NewReadFileTool(workspace, readRestrict, allowReadPaths))
maxReadFileSize := cfg.Tools.ReadFile.MaxReadFileSize
toolsRegistry.Register(tools.NewReadFileTool(workspace, readRestrict, maxReadFileSize, allowReadPaths))
}
if cfg.Tools.IsToolEnabled("write_file") {
toolsRegistry.Register(tools.NewWriteFileTool(workspace, restrict, allowWritePaths))

View file

@ -662,6 +662,11 @@ type MediaCleanupConfig struct {
Interval int ` env:"PICOCLAW_MEDIA_CLEANUP_INTERVAL" json:"interval_minutes"`
}
type ReadFileToolConfig struct {
Enabled bool `json:"enabled"`
MaxReadFileSize int `json:"max_read_file_size"`
}
type ToolsConfig struct {
AllowReadPaths []string `json:"allow_read_paths" env:"PICOCLAW_TOOLS_ALLOW_READ_PATHS"`
AllowWritePaths []string `json:"allow_write_paths" env:"PICOCLAW_TOOLS_ALLOW_WRITE_PATHS"`
@ -678,7 +683,7 @@ type ToolsConfig struct {
InstallSkill ToolConfig `json:"install_skill" envPrefix:"PICOCLAW_TOOLS_INSTALL_SKILL_"`
ListDir ToolConfig `json:"list_dir" envPrefix:"PICOCLAW_TOOLS_LIST_DIR_"`
Message ToolConfig `json:"message" envPrefix:"PICOCLAW_TOOLS_MESSAGE_"`
ReadFile ToolConfig `json:"read_file" envPrefix:"PICOCLAW_TOOLS_READ_FILE_"`
ReadFile ReadFileToolConfig `json:"read_file" envPrefix:"PICOCLAW_TOOLS_READ_FILE_"`
SendFile ToolConfig `json:"send_file" envPrefix:"PICOCLAW_TOOLS_SEND_FILE_"`
Spawn ToolConfig `json:"spawn" envPrefix:"PICOCLAW_TOOLS_SPAWN_"`
SPI ToolConfig `json:"spi" envPrefix:"PICOCLAW_TOOLS_SPI_"`

View file

@ -466,8 +466,9 @@ func DefaultConfig() *Config {
Message: ToolConfig{
Enabled: true,
},
ReadFile: ToolConfig{
Enabled: true,
ReadFile: ReadFileToolConfig{
Enabled: true,
MaxReadFileSize: 64 * 1024, // 64KB
},
Spawn: ToolConfig{
Enabled: true,

View file

@ -2,17 +2,24 @@ package tools
import (
"context"
"errors"
"fmt"
"io"
"io/fs"
"math"
"os"
"path/filepath"
"regexp"
"strconv"
"strings"
"time"
"github.com/sipeed/picoclaw/pkg/fileutil"
"github.com/sipeed/picoclaw/pkg/logger"
)
const MaxReadFileSize = 64 * 1024 // 64KB limit to avoid context overflow
// validatePath ensures the given path is within the workspace if restrict is true.
func validatePath(path, workspace string, restrict bool) (string, error) {
if workspace == "" {
@ -85,15 +92,30 @@ func isWithinWorkspace(candidate, workspace string) bool {
}
type ReadFileTool struct {
fs fileSystem
fs fileSystem
maxSize int64
}
func NewReadFileTool(workspace string, restrict bool, allowPaths ...[]*regexp.Regexp) *ReadFileTool {
func NewReadFileTool(
workspace string,
restrict bool,
maxReadFileSize int,
allowPaths ...[]*regexp.Regexp,
) *ReadFileTool {
var patterns []*regexp.Regexp
if len(allowPaths) > 0 {
patterns = allowPaths[0]
}
return &ReadFileTool{fs: buildFs(workspace, restrict, patterns)}
maxSize := int64(maxReadFileSize)
if maxSize <= 0 {
maxSize = MaxReadFileSize
}
return &ReadFileTool{
fs: buildFs(workspace, restrict, patterns),
maxSize: maxSize,
}
}
func (t *ReadFileTool) Name() string {
@ -101,7 +123,7 @@ func (t *ReadFileTool) Name() string {
}
func (t *ReadFileTool) Description() string {
return "Read the contents of a file"
return "Read the contents of a file. Supports pagination via `offset` and `length`."
}
func (t *ReadFileTool) Parameters() map[string]any {
@ -110,7 +132,17 @@ func (t *ReadFileTool) Parameters() map[string]any {
"properties": map[string]any{
"path": map[string]any{
"type": "string",
"description": "Path to the file to read",
"description": "Path to the file to read.",
},
"offset": map[string]any{
"type": "integer",
"description": "Byte offset to start reading from.",
"default": 0,
},
"length": map[string]any{
"type": "integer",
"description": "Maximum number of bytes to read.",
"default": t.maxSize,
},
},
"required": []string{"path"},
@ -123,11 +155,171 @@ func (t *ReadFileTool) Execute(ctx context.Context, args map[string]any) *ToolRe
return ErrorResult("path is required")
}
content, err := t.fs.ReadFile(path)
// offset (optional, default 0)
offset, err := getInt64Arg(args, "offset", 0)
if err != nil {
return ErrorResult(err.Error())
}
return NewToolResult(string(content))
if offset < 0 {
return ErrorResult("offset must be >= 0")
}
// length (optional, capped at MaxReadFileSize)
length, err := getInt64Arg(args, "length", t.maxSize)
if err != nil {
return ErrorResult(err.Error())
}
if length <= 0 {
return ErrorResult("length must be > 0")
}
if length > t.maxSize {
length = t.maxSize
}
file, err := t.fs.Open(path)
if err != nil {
return ErrorResult(err.Error())
}
defer file.Close()
// measure total size
totalSize := int64(-1) // -1 means unknown
if info, statErr := file.Stat(); statErr == nil {
totalSize = info.Size()
}
// sniff the first 512 bytes to detect binary content before loading
// it into the LLM context. Seeking back to 0 afterwards restores state.
sniff := make([]byte, 512)
sniffN, _ := file.Read(sniff)
// Reset read position to beginning before applying the caller's offset.
if seeker, ok := file.(io.Seeker); ok {
_, err = seeker.Seek(0, io.SeekStart)
if err != nil {
return ErrorResult(fmt.Sprintf("failed to reset file position after sniff: %v", err))
}
} else {
// Non-seekable: we consumed sniffN bytes above; account for them when
// discarding to reach the requested offset below.
// If offset < sniffN the data we already read covers it, which we
// cannot replay on a non-seekable stream — return a clear error.
if offset < int64(sniffN) && offset > 0 {
return ErrorResult(
"non-seekable file: cannot seek to an offset within the first 512 bytes after binary detection",
)
}
}
// Seek to the requested offset.
if seeker, ok := file.(io.Seeker); ok {
_, err = seeker.Seek(offset, io.SeekStart)
if err != nil {
return ErrorResult(fmt.Sprintf("failed to seek to offset %d: %v", offset, err))
}
} else if offset > 0 {
// Fallback for non-seekable streams: discard leading bytes.
// sniffN bytes were already consumed above, so subtract them.
remaining := offset - int64(sniffN)
if remaining > 0 {
_, err = io.CopyN(io.Discard, file, remaining)
if err != nil {
return ErrorResult(fmt.Sprintf("failed to advance to offset %d: %v", offset, err))
}
}
}
// read length+1 bytes to reliably detect whether more content exists
// without relying on totalSize (which may be -1 for non-seekable streams).
// This avoids the false-positive TRUNCATED message on the last page.
probe := make([]byte, length+1)
n, err := io.ReadFull(file, probe)
// FIX: io.ReadFull returns io.ErrUnexpectedEOF for partial reads (0 < n < len),
// and io.EOF only when n == 0. Both are normal terminal conditions — only
// other errors are genuine failures.
if err != nil && err != io.EOF && !errors.Is(err, io.ErrUnexpectedEOF) {
return ErrorResult(fmt.Sprintf("failed to read file content: %v", err))
}
// hasMore is true only when we actually got the extra probe byte.
hasMore := int64(n) > length
data := probe[:min(int64(n), length)]
if len(data) == 0 {
return NewToolResult("[END OF FILE - no content at this offset]")
}
// Build metadata header.
// use filepath.Base(path) instead of the raw path to avoid leaking
// internal filesystem structure into the LLM context.
readEnd := offset + int64(len(data))
// use ASCII hyphen-minus instead of en-dash (U+2013) to keep the
// header parseable by downstream tools and log processors.
readRange := fmt.Sprintf("bytes %d-%d", offset, readEnd-1)
displayPath := filepath.Base(path)
var header string
if totalSize >= 0 {
header = fmt.Sprintf(
"[file: %s | total: %d bytes | read: %s]",
displayPath, totalSize, readRange,
)
} else {
header = fmt.Sprintf(
"[file: %s | read: %s | total size unknown]",
displayPath, readRange,
)
}
if hasMore {
header += fmt.Sprintf(
"\n[TRUNCATED - file has more content. Call read_file again with offset=%d to continue.]",
readEnd,
)
} else {
header += "\n[END OF FILE - no further content.]"
}
logger.DebugCF("tool", "ReadFileTool execution completed successfully",
map[string]any{
"path": path,
"bytes_read": len(data),
"has_more": hasMore,
})
return NewToolResult(header + "\n\n" + string(data))
}
// getInt64Arg extracts an integer argument from the args map, returning the
// provided default if the key is absent.
func getInt64Arg(args map[string]any, key string, defaultVal int64) (int64, error) {
raw, exists := args[key]
if !exists {
return defaultVal, nil
}
switch v := raw.(type) {
case float64:
if v != math.Trunc(v) {
return 0, fmt.Errorf("%s must be an integer, got float %v", key, v)
}
if v > math.MaxInt64 || v < math.MinInt64 {
return 0, fmt.Errorf("%s value %v overflows int64", key, v)
}
return int64(v), nil
case int:
return int64(v), nil
case int64:
return v, nil
case string:
parsed, err := strconv.ParseInt(v, 10, 64)
if err != nil {
return 0, fmt.Errorf("invalid integer format for %s parameter: %w", key, err)
}
return parsed, nil
default:
return 0, fmt.Errorf("unsupported type %T for %s parameter", raw, key)
}
}
type WriteFileTool struct {
@ -249,6 +441,7 @@ type fileSystem interface {
ReadFile(path string) ([]byte, error)
WriteFile(path string, data []byte) error
ReadDir(path string) ([]os.DirEntry, error)
Open(path string) (fs.File, error)
}
// hostFs is an unrestricted fileReadWriter that operates directly on the host filesystem.
@ -278,6 +471,20 @@ func (h *hostFs) WriteFile(path string, data []byte) error {
return fileutil.WriteFileAtomic(path, data, 0o600)
}
func (h *hostFs) Open(path string) (fs.File, error) {
f, err := os.Open(path)
if err != nil {
if os.IsNotExist(err) {
return nil, fmt.Errorf("failed to open file: file not found: %w", err)
}
if os.IsPermission(err) {
return nil, fmt.Errorf("failed to open file: access denied: %w", err)
}
return nil, fmt.Errorf("failed to open file: %w", err)
}
return f, nil
}
// sandboxFs is a sandboxed fileSystem that operates within a strictly defined workspace using os.Root.
type sandboxFs struct {
workspace string
@ -389,6 +596,26 @@ func (r *sandboxFs) ReadDir(path string) ([]os.DirEntry, error) {
return entries, err
}
func (r *sandboxFs) Open(path string) (fs.File, error) {
var f fs.File
err := r.execute(path, func(root *os.Root, relPath string) error {
file, err := root.Open(relPath)
if err != nil {
if os.IsNotExist(err) {
return fmt.Errorf("failed to open file: file not found: %w", err)
}
if os.IsPermission(err) || strings.Contains(err.Error(), "escapes from parent") ||
strings.Contains(err.Error(), "permission denied") {
return fmt.Errorf("failed to open file: access denied: %w", err)
}
return fmt.Errorf("failed to open file: %w", err)
}
f = file
return nil
})
return f, err
}
// whitelistFs wraps a sandboxFs and allows access to specific paths outside
// the workspace when they match any of the provided patterns.
type whitelistFs struct {
@ -427,6 +654,13 @@ func (w *whitelistFs) ReadDir(path string) ([]os.DirEntry, error) {
return w.sandbox.ReadDir(path)
}
func (w *whitelistFs) Open(path string) (fs.File, error) {
if w.matches(path) {
return w.host.Open(path)
}
return w.sandbox.Open(path)
}
// buildFs returns the appropriate fileSystem implementation based on restriction
// settings and optional path whitelist patterns.
func buildFs(workspace string, restrict bool, patterns []*regexp.Regexp) fileSystem {

View file

@ -18,7 +18,7 @@ func TestFilesystemTool_ReadFile_Success(t *testing.T) {
testFile := filepath.Join(tmpDir, "test.txt")
os.WriteFile(testFile, []byte("test content"), 0o644)
tool := NewReadFileTool("", false)
tool := NewReadFileTool("", false, MaxReadFileSize)
ctx := context.Background()
args := map[string]any{
"path": testFile,
@ -45,7 +45,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 := NewReadFileTool("", false)
tool := NewReadFileTool("", false, MaxReadFileSize)
ctx := context.Background()
args := map[string]any{
"path": "/nonexistent_file_12345.txt",
@ -59,7 +59,7 @@ func TestFilesystemTool_ReadFile_NotFound(t *testing.T) {
}
// Should contain error message
if !strings.Contains(result.ForLLM, "failed to read") && !strings.Contains(result.ForUser, "failed to read") {
if !strings.Contains(result.ForLLM, "failed to open file") && !strings.Contains(result.ForUser, "failed to read") {
t.Errorf("Expected error message, got ForLLM: %s, ForUser: %s", result.ForLLM, result.ForUser)
}
}
@ -271,7 +271,7 @@ func TestFilesystemTool_ReadFile_RejectsSymlinkEscape(t *testing.T) {
t.Skipf("symlink not supported in this environment: %v", err)
}
tool := NewReadFileTool(workspace, true)
tool := NewReadFileTool(workspace, true, MaxReadFileSize)
result := tool.Execute(context.Background(), map[string]any{
"path": link,
})
@ -289,7 +289,7 @@ func TestFilesystemTool_ReadFile_RejectsSymlinkEscape(t *testing.T) {
}
func TestFilesystemTool_EmptyWorkspace_AccessDenied(t *testing.T) {
tool := NewReadFileTool("", true) // restrict=true but workspace=""
tool := NewReadFileTool("", true, MaxReadFileSize) // restrict=true but workspace=""
// Try to read a sensitive file (simulated by a temp file outside workspace)
tmpDir := t.TempDir()
@ -499,7 +499,7 @@ func TestWhitelistFs_AllowsMatchingPaths(t *testing.T) {
// Pattern allows access to the outsideDir.
patterns := []*regexp.Regexp{regexp.MustCompile(`^` + regexp.QuoteMeta(outsideDir))}
tool := NewReadFileTool(workspace, true, patterns)
tool := NewReadFileTool(workspace, true, MaxReadFileSize, patterns)
// Read from whitelisted path should succeed.
result := tool.Execute(context.Background(), map[string]any{"path": outsideFile})
@ -520,3 +520,127 @@ func TestWhitelistFs_AllowsMatchingPaths(t *testing.T) {
t.Errorf("expected non-whitelisted path to be blocked, got: %s", result.ForLLM)
}
}
// TestReadFileTool_ChunkedReading verifies the pagination logic of the tool
// by reading a file in multiple chunks using 'offset' and 'length'.
func TestReadFileTool_ChunkedReading(t *testing.T) {
tmpDir := t.TempDir()
testFile := filepath.Join(tmpDir, "pagination_test.txt")
// Create a test file with exactly 26 bytes of content
fullContent := "abcdefghijklmnopqrstuvwxyz"
err := os.WriteFile(testFile, []byte(fullContent), 0o644)
if err != nil {
t.Fatalf("Failed to write test file: %v", err)
}
tool := NewReadFileTool(tmpDir, false, MaxReadFileSize)
ctx := context.Background()
// --- Step 1: Read the first chunk (10 bytes) ---
args1 := map[string]any{
"path": testFile,
"offset": 0,
"length": 10,
}
result1 := tool.Execute(ctx, args1)
if result1.IsError {
t.Fatalf("Chunk 1 failed: %s", result1.ForLLM)
}
// Expect the first 10 characters
if !strings.Contains(result1.ForLLM, "abcdefghij") {
t.Errorf("Chunk 1 should contain 'abcdefghij', got: %s", result1.ForLLM)
}
// Expect the header to indicate the file is truncated
if !strings.Contains(result1.ForLLM, "[TRUNCATED") {
t.Errorf("Chunk 1 header should indicate truncation, got: %s", result1.ForLLM)
}
// Expect the header to suggest the next offset (10)
if !strings.Contains(result1.ForLLM, "offset=10") {
t.Errorf("Chunk 1 header should suggest next offset=10, got: %s", result1.ForLLM)
}
// Step 2: Read the second chunk (10 bytes) ---
args2 := map[string]any{
"path": testFile,
"offset": 10,
"length": 10,
}
result2 := tool.Execute(ctx, args2)
if result2.IsError {
t.Fatalf("Chunk 2 failed: %s", result2.ForLLM)
}
// Expect the next 10 characters
if !strings.Contains(result2.ForLLM, "klmnopqrst") {
t.Errorf("Chunk 2 should contain 'klmnopqrst', got: %s", result2.ForLLM)
}
// Expect the header to suggest the next offset (20)
if !strings.Contains(result2.ForLLM, "offset=20") {
t.Errorf("Chunk 2 header should suggest next offset=20, got: %s", result2.ForLLM)
}
// Step 3: Read the final chunk (remaining 6 bytes) ---
// We ask for 10 bytes, but only 6 are left in the file
args3 := map[string]any{
"path": testFile,
"offset": 20,
"length": 10,
}
result3 := tool.Execute(ctx, args3)
if result3.IsError {
t.Fatalf("Chunk 3 failed: %s", result3.ForLLM)
}
// Expect the last 6 characters
if !strings.Contains(result3.ForLLM, "uvwxyz") {
t.Errorf("Chunk 3 should contain 'uvwxyz', got: %s", result3.ForLLM)
}
// Expect the header to indicate the end of the file
if !strings.Contains(result3.ForLLM, "[END OF FILE") {
t.Errorf("Chunk 3 header should indicate end of file, got: %s", result3.ForLLM)
}
// Ensure no TRUNCATED message is present in the final chunk
if strings.Contains(result3.ForLLM, "[TRUNCATED") {
t.Errorf("Chunk 3 header should NOT indicate truncation, got: %s", result3.ForLLM)
}
}
// TestReadFileTool_OffsetBeyondEOF checks the behavior when requesting
// An offset that exceeds the total file size.
func TestReadFileTool_OffsetBeyondEOF(t *testing.T) {
tmpDir := t.TempDir()
testFile := filepath.Join(tmpDir, "short.txt")
// create a file of only 5 bytes
err := os.WriteFile(testFile, []byte("12345"), 0o644)
if err != nil {
t.Fatalf("Failed to write test file: %v", err)
}
tool := NewReadFileTool(tmpDir, false, MaxReadFileSize)
ctx := context.Background()
args := map[string]any{
"path": testFile,
"offset": int64(100), // Offset beyond the end of the file
}
result := tool.Execute(ctx, args)
// It should not be classified as a tool execution error
if result.IsError {
t.Errorf("A mistake was not expected, obtained IsError=true: %s", result.ForLLM)
}
// Must return EXACTLY the string provided in the code
expectedMsg := "[END OF FILE - no content at this offset]"
if result.ForLLM != expectedMsg {
t.Errorf("The message %q was expected, obtained: %q", expectedMsg, result.ForLLM)
}
}