From b73f66d99b175b33a5f63a56e470e8786da5dee6 Mon Sep 17 00:00:00 2001 From: Luan Vu Date: Thu, 12 Mar 2026 20:43:52 +0700 Subject: [PATCH] fix(tools): make MessageTool media path resolution workspace-aware (#169) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit MessageTool.parseMediaPath() was hardcoded to only allow files in /tmp/, while all other filesystem tools (read_file, write_file, edit, exec) use workspace-aware resolvePath() with restrict_to_workspace enforcement. This meant agents could create files in their workspace via write_file but couldn't send them as attachments — only /tmp/ files from create_image/create_audio worked. Replace parseMediaPath() with resolveMediaPath() that: - Reuses resolvePath() for consistent security (symlink, hardlink, traversal) - Honors per-agent workspace + restrict_to_workspace from context - Still allows /tmp/ as fallback (for create_image, create_audio, etc.) - Supports relative paths resolved against workspace - Updates tool description so LLM knows about MEDIA: prefix Co-authored-by: Luvu182 <208665161+Luvu182@users.noreply.github.com> --- cmd/gateway.go | 2 +- internal/tools/message.go | 59 +++++++++++---- internal/tools/message_test.go | 134 +++++++++++++++++++++++++++------ 3 files changed, 153 insertions(+), 42 deletions(-) diff --git a/cmd/gateway.go b/cmd/gateway.go index 945eaee7..583ab1ae 100644 --- a/cmd/gateway.go +++ b/cmd/gateway.go @@ -580,7 +580,7 @@ func runGateway() { toolsReg.Register(tools.NewSessionsSendTool()) // Message tool (send to channels) - toolsReg.Register(tools.NewMessageTool()) + toolsReg.Register(tools.NewMessageTool(workspace, agentCfg.RestrictToWorkspace)) slog.Info("session + message tools registered") // Register legacy tool aliases (backward-compat names from policy.go). diff --git a/internal/tools/message.go b/internal/tools/message.go index 71007517..44669e19 100644 --- a/internal/tools/message.go +++ b/internal/tools/message.go @@ -14,11 +14,15 @@ import ( // MessageTool allows the agent to proactively send messages to channels. type MessageTool struct { - sender ChannelSender - msgBus *bus.MessageBus + workspace string + restrict bool + sender ChannelSender + msgBus *bus.MessageBus } -func NewMessageTool() *MessageTool { return &MessageTool{} } +func NewMessageTool(workspace string, restrict bool) *MessageTool { + return &MessageTool{workspace: workspace, restrict: restrict} +} func (t *MessageTool) SetChannelSender(s ChannelSender) { t.sender = s } func (t *MessageTool) SetMessageBus(b *bus.MessageBus) { t.msgBus = b } @@ -47,7 +51,7 @@ func (t *MessageTool) Parameters() map[string]any { }, "message": map[string]any{ "type": "string", - "description": "Message content to send", + "description": "Message content to send. To send a file as attachment, use the prefix MEDIA: followed by the file path, e.g. 'MEDIA:docs/report.pdf' or 'MEDIA:/tmp/image.png'. The file will be uploaded as a document/photo/audio depending on its type.", }, }, "required": []string{"action", "message"}, @@ -82,7 +86,7 @@ func (t *MessageTool) Execute(ctx context.Context, args map[string]any) *Result } // Handle MEDIA: prefix — send file as attachment instead of text. - if filePath, ok := parseMediaPath(message); ok { + if filePath, ok := t.resolveMediaPath(ctx, message); ok { return t.sendMedia(ctx, channel, target, filePath) } @@ -160,24 +164,47 @@ func isGroupContext(ctx context.Context) bool { strings.HasPrefix(userID, "guild:") } -// parseMediaPath extracts a file path from a "MEDIA:/path/to/file" string. -// Only allows absolute paths within os.TempDir() to prevent path traversal. -func parseMediaPath(s string) (string, bool) { +// resolveMediaPath extracts and validates a file path from a "MEDIA:path" string. +// Uses the same workspace-aware path resolution as other filesystem tools: +// - When restrict_to_workspace is true: allows workspace dir + /tmp/ +// - When restrict_to_workspace is false: allows any valid path +// +// Relative paths are resolved against the agent's workspace. +func (t *MessageTool) resolveMediaPath(ctx context.Context, s string) (string, bool) { s = strings.TrimSpace(s) if !strings.HasPrefix(s, "MEDIA:") { return "", false } - path := filepath.Clean(strings.TrimSpace(s[len("MEDIA:"):])) - if path == "" || path == "." { + raw := strings.TrimSpace(s[len("MEDIA:"):]) + if raw == "" || raw == "." { return "", false } - if !filepath.IsAbs(path) { + + workspace := ToolWorkspaceFromCtx(ctx) + if workspace == "" { + workspace = t.workspace + } + restrict := effectiveRestrict(ctx, t.restrict) + + // resolvePath handles relative→absolute, symlink, hardlink, boundary checks. + resolved, err := resolvePath(raw, workspace, restrict) + if err != nil { + // When restricted, also allow /tmp/ (used by create_image, create_audio, etc.) + if restrict && isInTempDir(raw) { + return filepath.Clean(raw), true + } return "", false } - // Restrict to temp directory to prevent path traversal. + + return resolved, true +} + +// isInTempDir checks whether an absolute path is inside os.TempDir(). +func isInTempDir(path string) bool { + cleaned := filepath.Clean(path) + if !filepath.IsAbs(cleaned) { + return false + } tmpDir := filepath.Clean(os.TempDir()) - if !strings.HasPrefix(path, tmpDir+string(filepath.Separator)) && path != tmpDir { - return "", false - } - return path, true + return strings.HasPrefix(cleaned, tmpDir+string(filepath.Separator)) } diff --git a/internal/tools/message_test.go b/internal/tools/message_test.go index 91506eb7..27b340a6 100644 --- a/internal/tools/message_test.go +++ b/internal/tools/message_test.go @@ -1,42 +1,126 @@ package tools import ( + "context" "os" "path/filepath" "testing" ) -func TestParseMediaPath(t *testing.T) { +func TestResolveMediaPath(t *testing.T) { tmpDir := os.TempDir() - tests := []struct { - name string - input string - want string - wantOK bool - }{ - {"valid temp file", "MEDIA:" + filepath.Join(tmpDir, "test.png"), filepath.Join(tmpDir, "test.png"), true}, - {"valid nested temp file", "MEDIA:" + filepath.Join(tmpDir, "sub", "file.txt"), filepath.Join(tmpDir, "sub", "file.txt"), true}, - {"with spaces around prefix", " MEDIA:" + filepath.Join(tmpDir, "test.png") + " ", filepath.Join(tmpDir, "test.png"), true}, - {"no prefix", filepath.Join(tmpDir, "test.png"), "", false}, - {"empty after prefix", "MEDIA:", "", false}, - {"relative path", "MEDIA:relative/path.txt", "", false}, - {"dot path", "MEDIA:.", "", false}, - {"outside temp dir", "MEDIA:/etc/passwd", "", false}, - {"traversal attack", "MEDIA:" + filepath.Join(tmpDir, "..", "etc", "passwd"), "", false}, - {"home directory", "MEDIA:" + filepath.Join(os.Getenv("HOME"), "secret.txt"), "", false}, - {"empty string", "", "", false}, - {"just MEDIA", "MEDIA", "", false}, + // Create a temp workspace with a test file for workspace-relative tests. + workspace := t.TempDir() + docsDir := filepath.Join(workspace, "docs") + if err := os.MkdirAll(docsDir, 0o755); err != nil { + t.Fatal(err) + } + testFile := filepath.Join(docsDir, "report.pdf") + if err := os.WriteFile(testFile, []byte("test"), 0o644); err != nil { + t.Fatal(err) } + t.Run("restricted", func(t *testing.T) { + tool := NewMessageTool(workspace, true) + ctx := context.Background() + + tests := []struct { + name string + input string + want string + wantOK bool + }{ + // /tmp/ always allowed + {"valid temp file", "MEDIA:" + filepath.Join(tmpDir, "test.png"), filepath.Join(tmpDir, "test.png"), true}, + {"valid nested temp", "MEDIA:" + filepath.Join(tmpDir, "sub", "file.txt"), filepath.Join(tmpDir, "sub", "file.txt"), true}, + + // Workspace files allowed + {"workspace absolute", "MEDIA:" + testFile, testFile, true}, + {"workspace relative", "MEDIA:docs/report.pdf", testFile, true}, + + // Not a MEDIA: message + {"no prefix", filepath.Join(tmpDir, "test.png"), "", false}, + {"empty after prefix", "MEDIA:", "", false}, + {"dot path", "MEDIA:.", "", false}, + {"empty string", "", "", false}, + {"just MEDIA", "MEDIA", "", false}, + + // Outside workspace + outside /tmp/ → blocked + {"outside workspace", "MEDIA:/etc/passwd", "", false}, + {"traversal attack", "MEDIA:" + filepath.Join(workspace, "..", "etc", "passwd"), "", false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, ok := tool.resolveMediaPath(ctx, tt.input) + if ok != tt.wantOK { + t.Errorf("resolveMediaPath(%q) ok = %v, want %v", tt.input, ok, tt.wantOK) + } + if ok && got != tt.want { + t.Errorf("resolveMediaPath(%q) = %q, want %q", tt.input, got, tt.want) + } + }) + } + }) + + t.Run("unrestricted", func(t *testing.T) { + tool := NewMessageTool(workspace, false) + ctx := context.Background() + + tests := []struct { + name string + input string + wantOK bool + }{ + {"any absolute path", "MEDIA:/etc/hostname", true}, + {"workspace relative", "MEDIA:docs/report.pdf", true}, + {"temp file", "MEDIA:" + filepath.Join(tmpDir, "test.png"), true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + _, ok := tool.resolveMediaPath(ctx, tt.input) + if ok != tt.wantOK { + t.Errorf("resolveMediaPath(%q) ok = %v, want %v", tt.input, ok, tt.wantOK) + } + }) + } + }) + + t.Run("context workspace override", func(t *testing.T) { + // Tool has no workspace, but context provides one. + tool := NewMessageTool("", true) + ctx := WithToolWorkspace(context.Background(), workspace) + + got, ok := tool.resolveMediaPath(ctx, "MEDIA:docs/report.pdf") + if !ok { + t.Fatal("expected ok=true for workspace-relative path with context workspace") + } + if got != testFile { + t.Errorf("got %q, want %q", got, testFile) + } + }) +} + +func TestIsInTempDir(t *testing.T) { + tmpDir := os.TempDir() + tests := []struct { + name string + path string + want bool + }{ + {"in tmp", filepath.Join(tmpDir, "test.png"), true}, + {"nested in tmp", filepath.Join(tmpDir, "sub", "file.txt"), true}, + {"tmp itself", tmpDir, false}, // only files inside, not the dir itself + {"outside tmp", "/etc/passwd", false}, + {"relative path", "relative/path.txt", false}, + {"traversal", filepath.Join(tmpDir, "..", "etc", "passwd"), false}, + } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - got, ok := parseMediaPath(tt.input) - if ok != tt.wantOK { - t.Errorf("parseMediaPath(%q) ok = %v, want %v", tt.input, ok, tt.wantOK) - } - if got != tt.want { - t.Errorf("parseMediaPath(%q) = %q, want %q", tt.input, got, tt.want) + if got := isInTempDir(tt.path); got != tt.want { + t.Errorf("isInTempDir(%q) = %v, want %v", tt.path, got, tt.want) } }) }