From 6d4d2ed7a08d7435410458ec072ea8c7a476335d Mon Sep 17 00:00:00 2001
From: XZB-1248 <28593573+XZB-1248@users.noreply.github.com>
Date: Thu, 19 Feb 2026 22:18:26 +0800
Subject: [PATCH] feat: add workspace path validation for sending attachments

---
 pkg/agent/instance.go | 58 ++++++++++++++++++++++---------------------
 pkg/agent/loop.go     |  2 +-
 pkg/tools/message.go  | 38 ++++++++++++++++++++--------
 3 files changed, 58 insertions(+), 40 deletions(-)

diff --git a/pkg/agent/instance.go b/pkg/agent/instance.go
index 54a5396e7..bd9aa0ede 100644
--- a/pkg/agent/instance.go
+++ b/pkg/agent/instance.go
@@ -15,20 +15,21 @@ import (
 // AgentInstance represents a fully configured agent with its own workspace,
 // session manager, context builder, and tool registry.
 type AgentInstance struct {
-	ID             string
-	Name           string
-	Model          string
-	Fallbacks      []string
-	Workspace      string
-	MaxIterations  int
-	ContextWindow  int
-	Provider       providers.LLMProvider
-	Sessions       *session.SessionManager
-	ContextBuilder *ContextBuilder
-	Tools          *tools.ToolRegistry
-	Subagents      *config.SubagentsConfig
-	SkillsFilter   []string
-	Candidates     []providers.FallbackCandidate
+	ID                  string
+	Name                string
+	Model               string
+	Fallbacks           []string
+	Workspace           string
+	MaxIterations       int
+	ContextWindow       int
+	RestrictToWorkspace bool
+	Provider            providers.LLMProvider
+	Sessions            *session.SessionManager
+	ContextBuilder      *ContextBuilder
+	Tools               *tools.ToolRegistry
+	Subagents           *config.SubagentsConfig
+	SkillsFilter        []string
+	Candidates          []providers.FallbackCandidate
 }
 
 // NewAgentInstance creates an agent instance from config.
@@ -84,20 +85,21 @@ func NewAgentInstance(
 	candidates := providers.ResolveCandidates(modelCfg, defaults.Provider)
 
 	return &AgentInstance{
-		ID:             agentID,
-		Name:           agentName,
-		Model:          model,
-		Fallbacks:      fallbacks,
-		Workspace:      workspace,
-		MaxIterations:  maxIter,
-		ContextWindow:  defaults.MaxTokens,
-		Provider:       provider,
-		Sessions:       sessionsManager,
-		ContextBuilder: contextBuilder,
-		Tools:          toolsRegistry,
-		Subagents:      subagents,
-		SkillsFilter:   skillsFilter,
-		Candidates:     candidates,
+		ID:                  agentID,
+		Name:                agentName,
+		Model:               model,
+		Fallbacks:           fallbacks,
+		Workspace:           workspace,
+		MaxIterations:       maxIter,
+		ContextWindow:       defaults.MaxTokens,
+		RestrictToWorkspace: restrict,
+		Provider:            provider,
+		Sessions:            sessionsManager,
+		ContextBuilder:      contextBuilder,
+		Tools:               toolsRegistry,
+		Subagents:           subagents,
+		SkillsFilter:        skillsFilter,
+		Candidates:          candidates,
 	}
 }
 
diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go
index 18190f314..c3d2eb1f7 100644
--- a/pkg/agent/loop.go
+++ b/pkg/agent/loop.go
@@ -106,7 +106,7 @@ func registerSharedTools(cfg *config.Config, msgBus *bus.MessageBus, registry *A
 		agent.Tools.Register(tools.NewSPITool())
 
 		// Message tool
-		messageTool := tools.NewMessageTool()
+		messageTool := tools.NewMessageTool(agent.Workspace, agent.RestrictToWorkspace)
 		messageTool.SetSendCallback(func(channel, chatID, content string, attachments []bus.Attachment) error {
 			msgBus.PublishOutbound(bus.OutboundMessage{
 				Channel:     channel,
diff --git a/pkg/tools/message.go b/pkg/tools/message.go
index c2aedb308..e0e28b5bd 100644
--- a/pkg/tools/message.go
+++ b/pkg/tools/message.go
@@ -10,14 +10,20 @@ import (
 type SendCallback func(channel, chatID, content string, attachments []bus.Attachment) error
 
 type MessageTool struct {
+	allowedDir     string
+	restrict       bool
 	sendCallback   SendCallback
 	defaultChannel string
 	defaultChatID  string
 	sentInRound    bool // Tracks whether a message was sent in the current processing round
 }
 
-func NewMessageTool() *MessageTool {
-	return &MessageTool{}
+// NewMessageTool creates a new MessageTool with optional uploading directory restriction.
+func NewMessageTool(allowedDir string, restrict bool) *MessageTool {
+	return &MessageTool{
+		allowedDir: allowedDir,
+		restrict:   restrict,
+	}
 }
 
 func (t *MessageTool) Name() string {
@@ -110,16 +116,26 @@ func (t *MessageTool) Execute(ctx context.Context, args map[string]interface{})
 	var attachments []bus.Attachment
 	if attachmentsRaw, ok := args["attachments"].([]interface{}); ok {
 		for _, attachRaw := range attachmentsRaw {
-			if attachMap, ok := attachRaw.(map[string]interface{}); ok {
-				path, pathOk := attachMap["path"].(string)
-				filename, filenameOk := attachMap["filename"].(string)
-				if pathOk && filenameOk {
-					attachments = append(attachments, bus.Attachment{
-						Path:     path,
-						Filename: filename,
-					})
-				}
+			attachMap, ok := attachRaw.(map[string]interface{})
+			if !ok {
+				continue // Skip invalid attachment entries
 			}
+
+			path, pathOk := attachMap["path"].(string)
+			filename, filenameOk := attachMap["filename"].(string)
+			if !pathOk || !filenameOk {
+				continue // Skip invalid attachment entries
+			}
+
+			resolvedPath, err := validatePath(path, t.allowedDir, t.restrict)
+			if err != nil {
+				return ErrorResult(err.Error())
+			}
+
+			attachments = append(attachments, bus.Attachment{
+				Path:     resolvedPath,
+				Filename: filename,
+			})
 		}
 	}