diff --git a/pkg/agent/agent_utils.go b/pkg/agent/agent_utils.go index 9228b6d55..cb38fbefe 100644 --- a/pkg/agent/agent_utils.go +++ b/pkg/agent/agent_utils.go @@ -15,6 +15,7 @@ import ( "github.com/sipeed/picoclaw/pkg/config" "github.com/sipeed/picoclaw/pkg/providers" "github.com/sipeed/picoclaw/pkg/session" + "github.com/sipeed/picoclaw/pkg/tools" "github.com/sipeed/picoclaw/pkg/utils" ) @@ -39,6 +40,31 @@ func outboundContextFromInbound( return outboundCtx } +func withMCPHeadersFromRaw(ctx context.Context, raw map[string]string) context.Context { + if len(raw) == 0 { + return ctx + } + var headers map[string]string + for k, v := range raw { + after, ok := strings.CutPrefix(k, "mcp:") + if !ok { + continue + } + after = strings.TrimSpace(after) + if after == "" { + continue + } + if headers == nil { + headers = make(map[string]string) + } + headers[after] = v + } + if len(headers) == 0 { + return ctx + } + return tools.WithMCPHeaders(ctx, headers) +} + func outboundScopeFromSessionScope(scope *session.SessionScope) *bus.OutboundScope { if scope == nil { return nil diff --git a/pkg/agent/pipeline_execute.go b/pkg/agent/pipeline_execute.go index 567e56d17..957b69476 100644 --- a/pkg/agent/pipeline_execute.go +++ b/pkg/agent/pipeline_execute.go @@ -528,6 +528,9 @@ toolLoop: ts.sessionKey, ts.opts.Dispatch.SessionScope, ) + if inbound := ts.opts.Dispatch.InboundContext; inbound != nil { + execCtx = withMCPHeadersFromRaw(execCtx, inbound.Raw) + } toolResult := ts.agent.Tools.ExecuteWithContext( execCtx, toolName, diff --git a/pkg/channels/pico/pico.go b/pkg/channels/pico/pico.go index fefb64ee5..115f52d5a 100644 --- a/pkg/channels/pico/pico.go +++ b/pkg/channels/pico/pico.go @@ -44,6 +44,23 @@ var allowedInlineImageMIMETypes = map[string]struct{}{ "image/bmp": {}, } +// reservedRawKeys are internal metadata keys that cannot be overwritten by client-provided metadata. +// This prevents clients from spoofing internal identity or session information. +var reservedRawKeys = map[string]struct{}{ + "platform": {}, + "session_id": {}, + "conn_id": {}, + "message_id": {}, + "sender_id": {}, + "chat_id": {}, + "message_kind": {}, +} + +func isReservedRawKey(key string) bool { + _, reserved := reservedRawKeys[key] + return reserved +} + func outboundMessageIsThought(msg bus.OutboundMessage) bool { if len(msg.Context.Raw) == 0 { return false @@ -948,6 +965,17 @@ func (c *PicoChannel) handleMessageSend(pc *picoConn, msg PicoMessage) { "conn_id": pc.id, } + if payloadMeta, ok := msg.Payload["metadata"].(map[string]any); ok { + for k, v := range payloadMeta { + if s, ok := v.(string); ok && s != "" { + if isReservedRawKey(k) { + continue + } + metadata[k] = s + } + } + } + logger.DebugCF("pico", "Received message", map[string]any{ "session_id": sessionID, "preview": truncate(content, 50), diff --git a/pkg/channels/pico/pico_test.go b/pkg/channels/pico/pico_test.go index bbe73a222..73d6f2ed5 100644 --- a/pkg/channels/pico/pico_test.go +++ b/pkg/channels/pico/pico_test.go @@ -547,3 +547,32 @@ func newTestPicoWebSocket(t *testing.T) (*websocket.Conn, <-chan PicoMessage, fu defer resp.Body.Close() return clientConn, received, cleanup } + +func TestIsReservedRawKey(t *testing.T) { + reserved := []string{ + "platform", + "session_id", + "conn_id", + "message_id", + "sender_id", + "chat_id", + "message_kind", + } + for _, key := range reserved { + if !isReservedRawKey(key) { + t.Errorf("isReservedRawKey(%q) = false, want true", key) + } + } + + allowed := []string{ + "mcp:Authorization", + "mcp:X-Custom-Header", + "custom_field", + "user_data", + } + for _, key := range allowed { + if isReservedRawKey(key) { + t.Errorf("isReservedRawKey(%q) = true, want false", key) + } + } +} diff --git a/pkg/config/config.go b/pkg/config/config.go index 53eaef059..8c3cfceeb 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -1110,6 +1110,35 @@ func (c *SkillRegistryConfig) DecodeParam(target any) error { return json.Unmarshal(data, target) } +// DynamicHeadersConfig controls which headers can be forwarded from channel context to MCP servers. +// By default (when nil or empty), no dynamic headers are allowed - this is secure by default. +type DynamicHeadersConfig struct { + // Allowed is a list of header names that may be forwarded from channel context. + // Only headers explicitly listed here will be passed to this MCP server. + // Header names are matched case-insensitively. + Allowed []string `json:"allowed,omitempty"` + // MaxCount limits the number of dynamic headers per request (default: 10). + MaxCount int `json:"max_count,omitempty"` + // MaxValueLen limits the maximum length of each header value (default: 4096). + MaxValueLen int `json:"max_value_len,omitempty"` +} + +// GetMaxCount returns the max header count, defaulting to 10. +func (c *DynamicHeadersConfig) GetMaxCount() int { + if c == nil || c.MaxCount <= 0 { + return 10 + } + return c.MaxCount +} + +// GetMaxValueLen returns the max value length, defaulting to 4096. +func (c *DynamicHeadersConfig) GetMaxValueLen() int { + if c == nil || c.MaxValueLen <= 0 { + return 4096 + } + return c.MaxValueLen +} + // MCPServerConfig defines configuration for a single MCP server type MCPServerConfig struct { // Enabled indicates whether this MCP server is active @@ -1132,6 +1161,9 @@ type MCPServerConfig struct { URL string `json:"url,omitempty"` // Headers are HTTP headers to send with requests (sse/http only) Headers map[string]string `json:"headers,omitempty"` + // DynamicHeaders controls which headers from channel context can be forwarded to this server. + // When nil or empty, no dynamic headers are allowed (secure by default). + DynamicHeaders *DynamicHeadersConfig `json:"dynamic_headers,omitempty"` } // MCPConfig defines configuration for all MCP servers diff --git a/pkg/mcp/manager.go b/pkg/mcp/manager.go index 958927767..bd9194fb9 100644 --- a/pkg/mcp/manager.go +++ b/pkg/mcp/manager.go @@ -18,12 +18,14 @@ import ( "github.com/sipeed/picoclaw/pkg/config" runtimeevents "github.com/sipeed/picoclaw/pkg/events" "github.com/sipeed/picoclaw/pkg/logger" + toolshared "github.com/sipeed/picoclaw/pkg/tools/shared" ) // headerTransport is an http.RoundTripper that adds custom headers to requests type headerTransport struct { - base http.RoundTripper - headers map[string]string + base http.RoundTripper + headers map[string]string + dynamicHeaders *config.DynamicHeadersConfig } func expandHomeCommandPath(command string) string { @@ -45,15 +47,16 @@ func expandHomeCommandPath(command string) string { } func (t *headerTransport) RoundTrip(req *http.Request) (*http.Response, error) { - // Clone the request to avoid modifying the original req = req.Clone(req.Context()) - - // Add custom headers for key, value := range t.headers { req.Header.Set(key, value) } - - // Use the base transport + if dynamic := toolshared.MCPHeaders(req.Context()); len(dynamic) > 0 { + filtered := t.filterDynamicHeaders(dynamic) + for key, value := range filtered { + req.Header.Set(key, value) + } + } base := t.base if base == nil { base = http.DefaultTransport @@ -61,6 +64,39 @@ func (t *headerTransport) RoundTrip(req *http.Request) (*http.Response, error) { return base.RoundTrip(req) } +// filterDynamicHeaders applies the allowlist and limits from DynamicHeadersConfig. +// Returns nil if no dynamic headers config is set (secure by default). +func (t *headerTransport) filterDynamicHeaders(headers map[string]string) map[string]string { + if t.dynamicHeaders == nil || len(t.dynamicHeaders.Allowed) == 0 { + return nil + } + + allowedSet := make(map[string]struct{}, len(t.dynamicHeaders.Allowed)) + for _, h := range t.dynamicHeaders.Allowed { + allowedSet[strings.ToLower(h)] = struct{}{} + } + + maxCount := t.dynamicHeaders.GetMaxCount() + maxValueLen := t.dynamicHeaders.GetMaxValueLen() + + result := make(map[string]string) + count := 0 + for key, value := range headers { + if count >= maxCount { + break + } + if _, ok := allowedSet[strings.ToLower(key)]; !ok { + continue + } + if len(value) > maxValueLen { + value = value[:maxValueLen] + } + result[key] = value + count++ + } + return result +} + // loadEnvFile loads environment variables from a file in .env format // Each line should be in the format: KEY=value // Lines starting with # are comments @@ -381,15 +417,14 @@ func connectServer( DisableStandaloneSSE: disableStandaloneSSE, } - // Add custom headers if provided + sseTransport.HTTPClient = &http.Client{ + Transport: &headerTransport{ + base: http.DefaultTransport, + headers: cfg.Headers, + dynamicHeaders: cfg.DynamicHeaders, + }, + } if len(cfg.Headers) > 0 { - // Create a custom HTTP client with header-injecting transport - sseTransport.HTTPClient = &http.Client{ - Transport: &headerTransport{ - base: http.DefaultTransport, - headers: cfg.Headers, - }, - } logger.DebugCF("mcp", "Added custom HTTP headers", map[string]any{ "server": name, diff --git a/pkg/mcp/manager_test.go b/pkg/mcp/manager_test.go index 5789a37a9..b11df13e8 100644 --- a/pkg/mcp/manager_test.go +++ b/pkg/mcp/manager_test.go @@ -5,6 +5,7 @@ import ( "encoding/json" "fmt" "io" + "net/http" "os" "path/filepath" "strings" @@ -17,6 +18,7 @@ import ( "github.com/sipeed/picoclaw/pkg/config" runtimeevents "github.com/sipeed/picoclaw/pkg/events" + toolshared "github.com/sipeed/picoclaw/pkg/tools/shared" ) func TestLoadEnvFile(t *testing.T) { @@ -628,3 +630,300 @@ func (t *scriptedTransport) Close() error { func (t *scriptedTransport) SessionID() string { return t.sessionID } + +func TestHeaderTransport_DynamicHeaders(t *testing.T) { + captured := make(http.Header) + base := roundTripFunc(func(req *http.Request) (*http.Response, error) { + for k, v := range req.Header { + captured[k] = v + } + return &http.Response{StatusCode: 200, Body: http.NoBody}, nil + }) + + transport := &headerTransport{ + base: base, + headers: map[string]string{"X-Static": "from-config"}, + dynamicHeaders: &config.DynamicHeadersConfig{ + Allowed: []string{"Authorization", "X-Custom"}, + }, + } + + ctx := toolshared.WithMCPHeaders(context.Background(), map[string]string{ + "Authorization": "Bearer tok123", + "X-Custom": "dynamic-val", + }) + req, _ := http.NewRequestWithContext(ctx, http.MethodPost, "http://example.com", nil) + + resp, err := transport.RoundTrip(req) + if err != nil { + t.Fatalf("RoundTrip() error = %v", err) + } + resp.Body.Close() + + if got := captured.Get("X-Static"); got != "from-config" { + t.Errorf("X-Static = %q, want %q", got, "from-config") + } + if got := captured.Get("Authorization"); got != "Bearer tok123" { + t.Errorf("Authorization = %q, want %q", got, "Bearer tok123") + } + if got := captured.Get("X-Custom"); got != "dynamic-val" { + t.Errorf("X-Custom = %q, want %q", got, "dynamic-val") + } +} + +func TestHeaderTransport_DynamicOverridesStatic(t *testing.T) { + captured := make(http.Header) + base := roundTripFunc(func(req *http.Request) (*http.Response, error) { + for k, v := range req.Header { + captured[k] = v + } + return &http.Response{StatusCode: 200, Body: http.NoBody}, nil + }) + + transport := &headerTransport{ + base: base, + headers: map[string]string{"Authorization": "Bearer static"}, + dynamicHeaders: &config.DynamicHeadersConfig{ + Allowed: []string{"Authorization"}, + }, + } + + ctx := toolshared.WithMCPHeaders(context.Background(), map[string]string{ + "Authorization": "Bearer dynamic", + }) + req, _ := http.NewRequestWithContext(ctx, http.MethodPost, "http://example.com", nil) + + resp, err := transport.RoundTrip(req) + if err != nil { + t.Fatalf("RoundTrip() error = %v", err) + } + resp.Body.Close() + + if got := captured.Get("Authorization"); got != "Bearer dynamic" { + t.Errorf("Authorization = %q, want dynamic to override static %q", got, "Bearer dynamic") + } +} + +func TestHeaderTransport_NoDynamicHeaders(t *testing.T) { + captured := make(http.Header) + base := roundTripFunc(func(req *http.Request) (*http.Response, error) { + for k, v := range req.Header { + captured[k] = v + } + return &http.Response{StatusCode: 200, Body: http.NoBody}, nil + }) + + transport := &headerTransport{ + base: base, + headers: map[string]string{"X-Static": "val"}, + } + + req, _ := http.NewRequestWithContext(context.Background(), http.MethodPost, "http://example.com", nil) + resp, err := transport.RoundTrip(req) + if err != nil { + t.Fatalf("RoundTrip() error = %v", err) + } + resp.Body.Close() + + if got := captured.Get("X-Static"); got != "val" { + t.Errorf("X-Static = %q, want %q", got, "val") + } + if got := captured.Get("Authorization"); got != "" { + t.Errorf("Authorization should be empty, got %q", got) + } +} + +func TestHeaderTransport_NoAllowlist_BlocksDynamic(t *testing.T) { + captured := make(http.Header) + base := roundTripFunc(func(req *http.Request) (*http.Response, error) { + for k, v := range req.Header { + captured[k] = v + } + return &http.Response{StatusCode: 200, Body: http.NoBody}, nil + }) + + // No dynamicHeaders config = secure by default, blocks all dynamic headers + transport := &headerTransport{ + base: base, + headers: map[string]string{"X-Static": "val"}, + } + + ctx := toolshared.WithMCPHeaders(context.Background(), map[string]string{ + "Authorization": "Bearer should-be-blocked", + "X-Custom": "also-blocked", + }) + req, _ := http.NewRequestWithContext(ctx, http.MethodPost, "http://example.com", nil) + + resp, err := transport.RoundTrip(req) + if err != nil { + t.Fatalf("RoundTrip() error = %v", err) + } + resp.Body.Close() + + if got := captured.Get("X-Static"); got != "val" { + t.Errorf("X-Static = %q, want %q", got, "val") + } + if got := captured.Get("Authorization"); got != "" { + t.Errorf("Authorization should be blocked without allowlist, got %q", got) + } + if got := captured.Get("X-Custom"); got != "" { + t.Errorf("X-Custom should be blocked without allowlist, got %q", got) + } +} + +func TestHeaderTransport_AllowlistFilters(t *testing.T) { + captured := make(http.Header) + base := roundTripFunc(func(req *http.Request) (*http.Response, error) { + for k, v := range req.Header { + captured[k] = v + } + return &http.Response{StatusCode: 200, Body: http.NoBody}, nil + }) + + transport := &headerTransport{ + base: base, + dynamicHeaders: &config.DynamicHeadersConfig{ + Allowed: []string{"X-Allowed-Header"}, + }, + } + + ctx := toolshared.WithMCPHeaders(context.Background(), map[string]string{ + "X-Allowed-Header": "this-passes", + "X-Blocked-Header": "this-is-blocked", + "Authorization": "also-blocked", + }) + req, _ := http.NewRequestWithContext(ctx, http.MethodPost, "http://example.com", nil) + + resp, err := transport.RoundTrip(req) + if err != nil { + t.Fatalf("RoundTrip() error = %v", err) + } + resp.Body.Close() + + if got := captured.Get("X-Allowed-Header"); got != "this-passes" { + t.Errorf("X-Allowed-Header = %q, want %q", got, "this-passes") + } + if got := captured.Get("X-Blocked-Header"); got != "" { + t.Errorf("X-Blocked-Header should be blocked, got %q", got) + } + if got := captured.Get("Authorization"); got != "" { + t.Errorf("Authorization should be blocked, got %q", got) + } +} + +func TestHeaderTransport_AllowlistCaseInsensitive(t *testing.T) { + captured := make(http.Header) + base := roundTripFunc(func(req *http.Request) (*http.Response, error) { + for k, v := range req.Header { + captured[k] = v + } + return &http.Response{StatusCode: 200, Body: http.NoBody}, nil + }) + + transport := &headerTransport{ + base: base, + dynamicHeaders: &config.DynamicHeadersConfig{ + Allowed: []string{"X-Grafana-Token"}, // lowercase in allowlist + }, + } + + ctx := toolshared.WithMCPHeaders(context.Background(), map[string]string{ + "x-grafana-token": "token-value", // different case in request + }) + req, _ := http.NewRequestWithContext(ctx, http.MethodPost, "http://example.com", nil) + + resp, err := transport.RoundTrip(req) + if err != nil { + t.Fatalf("RoundTrip() error = %v", err) + } + resp.Body.Close() + + if got := captured.Get("X-Grafana-Token"); got != "token-value" { + t.Errorf("X-Grafana-Token = %q, want %q (case-insensitive match)", got, "token-value") + } +} + +func TestHeaderTransport_MaxCountLimit(t *testing.T) { + captured := make(http.Header) + base := roundTripFunc(func(req *http.Request) (*http.Response, error) { + for k, v := range req.Header { + captured[k] = v + } + return &http.Response{StatusCode: 200, Body: http.NoBody}, nil + }) + + transport := &headerTransport{ + base: base, + dynamicHeaders: &config.DynamicHeadersConfig{ + Allowed: []string{"X-Header-1", "X-Header-2", "X-Header-3"}, + MaxCount: 2, + }, + } + + ctx := toolshared.WithMCPHeaders(context.Background(), map[string]string{ + "X-Header-1": "val1", + "X-Header-2": "val2", + "X-Header-3": "val3", + }) + req, _ := http.NewRequestWithContext(ctx, http.MethodPost, "http://example.com", nil) + + resp, err := transport.RoundTrip(req) + if err != nil { + t.Fatalf("RoundTrip() error = %v", err) + } + resp.Body.Close() + + // Only 2 headers should be set (MaxCount=2), but we can't predict which due to map iteration order + count := 0 + for _, h := range []string{"X-Header-1", "X-Header-2", "X-Header-3"} { + if captured.Get(h) != "" { + count++ + } + } + if count != 2 { + t.Errorf("Expected exactly 2 headers due to MaxCount=2, got %d", count) + } +} + +func TestHeaderTransport_MaxValueLenLimit(t *testing.T) { + captured := make(http.Header) + base := roundTripFunc(func(req *http.Request) (*http.Response, error) { + for k, v := range req.Header { + captured[k] = v + } + return &http.Response{StatusCode: 200, Body: http.NoBody}, nil + }) + + transport := &headerTransport{ + base: base, + dynamicHeaders: &config.DynamicHeadersConfig{ + Allowed: []string{"X-Long-Header"}, + MaxValueLen: 10, + }, + } + + ctx := toolshared.WithMCPHeaders(context.Background(), map[string]string{ + "X-Long-Header": "this-value-is-way-too-long-and-should-be-truncated", + }) + req, _ := http.NewRequestWithContext(ctx, http.MethodPost, "http://example.com", nil) + + resp, err := transport.RoundTrip(req) + if err != nil { + t.Fatalf("RoundTrip() error = %v", err) + } + resp.Body.Close() + + got := captured.Get("X-Long-Header") + if len(got) != 10 { + t.Errorf("X-Long-Header length = %d, want 10 (truncated)", len(got)) + } + if got != "this-value" { + t.Errorf("X-Long-Header = %q, want %q", got, "this-value") + } +} + +type roundTripFunc func(*http.Request) (*http.Response, error) + +func (f roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) { + return f(req) +} diff --git a/pkg/tools/shared/base.go b/pkg/tools/shared/base.go index 298e1b478..ab8cb7ffa 100644 --- a/pkg/tools/shared/base.go +++ b/pkg/tools/shared/base.go @@ -50,6 +50,7 @@ var ( ctxKeyAgentID = &toolCtxKey{"agentID"} ctxKeySessionKey = &toolCtxKey{"sessionKey"} ctxKeySessionScope = &toolCtxKey{"sessionScope"} + ctxKeyMCPHeaders = &toolCtxKey{"mcpHeaders"} ) // WithToolContext returns a child context carrying channel and chatID. @@ -130,6 +131,33 @@ func ToolSessionScope(ctx context.Context) *session.SessionScope { return session.CloneScope(scope) } +// WithMCPHeaders returns a child context carrying per-request headers for MCP HTTP transports. +// The map is cloned to prevent callers from mutating it after storage. +func WithMCPHeaders(ctx context.Context, headers map[string]string) context.Context { + if len(headers) == 0 { + return ctx + } + clone := make(map[string]string, len(headers)) + for k, v := range headers { + clone[k] = v + } + return context.WithValue(ctx, ctxKeyMCPHeaders, clone) +} + +// MCPHeaders extracts per-request MCP headers from ctx, or nil if unset. +// Returns a clone so callers cannot mutate the stored map. +func MCPHeaders(ctx context.Context) map[string]string { + v, _ := ctx.Value(ctxKeyMCPHeaders).(map[string]string) + if len(v) == 0 { + return nil + } + clone := make(map[string]string, len(v)) + for k, val := range v { + clone[k] = val + } + return clone +} + // AsyncCallback is a function type that async tools use to notify completion. // When an async tool finishes its work, it calls this callback with the result. // diff --git a/pkg/tools/shared_facade.go b/pkg/tools/shared_facade.go index 85bac140a..1604b25e7 100644 --- a/pkg/tools/shared_facade.go +++ b/pkg/tools/shared_facade.go @@ -61,6 +61,14 @@ func WithToolSessionContext( return toolshared.WithToolSessionContext(ctx, agentID, sessionKey, scope) } +func WithMCPHeaders(ctx context.Context, headers map[string]string) context.Context { + return toolshared.WithMCPHeaders(ctx, headers) +} + +func MCPHeaders(ctx context.Context) map[string]string { + return toolshared.MCPHeaders(ctx) +} + func ToolChannel(ctx context.Context) string { return toolshared.ToolChannel(ctx) }