From 6e94710384ee2be0cf8d3e47f073d41e344702ce Mon Sep 17 00:00:00 2001 From: Vasanthdev2004 Date: Sat, 7 Nov 2026 14:47:58 +0530 Subject: [PATCH 1/2] fix(mcp): scrub the injected AI/ML API key from server responses Zero injects the user's AI/ML API key into requests to mcp.aimlapi.com (#1193) and keeps it out of config, but only an HTTP error status left the server's text out. A JSON-RPC error message, the method of an unexpected message, a tool result (MCP reports a failed tool call as a result with isError, not as an error) and a legacy SSE endpoint that fails to parse all reached the caller as the server wrote them, so a server that echoed the key put it in errors, in tool output sent to the model, and in anything that prints them. A response to a request that carried the injected key now has the key replaced with [REDACTED] in its error text, its method and every string of its result before anything reads it, on streamable HTTP and on the legacy SSE transport. The SSE client scrubs both the key its stream request carried and the one on the POST, since the key is resolved per request. A user Authorization header, OAuth and project-configured servers are unchanged, since nothing was injected. The scrub uses a new redaction.RedactSecretValues: the split-aware known-value pass RedactString applies to ExtraSecretValues, without the shape patterns, which could otherwise rewrite tool output or base64 payloads that never held the key. A result without the key keeps its original bytes. --- docs/EXTENDING.md | 4 + internal/mcp/aimlapi.go | 126 +++++- internal/mcp/aimlapi_scrub_test.go | 414 ++++++++++++++++++ internal/mcp/network_client.go | 50 ++- .../redaction/redact_secret_values_test.go | 40 ++ internal/redaction/redaction.go | 13 + 6 files changed, 626 insertions(+), 21 deletions(-) create mode 100644 internal/mcp/aimlapi_scrub_test.go create mode 100644 internal/redaction/redact_secret_values_test.go diff --git a/docs/EXTENDING.md b/docs/EXTENDING.md index a2906f39e..d193b4c7d 100644 --- a/docs/EXTENDING.md +++ b/docs/EXTENDING.md @@ -303,6 +303,10 @@ attribution headers `X-AIMLAPI-Source: agent/zero` and `X-AIMLAPI-Partner-ID`, and gets the key under the rules above. Other MCP servers never receive these headers or the key. +If the server echoes back the key Zero sent it, Zero replaces the key with +`[REDACTED]` in that server's error messages and tool results before anything +reads them, and reports an HTTP error status without the response body. + ### As a server — expose Zero's tools to another agent ```bash diff --git a/internal/mcp/aimlapi.go b/internal/mcp/aimlapi.go index c3fa73701..98b36bf06 100644 --- a/internal/mcp/aimlapi.go +++ b/internal/mcp/aimlapi.go @@ -1,6 +1,9 @@ package mcp import ( + "bytes" + "encoding/json" + "errors" "fmt" "net/http" "net/url" @@ -8,6 +11,7 @@ import ( "github.com/Twigpine/zero/internal/aimlapi" "github.com/Twigpine/zero/internal/config" + "github.com/Twigpine/zero/internal/redaction" ) // The AI/ML API MCP server authenticates with the user's AI/ML API key. @@ -21,10 +25,12 @@ import ( // be sent (no proxy profile, a single token), so a provider change applies to // the next request and a withheld or invalid key means no Authorization // header. It lives only in that request's headers: never in Server, its -// Identity (OAuth token binding), config.json, list output or errors (an -// error response to such a request is reported without its body). A -// user-configured Authorization header, OAuth, or a project-touched server -// always goes without it. +// Identity (OAuth token binding), config.json, list output, errors or tool +// output. An error status to such a request is reported without its body, and +// any other response to it has the key scrubbed from its error text, method +// and result before anything reads them (scrubAIMLAPIKeys), since that text is +// the server's own. A user-configured Authorization header, OAuth, or a +// project-touched server always goes without it. var loadAIMLAPIServiceKey = config.LoadAIMLAPIServiceKey var aimlapiMCPHost = func() string { @@ -101,6 +107,118 @@ func hasUserAuthorizationHeader(server Server) bool { return false } +// injectedAIMLAPIKey returns the key applyAIMLAPIRequestHeaders put on request, +// or "" when request went out without it. Pass the request a response answers +// (response.Request, the last hop of any redirect): that is what the responding +// server was sent. +func injectedAIMLAPIKey(server Server, request *http.Request) string { + if !sentAIMLAPIServiceKey(server, request) { + return "" + } + value := strings.TrimSpace(request.Header.Get("Authorization")) + if scheme, key, ok := strings.Cut(value, " "); ok && strings.EqualFold(scheme, "Bearer") { + value = strings.TrimSpace(key) + } + return value +} + +// scrubAIMLAPIKeys removes the injected keys from a response before any of it +// is read: the error text, the method name and every string in the result, +// object member names included. A server can echo the credential it was sent, +// and MCP reports a failed tool call as a result rather than an error, so the +// result needs it as much as the error text. Only the known values go, split +// forms included (redaction.RedactSecretValues); a result that holds none of +// them keeps its original bytes. +func scrubAIMLAPIKeys(message rpcMessage, keys ...string) rpcMessage { + var sent []string + for _, key := range keys { + if strings.TrimSpace(key) != "" { + sent = append(sent, key) + } + } + if len(sent) == 0 { + return message + } + if message.Error != nil { + scrubbed := *message.Error + scrubbed.Message = redaction.RedactSecretValues(scrubbed.Message, sent...) + message.Error = &scrubbed + } + message.Method = redaction.RedactSecretValues(message.Method, sent...) + message.Result = scrubAIMLAPIKeysFromJSON(message.Result, sent) + return message +} + +func scrubAIMLAPIKeysFromJSON(raw json.RawMessage, keys []string) json.RawMessage { + if len(raw) == 0 { + return raw + } + decoder := json.NewDecoder(bytes.NewReader(raw)) + decoder.UseNumber() + var value any + if err := decoder.Decode(&value); err != nil { + // Decoding it into the caller's type fails the same way, and that + // error does not quote the text. + return raw + } + scrubbed, changed := scrubAIMLAPIKeysFromValue(value, keys) + if !changed { + return raw + } + encoded, err := json.Marshal(scrubbed) + if err != nil { + // Never hand back the unscrubbed bytes: an empty result reads as + // nothing rather than as the key. + return nil + } + return encoded +} + +func scrubAIMLAPIKeysFromValue(value any, keys []string) (any, bool) { + switch typed := value.(type) { + case string: + scrubbed := redaction.RedactSecretValues(typed, keys...) + return scrubbed, scrubbed != typed + case []any: + changed := false + for index, item := range typed { + if scrubbed, itemChanged := scrubAIMLAPIKeysFromValue(item, keys); itemChanged { + typed[index] = scrubbed + changed = true + } + } + return typed, changed + case map[string]any: + changed := false + scrubbed := make(map[string]any, len(typed)) + for name, item := range typed { + scrubbedName := redaction.RedactSecretValues(name, keys...) + scrubbedItem, itemChanged := scrubAIMLAPIKeysFromValue(item, keys) + if itemChanged || scrubbedName != name { + changed = true + } + scrubbed[scrubbedName] = scrubbedItem + } + return scrubbed, changed + } + return value, false +} + +// scrubAIMLAPIKeyError returns err with the injected keys removed from its text. +// An error the keys do not appear in keeps its chain; one they do appear in is +// replaced by its scrubbed text, because its chain still holds the key. +func scrubAIMLAPIKeyError(err error, keys ...string) error { + if err == nil { + return nil + } + text := err.Error() + scrubbed := redaction.RedactSecretValues(text, keys...) + if scrubbed == text { + return err + } + return errors.New(scrubbed) +} + // aimlapiUnauthorizedHint explains how to authenticate when the AI/ML API MCP // server rejects a request without OAuth. It names no key material. func aimlapiUnauthorizedHint(server Server, status int) string { diff --git a/internal/mcp/aimlapi_scrub_test.go b/internal/mcp/aimlapi_scrub_test.go new file mode 100644 index 000000000..1d1fe9014 --- /dev/null +++ b/internal/mcp/aimlapi_scrub_test.go @@ -0,0 +1,414 @@ +package mcp + +import ( + "context" + "encoding/json" + "io" + "net/http" + "strings" + "sync" + "testing" + "time" + + "github.com/Twigpine/zero/internal/aimlapi" +) + +// scrubTestKey has no credential shape, so only the known-value scrub can find +// it; a shape pattern would leave it alone. +const scrubTestKey = "aiml0123456789abcdef0123456789abcd" + +// aimlapiEcho stands in for mcp.aimlapi.com over streamable HTTP. It answers +// each JSON-RPC request with reply(method, id), as a JSON body or as one SSE +// event, and records the Authorization header every request carried. +type aimlapiEcho struct { + sse bool + reply func(method string, id json.RawMessage) string + mu sync.Mutex + auth []string +} + +func (echo *aimlapiEcho) RoundTrip(request *http.Request) (*http.Response, error) { + var body []byte + if request.Body != nil { + body, _ = io.ReadAll(request.Body) + _ = request.Body.Close() + } + echo.mu.Lock() + echo.auth = append(echo.auth, request.Header.Get("Authorization")) + echo.mu.Unlock() + var message struct { + ID json.RawMessage `json:"id"` + Method string `json:"method"` + } + _ = json.Unmarshal(body, &message) + if len(message.ID) == 0 { + return fakeHTTPResponse(request, http.StatusAccepted, ""), nil + } + reply := echo.reply(message.Method, message.ID) + response := fakeHTTPResponse(request, http.StatusOK, reply) + if echo.sse { + response.Header.Set("Content-Type", "text/event-stream") + response.Body = io.NopCloser(strings.NewReader("event: message\ndata: " + reply + "\n\n")) + } + return response, nil +} + +// assertSent fails unless every request carried want as its Authorization, so a +// test cannot pass by never reaching the injected-key path. +func (echo *aimlapiEcho) assertSent(t *testing.T, want string) { + t.Helper() + echo.mu.Lock() + defer echo.mu.Unlock() + if len(echo.auth) == 0 { + t.Fatal("the fake server was never called") + } + for _, got := range echo.auth { + if got != want { + t.Fatalf("Authorization = %q, want %q", got, want) + } + } +} + +func rpcReply(id json.RawMessage, field string, value any) string { + encoded, err := json.Marshal(map[string]any{"jsonrpc": "2.0", "id": id, field: value}) + if err != nil { + panic(err) + } + return string(encoded) +} + +// echoedKeys are the spellings of the key a server echoes back: whole, and +// split by a NUL and by a zero-width space, which a plain replace would miss. +func echoedKeys() []string { + return []string{ + scrubTestKey, + scrubTestKey[:10] + string(rune(0)) + scrubTestKey[10:], + scrubTestKey[:20] + string(rune(0x200b)) + scrubTestKey[20:], + } +} + +// assertNoPartOf fails if any eight-byte run of key survives in text. +func assertNoPartOf(t *testing.T, text string, key string) { + t.Helper() + for start := 0; start+8 <= len(key); start++ { + if strings.Contains(text, key[start:start+8]) { + t.Fatalf("text still holds part of the key (%q): %q", key[start:start+8], text) + } + } +} + +// assertKeyScrubbed fails if any part of scrubTestKey survives in text, or if +// text does not show the replacement. +func assertKeyScrubbed(t *testing.T, text string) { + t.Helper() + assertNoPartOf(t, text, scrubTestKey) + if !strings.Contains(text, "[REDACTED]") { + t.Fatalf("text = %q, want the key replaced by [REDACTED]", text) + } +} + +func TestScrubAIMLAPIKeysKeepsAResultWithoutTheKeyByteForByte(t *testing.T) { + raw := json.RawMessage(`{"tools": [ {"name": "b", "description": "password=hunter2"} ], "a": 1.50}`) + if got := scrubAIMLAPIKeysFromJSON(raw, []string{scrubTestKey}); string(got) != string(raw) { + t.Fatalf("result without the key = %s, want the original bytes %s", got, raw) + } + + withKey := json.RawMessage(`{"content": [ {"type": "text", "text": "key ` + scrubTestKey + `"} ], "n": 1.50}`) + got := scrubAIMLAPIKeysFromJSON(withKey, []string{scrubTestKey}) + assertKeyScrubbed(t, string(got)) + var decoded struct { + Content []Content `json:"content"` + N json.Number `json:"n"` + } + if err := json.Unmarshal(got, &decoded); err != nil || len(decoded.Content) != 1 || decoded.Content[0].Text != "key [REDACTED]" || decoded.N != "1.50" { + t.Fatalf("scrubbed result = %s (%v), want valid JSON with only the key replaced and the number as written", got, err) + } +} + +func newEchoNetworkClient(server Server, echo *aimlapiEcho) *networkClient { + return &networkClient{server: server, client: mcpHTTPClient(server, echo), nextID: 1} +} + +func TestAIMLAPIMCPScrubsTheInjectedKeyFromResponses(t *testing.T) { + isolateAIMLAPIKeyLookup(t, scrubTestKey) + server := aimlapiTestServer("aimlapi", aimlapi.MCPServerURL) + for _, sse := range []bool{false, true} { + transport := "json" + if sse { + transport = "sse" + } + t.Run(transport+"/error message", func(t *testing.T) { + for _, echoed := range echoedKeys() { + echo := &aimlapiEcho{sse: sse, reply: func(_ string, id json.RawMessage) string { + return rpcReply(id, "error", map[string]any{"code": -32001, "message": "invalid credential " + echoed + " for this account"}) + }} + err := newEchoNetworkClient(server, echo).initialize(context.Background()) + if err == nil { + t.Fatal("initialize succeeded against a server that answered with an error") + } + echo.assertSent(t, "Bearer "+scrubTestKey) + assertKeyScrubbed(t, err.Error()) + if !strings.Contains(err.Error(), "invalid credential [REDACTED] for this account") { + t.Fatalf("error = %q, want the rest of the server text kept", err) + } + } + }) + t.Run(transport+"/tool result", func(t *testing.T) { + const image = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg==" + for _, echoed := range echoedKeys() { + echo := &aimlapiEcho{sse: sse, reply: func(_ string, id json.RawMessage) string { + return rpcReply(id, "result", map[string]any{ + "isError": true, + "content": []any{ + map[string]any{"type": "text", "text": "upstream rejected key " + echoed}, + map[string]any{"type": "image", "mimeType": "image/png", "data": image}, + }, + }) + }} + result, err := newEchoNetworkClient(server, echo).CallTool(context.Background(), "generate", map[string]any{"prompt": "x"}) + if err != nil { + t.Fatalf("CallTool: %v", err) + } + echo.assertSent(t, "Bearer "+scrubTestKey) + if len(result.Content) != 2 || !result.IsError { + t.Fatalf("result = %+v, want the error result with both blocks", result) + } + assertKeyScrubbed(t, result.Content[0].Text) + if result.Content[0].Text != "upstream rejected key [REDACTED]" { + t.Fatalf("text block = %q, want only the key replaced", result.Content[0].Text) + } + if result.Content[1].Data != image { + t.Fatalf("image data changed: %q", result.Content[1].Data) + } + } + }) + t.Run(transport+"/result without the key", func(t *testing.T) { + // Shapes stay as the server wrote them: only the known value goes. + const description = "Set password=hunter2 or pass Authorization: Bearer other-token" + echo := &aimlapiEcho{sse: sse, reply: func(_ string, id json.RawMessage) string { + return rpcReply(id, "result", map[string]any{"tools": []any{map[string]any{ + "name": "generate", + "description": description, + "inputSchema": map[string]any{"type": "object", "properties": map[string]any{"api_key": map[string]any{"type": "string"}}}, + }}}) + }} + tools, err := newEchoNetworkClient(server, echo).ListTools(context.Background()) + if err != nil { + t.Fatalf("ListTools: %v", err) + } + echo.assertSent(t, "Bearer "+scrubTestKey) + if len(tools) != 1 || tools[0].Description != description { + t.Fatalf("tools = %+v, want the description unchanged", tools) + } + properties, _ := tools[0].InputSchema["properties"].(map[string]any) + if _, ok := properties["api_key"]; !ok { + t.Fatalf("inputSchema = %v, want the api_key property kept", tools[0].InputSchema) + } + }) + } + + t.Run("json/unexpected method", func(t *testing.T) { + echo := &aimlapiEcho{reply: func(string, json.RawMessage) string { + return `{"jsonrpc":"2.0","method":"notifications/` + scrubTestKey + `"}` + }} + err := newEchoNetworkClient(server, echo).initialize(context.Background()) + if err == nil || !strings.Contains(err.Error(), "got method") { + t.Fatalf("initialize error = %v, want the unexpected method reported", err) + } + echo.assertSent(t, "Bearer "+scrubTestKey) + assertKeyScrubbed(t, err.Error()) + }) +} + +func TestAIMLAPIMCPLeavesResponsesAloneWithoutTheInjectedKey(t *testing.T) { + isolateAIMLAPIKeyLookup(t, scrubTestKey) + // A user-configured Authorization header means nothing was injected, so the + // server text is reported as written; this change scrubs only what Zero + // itself attached. + server := aimlapiTestServer("aimlapi", aimlapi.MCPServerURL) + server.Headers = map[string]string{"Authorization": "Bearer " + scrubTestKey} + echo := &aimlapiEcho{reply: func(_ string, id json.RawMessage) string { + return rpcReply(id, "error", map[string]any{"code": -32001, "message": "invalid credential " + scrubTestKey}) + }} + err := newEchoNetworkClient(server, echo).initialize(context.Background()) + echo.assertSent(t, "Bearer "+scrubTestKey) + if err == nil || !strings.Contains(err.Error(), "invalid credential "+scrubTestKey) { + t.Fatalf("error = %v, want the server text unchanged when no key was injected", err) + } +} + +// aimlapiSSEEcho stands in for mcp.aimlapi.com over the legacy SSE transport: +// the GET opens a stream that starts with the endpoint event, and each POST is +// answered on that stream with reply(method, id). +type aimlapiSSEEcho struct { + endpoint string + reply func(method string, id json.RawMessage) string + events chan string + mu sync.Mutex + auth []string +} + +func newAIMLAPISSEEcho(endpoint string, reply func(string, json.RawMessage) string) *aimlapiSSEEcho { + return &aimlapiSSEEcho{endpoint: endpoint, reply: reply, events: make(chan string, 8)} +} + +func (echo *aimlapiSSEEcho) RoundTrip(request *http.Request) (*http.Response, error) { + echo.mu.Lock() + echo.auth = append(echo.auth, request.Header.Get("Authorization")) + echo.mu.Unlock() + if request.Method == http.MethodGet { + reader, writer := io.Pipe() + go func() { + defer writer.Close() + if _, err := io.WriteString(writer, "event: endpoint\ndata: "+echo.endpoint+"\n\n"); err != nil { + return + } + for { + select { + case event := <-echo.events: + if _, err := io.WriteString(writer, event); err != nil { + return + } + case <-request.Context().Done(): + return + } + } + }() + response := fakeHTTPResponse(request, http.StatusOK, "") + response.Header.Set("Content-Type", "text/event-stream") + response.Body = reader + return response, nil + } + var body []byte + if request.Body != nil { + body, _ = io.ReadAll(request.Body) + _ = request.Body.Close() + } + var message struct { + ID json.RawMessage `json:"id"` + Method string `json:"method"` + } + _ = json.Unmarshal(body, &message) + if len(message.ID) > 0 { + echo.events <- "event: message\ndata: " + echo.reply(message.Method, message.ID) + "\n\n" + } + return fakeHTTPResponse(request, http.StatusAccepted, ""), nil +} + +func (echo *aimlapiSSEEcho) assertSent(t *testing.T, want string, requests int) { + t.Helper() + echo.mu.Lock() + defer echo.mu.Unlock() + if len(echo.auth) < requests { + t.Fatalf("the fake server saw %d requests, want at least %d", len(echo.auth), requests) + } + for _, got := range echo.auth { + if got != want { + t.Fatalf("Authorization = %q, want %q", got, want) + } + } +} + +func newEchoSSEClient(echo *aimlapiSSEEcho) *remoteSSEClient { + server := Server{Name: "aimlapi", Type: ServerTypeSSE, URL: aimlapi.MCPServerURL} + return &remoteSSEClient{ + server: server, + client: mcpHTTPClient(server, echo), + nextID: 1, + pending: map[string]chan ssePendingResponse{}, + } +} + +func TestAIMLAPIMCPSSEScrubsTheInjectedKeyFromStreamAnswers(t *testing.T) { + isolateAIMLAPIKeyLookup(t, scrubTestKey) + for _, echoed := range echoedKeys() { + echo := newAIMLAPISSEEcho("/messages", func(method string, id json.RawMessage) string { + if method == "tools/call" { + return rpcReply(id, "result", map[string]any{"isError": true, "content": []any{map[string]any{"type": "text", "text": "rejected " + echoed}}}) + } + return rpcReply(id, "error", map[string]any{"code": -32001, "message": "invalid credential " + echoed}) + }) + client := newEchoSSEClient(echo) + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + if err := client.openStream(ctx); err != nil { + cancel() + t.Fatalf("openStream: %v", err) + } + _, listErr := client.ListTools(ctx) + result, callErr := client.CallTool(ctx, "generate", map[string]any{}) + _ = client.Close() + cancel() + + echo.assertSent(t, "Bearer "+scrubTestKey, 3) + if listErr == nil { + t.Fatal("ListTools succeeded against a server that answered with an error") + } + assertKeyScrubbed(t, listErr.Error()) + if callErr != nil || len(result.Content) != 1 { + t.Fatalf("CallTool = %+v, %v, want the error result", result, callErr) + } + assertKeyScrubbed(t, result.Content[0].Text) + } +} + +func TestAIMLAPIMCPSSEScrubsTheStreamKeyAndThePostKey(t *testing.T) { + // The key is resolved for every request, so the stream and a later POST can + // carry different keys when the provider changes in between. An answer on + // the stream comes from a server that was sent both. + const streamKey = "strm9f8e7d6c5b4a39281706f5e4d3c2b1" + const postKey = "post1a2b3c4d5e6f708192a3b4c5d6e7f8" + keys := []string{streamKey, postKey} + previous := loadAIMLAPIServiceKey + loadAIMLAPIServiceKey = func() (string, error) { + key := keys[0] + if len(keys) > 1 { + keys = keys[1:] + } + return key, nil + } + t.Cleanup(func() { loadAIMLAPIServiceKey = previous }) + + echo := newAIMLAPISSEEcho("/messages", func(_ string, id json.RawMessage) string { + return rpcReply(id, "error", map[string]any{"code": -32001, "message": "stream key " + streamKey + ", post key " + postKey}) + }) + client := newEchoSSEClient(echo) + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + if err := client.openStream(ctx); err != nil { + t.Fatalf("openStream: %v", err) + } + _, err := client.ListTools(ctx) + _ = client.Close() + + echo.mu.Lock() + auth := append([]string(nil), echo.auth...) + echo.mu.Unlock() + if len(auth) != 2 || auth[0] != "Bearer "+streamKey || auth[1] != "Bearer "+postKey { + t.Fatalf("Authorization headers = %q, want the stream key and then the post key", auth) + } + if err == nil { + t.Fatal("ListTools succeeded against a server that answered with an error") + } + assertNoPartOf(t, err.Error(), streamKey) + assertNoPartOf(t, err.Error(), postKey) + if !strings.Contains(err.Error(), "stream key [REDACTED], post key [REDACTED]") { + t.Fatalf("error = %q, want both keys replaced", err) + } +} + +func TestAIMLAPIMCPSSEScrubsTheInjectedKeyFromAnEndpointError(t *testing.T) { + isolateAIMLAPIKeyLookup(t, scrubTestKey) + // An invalid escape makes the endpoint fail to parse, and the parse error + // quotes the endpoint the server sent. + echo := newAIMLAPISSEEcho("/messages/"+scrubTestKey+"%zz", func(string, json.RawMessage) string { return "" }) + client := newEchoSSEClient(echo) + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + err := client.openStream(ctx) + _ = client.Close() + echo.assertSent(t, "Bearer "+scrubTestKey, 1) + if err == nil || !strings.Contains(err.Error(), "endpoint") { + t.Fatalf("openStream error = %v, want the endpoint parse failure", err) + } + assertKeyScrubbed(t, err.Error()) +} diff --git a/internal/mcp/network_client.go b/internal/mcp/network_client.go index deefc164e..2683aced0 100644 --- a/internal/mcp/network_client.go +++ b/internal/mcp/network_client.go @@ -35,6 +35,10 @@ type remoteSSEClient struct { pending map[string]chan ssePendingResponse streamErr error closed bool + // streamKey is the injected AI/ML API key the stream request carried, if + // any. openStream sets it once, before the reader starts; the messages on + // the stream are answers from a server that was sent it. + streamKey string } type ssePendingResponse struct { @@ -259,12 +263,13 @@ func (client *remoteSSEClient) request(ctx context.Context, method string, param client.pending[key] = responseChannel client.mu.Unlock() - if err := client.post(ctx, rpcMessage{ + sentKey, err := client.post(ctx, rpcMessage{ JSONRPC: "2.0", ID: id, Method: method, Params: rawParams, - }); err != nil { + }) + if err != nil { client.removePending(key) return err } @@ -274,14 +279,15 @@ func (client *remoteSSEClient) request(ctx context.Context, method string, param if response.err != nil { return response.err } - if !rpcIDMatches(response.message.ID, id) { + message := scrubAIMLAPIKeys(response.message, client.streamKey, sentKey) + if !rpcIDMatches(message.ID, id) { return fmt.Errorf("MCP %s response id mismatch for server %s", method, client.server.Name) } - if response.message.Error != nil { - return fmt.Errorf("MCP %s failed: %s", method, response.message.Error.Message) + if message.Error != nil { + return fmt.Errorf("MCP %s failed: %s", method, message.Error.Message) } - if target != nil && len(response.message.Result) > 0 { - if err := json.Unmarshal(response.message.Result, target); err != nil { + if target != nil && len(message.Result) > 0 { + if err := json.Unmarshal(message.Result, target); err != nil { return fmt.Errorf("decode MCP %s result: %w", method, err) } } @@ -320,11 +326,12 @@ func (client *remoteSSEClient) notify(ctx context.Context, method string, params if err != nil { return err } - return client.post(ctx, rpcMessage{ + _, err = client.post(ctx, rpcMessage{ JSONRPC: "2.0", Method: method, Params: rawParams, }) + return err } func (client *networkClient) post(ctx context.Context, message rpcMessage, expectResponse bool) (result rpcMessage, err error) { @@ -357,7 +364,11 @@ func (client *networkClient) post(ctx context.Context, message rpcMessage, expec if response.StatusCode == http.StatusNoContent { return rpcMessage{}, fmt.Errorf("MCP %s server %s returned no response body", client.server.Type, client.server.Name) } - return client.decodeResponse(response) + decoded, err := client.decodeResponse(response) + if err != nil { + return rpcMessage{}, err + } + return scrubAIMLAPIKeys(decoded, injectedAIMLAPIKey(client.server, response.Request)), nil } func (client *remoteSSEClient) openStream(ctx context.Context) error { @@ -380,6 +391,7 @@ func (client *remoteSSEClient) openStream(ctx context.Context) error { return err } client.streamBody = response.Body + client.streamKey = injectedAIMLAPIKey(client.server, response.Request) ready := make(chan error, 1) go client.readStream(response.Body, ready) @@ -391,31 +403,34 @@ func (client *remoteSSEClient) openStream(ctx context.Context) error { } } -func (client *remoteSSEClient) post(ctx context.Context, message rpcMessage) (err error) { +// post sends message to the endpoint and returns the injected AI/ML API key the +// request carried, if any, since the answer arrives on the stream from a server +// that was sent it. +func (client *remoteSSEClient) post(ctx context.Context, message rpcMessage) (sentKey string, err error) { endpointURL, err := client.currentEndpoint() if err != nil { - return err + return "", err } body, err := json.Marshal(message) if err != nil { - return err + return "", err } request, err := http.NewRequestWithContext(ctx, http.MethodPost, endpointURL, bytes.NewReader(body)) if err != nil { - return fmt.Errorf("create MCP SSE post for %s: %w", client.server.Name, err) + return "", fmt.Errorf("create MCP SSE post for %s: %w", client.server.Name, err) } client.applyPostHeaders(request) response, err := client.client.Do(request) if err != nil { - return fmt.Errorf("post MCP SSE message to %s: %w", client.server.Name, err) + return "", fmt.Errorf("post MCP SSE message to %s: %w", client.server.Name, err) } defer closeResponseBody(&err, client.server, response.Body) if response.StatusCode < http.StatusOK || response.StatusCode >= http.StatusMultipleChoices { - return httpStatusError(client.server, response) + return "", httpStatusError(client.server, response) } _, _ = io.Copy(io.Discard, response.Body) - return nil + return injectedAIMLAPIKey(client.server, response.Request), nil } func (client *networkClient) applyHeaders(request *http.Request) { @@ -526,7 +541,8 @@ func (client *remoteSSEClient) readStream(reader io.Reader, ready chan<- error) func (client *remoteSSEClient) setEndpoint(value string) error { endpointURL, err := resolveSSEEndpointURL(client.server.URL, strings.TrimSpace(value)) if err != nil { - return err + // A URL parse error quotes the endpoint, which is the server's text. + return scrubAIMLAPIKeyError(err, client.streamKey) } client.mu.Lock() client.endpointURL = endpointURL diff --git a/internal/redaction/redact_secret_values_test.go b/internal/redaction/redact_secret_values_test.go new file mode 100644 index 000000000..eb9d03808 --- /dev/null +++ b/internal/redaction/redact_secret_values_test.go @@ -0,0 +1,40 @@ +package redaction + +import "testing" + +func TestRedactSecretValuesRemovesOnlyTheKnownValues(t *testing.T) { + const key = "aiml-0123456789abcdef0123456789abcdef" + nul := string(rune(0)) + zeroWidthSpace := string(rune(0x200b)) + for _, tc := range []struct { + name string + value string + want string + // shaped marks text RedactString would rewrite by shape alone, so the + // case shows that RedactSecretValues leaves shapes to RedactString. + shaped bool + }{ + {name: "whole", value: "invalid key " + key + " for this account", want: "invalid key " + RedactedSecret + " for this account"}, + {name: "bearer form", value: "got Bearer " + key, want: "got Bearer " + RedactedSecret}, + {name: "split by NUL", value: "invalid key " + key[:9] + nul + key[9:], want: "invalid key " + RedactedSecret}, + {name: "split by a zero-width space", value: "value " + key[:20] + zeroWidthSpace + key[20:] + ".", want: "value " + RedactedSecret + "."}, + {name: "password assignment", value: "password=hunter2", want: "password=hunter2", shaped: true}, + {name: "another bearer token", value: "Authorization: Bearer some-other-token", want: "Authorization: Bearer some-other-token", shaped: true}, + {name: "openai shaped key", value: "use sk-proj-abcdefghijklmnopqrstuvwxyz0123456789", want: "use sk-proj-abcdefghijklmnopqrstuvwxyz0123456789", shaped: true}, + } { + t.Run(tc.name, func(t *testing.T) { + if got := RedactSecretValues(tc.value, key); got != tc.want { + t.Fatalf("RedactSecretValues(%q) = %q, want %q", tc.value, got, tc.want) + } + if tc.shaped && RedactString(tc.value, Options{}) == tc.value { + t.Fatalf("RedactString(%q) left it unchanged, so this case does not show the shape patterns are off", tc.value) + } + }) + } + + for _, secrets := range [][]string{nil, {""}, {" "}} { + if got := RedactSecretValues("token "+key, secrets...); got != "token "+key { + t.Fatalf("RedactSecretValues with secrets %q = %q, want the text unchanged", secrets, got) + } + } +} diff --git a/internal/redaction/redaction.go b/internal/redaction/redaction.go index afe7e71a9..6581ea5f5 100644 --- a/internal/redaction/redaction.go +++ b/internal/redaction/redaction.go @@ -628,6 +628,19 @@ func secretMatchHasDigit(s string) bool { return false } +// RedactSecretValues replaces every occurrence of each of secrets in value with +// RedactedSecret, whole or split by invisible separators, exactly as +// RedactString treats Options.ExtraSecretValues. Unlike RedactString it runs +// none of the shape patterns, so text that does not contain one of the values +// comes back unchanged. That suits content where only known values must go, +// such as tool output, which a shape pattern could otherwise rewrite. +func RedactSecretValues(value string, secrets ...string) string { + if len(secrets) == 0 { + return value + } + return redactExtraSecretValues(value, secrets, RedactedSecret) +} + func RedactValue(value any, options Options) any { return redactReflect(reflect.ValueOf(value), redactionContext{ options: options, From b36246d05c890ed5b125ae5f45f5392f8d96d68a Mon Sep 17 00:00:00 2001 From: Vasanthdev2004 Date: Mon, 9 Nov 2026 11:35:36 +0530 Subject: [PATCH 2/2] fix(mcp): scrub every key a call sent, and transport errors too Review found two gaps. The key is resolved again for each redirect hop, but the scrub used only the final hop's key, so a same-origin server answering the last hop could echo an earlier hop's key. And a legacy SSE POST that failed at the transport quoted the endpoint URL the server chose, which can hold the stream's key, with no scrub at all. A collector in the request context now records every key applyAIMLAPIRequestHeaders injects during a call. net/http gives each redirect hop the first request's context, so a later hop's key is recorded too. Both transports scrub the response and every error the call returns with those keys, the stream request's own errors included, and the SSE client adds the keys its stream request sent. A scrubbed error unwraps only to the context sentinel the original matched, so a cancellation still reads as one without the chain that held the key. Reported by @jatmn. --- internal/mcp/aimlapi.go | 96 ++++++++--- internal/mcp/aimlapi_scrub_test.go | 266 +++++++++++++++++++++++++++-- internal/mcp/network_client.go | 55 +++--- 3 files changed, 363 insertions(+), 54 deletions(-) diff --git a/internal/mcp/aimlapi.go b/internal/mcp/aimlapi.go index 98b36bf06..8dfbc6f6f 100644 --- a/internal/mcp/aimlapi.go +++ b/internal/mcp/aimlapi.go @@ -2,12 +2,14 @@ package mcp import ( "bytes" + "context" "encoding/json" "errors" "fmt" "net/http" "net/url" "strings" + "sync" "github.com/Twigpine/zero/internal/aimlapi" "github.com/Twigpine/zero/internal/config" @@ -24,12 +26,16 @@ import ( // config.LoadAIMLAPIServiceKey, which in the same call decides whether it may // be sent (no proxy profile, a single token), so a provider change applies to // the next request and a withheld or invalid key means no Authorization -// header. It lives only in that request's headers: never in Server, its -// Identity (OAuth token binding), config.json, list output, errors or tool -// output. An error status to such a request is reported without its body, and -// any other response to it has the key scrubbed from its error text, method -// and result before anything reads them (scrubAIMLAPIKeys), since that text is -// the server's own. A user-configured Authorization header, OAuth, or a +// header. It lives only in that request's headers, and for the length of the +// call in the collector its context carries (aimlapiSentKeys), so the reply can +// be scrubbed: never in Server, its Identity (OAuth token binding), config.json, +// list output, errors or tool output. Every key the call sent, redirect hops +// included, and on the SSE transport the keys the stream request sent, is +// scrubbed from the response's error text, method and result before anything +// reads them (scrubAIMLAPIKeys) and from any error the call returns, a +// transport error quoting a URL the server chose among them +// (scrubAIMLAPIKeyError). An error status to such a request is also reported +// without its body. A user-configured Authorization header, OAuth, or a // project-touched server always goes without it. var loadAIMLAPIServiceKey = config.LoadAIMLAPIServiceKey @@ -62,6 +68,9 @@ func applyAIMLAPIRequestHeaders(server Server, request *http.Request) { return } request.Header.Set("Authorization", "Bearer "+key) + if sent, ok := request.Context().Value(aimlapiSentKeysContextKey{}).(*aimlapiSentKeys); ok { + sent.add(key) + } } // reapplyAIMLAPIRedirectHeaders gives a redirect hop the automatic headers its @@ -107,19 +116,40 @@ func hasUserAuthorizationHeader(server Server) bool { return false } -// injectedAIMLAPIKey returns the key applyAIMLAPIRequestHeaders put on request, -// or "" when request went out without it. Pass the request a response answers -// (response.Request, the last hop of any redirect): that is what the responding -// server was sent. -func injectedAIMLAPIKey(server Server, request *http.Request) string { - if !sentAIMLAPIServiceKey(server, request) { - return "" - } - value := strings.TrimSpace(request.Header.Get("Authorization")) - if scheme, key, ok := strings.Cut(value, " "); ok && strings.EqualFold(scheme, "Bearer") { - value = strings.TrimSpace(key) +// aimlapiSentKeys collects every key applyAIMLAPIRequestHeaders puts on the +// requests of one call. The first request's context carries it, and net/http +// gives each redirect hop that same context, so a key resolved again for a later +// hop is collected too. The final hop's key alone is not enough: a same-origin +// server answering the last hop has seen every earlier one. +type aimlapiSentKeys struct { + mu sync.Mutex + keys []string +} + +type aimlapiSentKeysContextKey struct{} + +// withAIMLAPISentKeys returns ctx carrying a fresh collector for one call. +func withAIMLAPISentKeys(ctx context.Context) (context.Context, *aimlapiSentKeys) { + sent := &aimlapiSentKeys{} + return context.WithValue(ctx, aimlapiSentKeysContextKey{}, sent), sent +} + +func (sent *aimlapiSentKeys) add(key string) { + sent.mu.Lock() + defer sent.mu.Unlock() + for _, existing := range sent.keys { + if existing == key { + return + } } - return value + sent.keys = append(sent.keys, key) +} + +// list returns a copy of the keys collected so far. +func (sent *aimlapiSentKeys) list() []string { + sent.mu.Lock() + defer sent.mu.Unlock() + return append([]string(nil), sent.keys...) } // scrubAIMLAPIKeys removes the injected keys from a response before any of it @@ -205,20 +235,40 @@ func scrubAIMLAPIKeysFromValue(value any, keys []string) (any, bool) { } // scrubAIMLAPIKeyError returns err with the injected keys removed from its text. -// An error the keys do not appear in keeps its chain; one they do appear in is -// replaced by its scrubbed text, because its chain still holds the key. +// An error the keys do not appear in keeps its chain. One they do appear in is +// replaced by its scrubbed text, because its chain still holds the key (a +// *url.Error keeps the URL the server chose); the replacement unwraps only to +// the context sentinel the original matched, so a cancellation still reads as +// one. func scrubAIMLAPIKeyError(err error, keys ...string) error { - if err == nil { - return nil + if err == nil || len(keys) == 0 { + return err } text := err.Error() scrubbed := redaction.RedactSecretValues(text, keys...) if scrubbed == text { return err } - return errors.New(scrubbed) + replacement := &aimlapiScrubbedError{text: scrubbed} + for _, sentinel := range []error{context.Canceled, context.DeadlineExceeded} { + if errors.Is(err, sentinel) { + replacement.cause = sentinel + break + } + } + return replacement } +// aimlapiScrubbedError is an error whose text had an injected key removed. +type aimlapiScrubbedError struct { + text string + cause error +} + +func (scrubbed *aimlapiScrubbedError) Error() string { return scrubbed.text } + +func (scrubbed *aimlapiScrubbedError) Unwrap() error { return scrubbed.cause } + // aimlapiUnauthorizedHint explains how to authenticate when the AI/ML API MCP // server rejects a request without OAuth. It names no key material. func aimlapiUnauthorizedHint(server Server, status int) string { diff --git a/internal/mcp/aimlapi_scrub_test.go b/internal/mcp/aimlapi_scrub_test.go index 1d1fe9014..b36705c87 100644 --- a/internal/mcp/aimlapi_scrub_test.go +++ b/internal/mcp/aimlapi_scrub_test.go @@ -3,8 +3,11 @@ package mcp import ( "context" "encoding/json" + "errors" + "fmt" "io" "net/http" + "net/url" "strings" "sync" "testing" @@ -243,6 +246,11 @@ func TestAIMLAPIMCPLeavesResponsesAloneWithoutTheInjectedKey(t *testing.T) { type aimlapiSSEEcho struct { endpoint string reply func(method string, id json.RawMessage) string + // postErr, when set, fails every POST at the transport. + postErr error + // rawEvent, when set, writes its whole event to the stream instead of a + // message event built from reply. + rawEvent func(method string, id json.RawMessage) string events chan string mu sync.Mutex auth []string @@ -279,6 +287,9 @@ func (echo *aimlapiSSEEcho) RoundTrip(request *http.Request) (*http.Response, er response.Body = reader return response, nil } + if echo.postErr != nil { + return nil, echo.postErr + } var body []byte if request.Body != nil { body, _ = io.ReadAll(request.Body) @@ -290,11 +301,33 @@ func (echo *aimlapiSSEEcho) RoundTrip(request *http.Request) (*http.Response, er } _ = json.Unmarshal(body, &message) if len(message.ID) > 0 { - echo.events <- "event: message\ndata: " + echo.reply(message.Method, message.ID) + "\n\n" + if echo.rawEvent != nil { + echo.events <- echo.rawEvent(message.Method, message.ID) + } else { + echo.events <- "event: message\ndata: " + echo.reply(message.Method, message.ID) + "\n\n" + } } return fakeHTTPResponse(request, http.StatusAccepted, ""), nil } +// rotateAIMLAPIKeys makes the key resolver hand out keys in order, repeating the +// last, the way a provider change between two requests of one call would. +func rotateAIMLAPIKeys(t *testing.T, keys ...string) { + t.Helper() + var mu sync.Mutex + previous := loadAIMLAPIServiceKey + loadAIMLAPIServiceKey = func() (string, error) { + mu.Lock() + defer mu.Unlock() + key := keys[0] + if len(keys) > 1 { + keys = keys[1:] + } + return key, nil + } + t.Cleanup(func() { loadAIMLAPIServiceKey = previous }) +} + func (echo *aimlapiSSEEcho) assertSent(t *testing.T, want string, requests int) { t.Helper() echo.mu.Lock() @@ -357,16 +390,7 @@ func TestAIMLAPIMCPSSEScrubsTheStreamKeyAndThePostKey(t *testing.T) { // the stream comes from a server that was sent both. const streamKey = "strm9f8e7d6c5b4a39281706f5e4d3c2b1" const postKey = "post1a2b3c4d5e6f708192a3b4c5d6e7f8" - keys := []string{streamKey, postKey} - previous := loadAIMLAPIServiceKey - loadAIMLAPIServiceKey = func() (string, error) { - key := keys[0] - if len(keys) > 1 { - keys = keys[1:] - } - return key, nil - } - t.Cleanup(func() { loadAIMLAPIServiceKey = previous }) + rotateAIMLAPIKeys(t, streamKey, postKey) echo := newAIMLAPISSEEcho("/messages", func(_ string, id json.RawMessage) string { return rpcReply(id, "error", map[string]any{"code": -32001, "message": "stream key " + streamKey + ", post key " + postKey}) @@ -412,3 +436,223 @@ func TestAIMLAPIMCPSSEScrubsTheInjectedKeyFromAnEndpointError(t *testing.T) { } assertKeyScrubbed(t, err.Error()) } + +// aimlapiHopEcho answers the first request with a 307 to location and hands the +// second hop to second, recording the Authorization each hop carried. +type aimlapiHopEcho struct { + location string + second func(request *http.Request, id json.RawMessage) (*http.Response, error) + mu sync.Mutex + auth []string +} + +func (echo *aimlapiHopEcho) RoundTrip(request *http.Request) (*http.Response, error) { + var body []byte + if request.Body != nil { + body, _ = io.ReadAll(request.Body) + _ = request.Body.Close() + } + echo.mu.Lock() + echo.auth = append(echo.auth, request.Header.Get("Authorization")) + hop := len(echo.auth) + echo.mu.Unlock() + if hop == 1 { + response := fakeHTTPResponse(request, http.StatusTemporaryRedirect, "") + response.Header.Set("Location", echo.location) + return response, nil + } + var message struct { + ID json.RawMessage `json:"id"` + } + _ = json.Unmarshal(body, &message) + return echo.second(request, message.ID) +} + +func TestAIMLAPIMCPScrubsEveryKeySentAcrossRedirectHops(t *testing.T) { + // The key is resolved again for a redirect hop, so the hops of one call can + // carry different keys, and the same-origin server answering the last hop + // has seen them all. + const first = "aimlfirst0123456789abcdef01234567" + const second = "aimlsecond0123456789abcdef0123456" + for _, tc := range []struct { + name string + location string + // hopKey is the key the redirect hop must carry, "" for none. + hopKey string + second func(*http.Request, json.RawMessage) (*http.Response, error) + }{ + { + name: "error echoes the first hop's key", + location: "https://mcp.aimlapi.com/mcp/", + hopKey: second, + second: func(request *http.Request, id json.RawMessage) (*http.Response, error) { + return fakeHTTPResponse(request, http.StatusOK, rpcReply(id, "error", map[string]any{"code": -32001, "message": "stale " + first})), nil + }, + }, + { + name: "status body after a hop that dropped the key", + location: "https://user@mcp.aimlapi.com/mcp", + second: func(request *http.Request, _ json.RawMessage) (*http.Response, error) { + return fakeHTTPResponse(request, http.StatusPaymentRequired, `{"error":"invalid credential `+first+`"}`), nil + }, + }, + { + name: "transport error quoting the redirect target", + location: "https://mcp.aimlapi.com/mcp?session=" + first, + hopKey: second, + second: func(*http.Request, json.RawMessage) (*http.Response, error) { + return nil, errors.New("connection reset by peer") + }, + }, + } { + t.Run(tc.name, func(t *testing.T) { + rotateAIMLAPIKeys(t, first, second) + echo := &aimlapiHopEcho{location: tc.location, second: tc.second} + server := aimlapiTestServer("aimlapi", aimlapi.MCPServerURL) + client := &networkClient{server: server, client: mcpHTTPClient(server, echo), nextID: 1} + err := client.initialize(context.Background()) + + echo.mu.Lock() + auth := append([]string(nil), echo.auth...) + echo.mu.Unlock() + if len(auth) != 2 || auth[0] != "Bearer "+first { + t.Fatalf("Authorization per hop = %q, want the first key on the first request and a redirect hop", auth) + } + if tc.hopKey != "" && auth[1] != "Bearer "+tc.hopKey { + t.Fatalf("redirect hop Authorization = %q, want the key resolved again for it", auth[1]) + } + // With no key, net/http fills Basic auth from the hop's own userinfo. + if tc.hopKey == "" && (strings.Contains(auth[1], first) || strings.Contains(auth[1], second)) { + t.Fatalf("redirect hop Authorization = %q, want no automatic key", auth[1]) + } + if err == nil { + t.Fatal("initialize succeeded against a failing server") + } + assertNoPartOf(t, err.Error(), first) + assertNoPartOf(t, err.Error(), second) + if !strings.Contains(err.Error(), "[REDACTED]") { + t.Fatalf("error = %q, want the echoed key replaced", err) + } + }) + } +} + +func TestAIMLAPIMCPSSEScrubsTheStreamKeyFromAFailedPost(t *testing.T) { + // The server chose an endpoint holding the key its stream request carried. + // A POST to it that fails at the transport quotes that URL, and the POST + // itself carried a different key, so only the stream's keys can scrub it. + const streamKey = "strm0a1b2c3d4e5f60718293a4b5c6d7e8" + const postKey = "post9e8d7c6b5a4f30291807f6e5d4c3b2" + rotateAIMLAPIKeys(t, streamKey, postKey) + echo := newAIMLAPISSEEcho("/messages/"+streamKey, func(string, json.RawMessage) string { return "" }) + echo.postErr = errors.New("connection reset by peer") + client := newEchoSSEClient(echo) + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + if err := client.openStream(ctx); err != nil { + t.Fatalf("openStream: %v", err) + } + _, err := client.ListTools(ctx) + _ = client.Close() + + echo.mu.Lock() + auth := append([]string(nil), echo.auth...) + echo.mu.Unlock() + if len(auth) != 2 || auth[0] != "Bearer "+streamKey || auth[1] != "Bearer "+postKey { + t.Fatalf("Authorization headers = %q, want the stream key and then the post key", auth) + } + if err == nil || !strings.Contains(err.Error(), "post MCP SSE message") { + t.Fatalf("ListTools error = %v, want the failed POST", err) + } + assertNoPartOf(t, err.Error(), streamKey) + assertNoPartOf(t, err.Error(), postKey) + if !strings.Contains(err.Error(), "/messages/[REDACTED]") { + t.Fatalf("error = %q, want the endpoint quoted with the key replaced", err) + } +} + +func TestAIMLAPIMCPSSEScrubsTheKeyFromALaterEndpointError(t *testing.T) { + // An endpoint event can arrive after the stream is open. Its parse error + // then fails the pending requests directly rather than returning from + // openStream, so it has to be scrubbed where it is made. + isolateAIMLAPIKeyLookup(t, scrubTestKey) + echo := newAIMLAPISSEEcho("/messages", nil) + echo.rawEvent = func(string, json.RawMessage) string { + return "event: endpoint\ndata: /messages/" + scrubTestKey + "%zz\n\n" + } + client := newEchoSSEClient(echo) + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + if err := client.openStream(ctx); err != nil { + t.Fatalf("openStream: %v", err) + } + _, err := client.ListTools(ctx) + _ = client.Close() + echo.assertSent(t, "Bearer "+scrubTestKey, 2) + if err == nil || !strings.Contains(err.Error(), "endpoint") { + t.Fatalf("ListTools error = %v, want the later endpoint's parse failure", err) + } + assertKeyScrubbed(t, err.Error()) +} + +func TestAIMLAPIMCPSSEScrubsKeysFromAFailedStreamRedirect(t *testing.T) { + // The stream request follows redirects like any other, so a hop that fails + // at the transport quotes the target the server chose, which can hold the + // key the first hop carried. + const first = "strmfirst0123456789abcdef01234567" + const second = "strmsecond0123456789abcdef0123456" + rotateAIMLAPIKeys(t, first, second) + echo := &aimlapiHopEcho{ + location: "https://mcp.aimlapi.com/mcp?session=" + first, + second: func(*http.Request, json.RawMessage) (*http.Response, error) { + return nil, errors.New("connection reset by peer") + }, + } + server := Server{Name: "aimlapi", Type: ServerTypeSSE, URL: aimlapi.MCPServerURL} + client := &remoteSSEClient{server: server, client: mcpHTTPClient(server, echo), nextID: 1, pending: map[string]chan ssePendingResponse{}} + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + err := client.openStream(ctx) + _ = client.Close() + + echo.mu.Lock() + auth := append([]string(nil), echo.auth...) + echo.mu.Unlock() + if len(auth) != 2 || auth[0] != "Bearer "+first || auth[1] != "Bearer "+second { + t.Fatalf("Authorization per hop = %q, want the first key and then the second", auth) + } + if err == nil || !strings.Contains(err.Error(), "open MCP SSE stream") { + t.Fatalf("openStream error = %v, want the failed stream request", err) + } + assertNoPartOf(t, err.Error(), first) + assertNoPartOf(t, err.Error(), second) + if !strings.Contains(err.Error(), "session=[REDACTED]") { + t.Fatalf("error = %q, want the redirect target quoted with the key replaced", err) + } +} + +func TestScrubAIMLAPIKeyErrorKeepsOnlyTheContextSentinel(t *testing.T) { + for _, sentinel := range []error{context.Canceled, context.DeadlineExceeded} { + original := fmt.Errorf("post MCP SSE message to aimlapi: %w", &url.Error{Op: "Post", URL: "https://mcp.aimlapi.com/messages/" + scrubTestKey, Err: sentinel}) + scrubbed := scrubAIMLAPIKeyError(original, scrubTestKey) + assertKeyScrubbed(t, scrubbed.Error()) + if !errors.Is(scrubbed, sentinel) { + t.Fatalf("scrubbed error %q no longer reads as %v", scrubbed, sentinel) + } + if unwrapped := errors.Unwrap(scrubbed); unwrapped != sentinel { + t.Fatalf("scrubbed error unwraps to %v, want only %v", unwrapped, sentinel) + } + var urlErr *url.Error + if errors.As(scrubbed, &urlErr) { + t.Fatalf("scrubbed error still exposes the *url.Error holding the key: %v", urlErr.URL) + } + } + + if scrubbed := scrubAIMLAPIKeyError(fmt.Errorf("decode: %s", scrubTestKey), scrubTestKey); errors.Unwrap(scrubbed) != nil { + t.Fatalf("scrubbed error without a context sentinel unwraps to %v, want nothing", errors.Unwrap(scrubbed)) + } + untouched := fmt.Errorf("call MCP: %w", context.Canceled) + if got := scrubAIMLAPIKeyError(untouched, scrubTestKey); got != untouched { + t.Fatalf("an error without the key was replaced: %v", got) + } +} diff --git a/internal/mcp/network_client.go b/internal/mcp/network_client.go index 2683aced0..f67b3b3bb 100644 --- a/internal/mcp/network_client.go +++ b/internal/mcp/network_client.go @@ -35,10 +35,11 @@ type remoteSSEClient struct { pending map[string]chan ssePendingResponse streamErr error closed bool - // streamKey is the injected AI/ML API key the stream request carried, if - // any. openStream sets it once, before the reader starts; the messages on - // the stream are answers from a server that was sent it. - streamKey string + // streamKeys are the injected AI/ML API keys the stream request carried, + // redirect hops included. openStream sets them once, before the reader + // starts; the endpoint and the messages on the stream come from a server + // that was sent them. + streamKeys []string } type ssePendingResponse struct { @@ -263,7 +264,7 @@ func (client *remoteSSEClient) request(ctx context.Context, method string, param client.pending[key] = responseChannel client.mu.Unlock() - sentKey, err := client.post(ctx, rpcMessage{ + sentKeys, err := client.post(ctx, rpcMessage{ JSONRPC: "2.0", ID: id, Method: method, @@ -279,7 +280,7 @@ func (client *remoteSSEClient) request(ctx context.Context, method string, param if response.err != nil { return response.err } - message := scrubAIMLAPIKeys(response.message, client.streamKey, sentKey) + message := scrubAIMLAPIKeys(response.message, append(append([]string(nil), client.streamKeys...), sentKeys...)...) if !rpcIDMatches(message.ID, id) { return fmt.Errorf("MCP %s response id mismatch for server %s", method, client.server.Name) } @@ -339,6 +340,11 @@ func (client *networkClient) post(ctx context.Context, message rpcMessage, expec if err != nil { return rpcMessage{}, err } + ctx, sent := withAIMLAPISentKeys(ctx) + // From here an error can quote the server or a URL it chose, a redirect + // target among them, so every key the call sent, on any hop, is scrubbed + // from it. This runs last, after the body is closed. + defer func() { err = scrubAIMLAPIKeyError(err, sent.list()...) }() request, err := http.NewRequestWithContext(ctx, http.MethodPost, client.server.URL, bytes.NewReader(body)) if err != nil { return rpcMessage{}, fmt.Errorf("create MCP %s request for %s: %w", client.server.Type, client.server.Name, err) @@ -368,12 +374,14 @@ func (client *networkClient) post(ctx context.Context, message rpcMessage, expec if err != nil { return rpcMessage{}, err } - return scrubAIMLAPIKeys(decoded, injectedAIMLAPIKey(client.server, response.Request)), nil + return scrubAIMLAPIKeys(decoded, sent.list()...), nil } -func (client *remoteSSEClient) openStream(ctx context.Context) error { +func (client *remoteSSEClient) openStream(ctx context.Context) (err error) { streamCtx, cancel := context.WithCancel(ctx) client.streamCancel = cancel + streamCtx, sent := withAIMLAPISentKeys(streamCtx) + defer func() { err = scrubAIMLAPIKeyError(err, sent.list()...) }() request, err := http.NewRequestWithContext(streamCtx, http.MethodGet, client.server.URL, nil) if err != nil { @@ -391,7 +399,7 @@ func (client *remoteSSEClient) openStream(ctx context.Context) error { return err } client.streamBody = response.Body - client.streamKey = injectedAIMLAPIKey(client.server, response.Request) + client.streamKeys = sent.list() ready := make(chan error, 1) go client.readStream(response.Body, ready) @@ -403,34 +411,41 @@ func (client *remoteSSEClient) openStream(ctx context.Context) error { } } -// post sends message to the endpoint and returns the injected AI/ML API key the -// request carried, if any, since the answer arrives on the stream from a server -// that was sent it. -func (client *remoteSSEClient) post(ctx context.Context, message rpcMessage) (sentKey string, err error) { +// post sends message to the endpoint and returns the injected AI/ML API keys the +// request carried, redirect hops included, since the answer arrives on the +// stream from a server that was sent them. The endpoint is the server's own +// choice and can hold one of the stream's keys, and a transport error quotes it, +// so the stream's keys are scrubbed from any error as well as the POST's. +func (client *remoteSSEClient) post(ctx context.Context, message rpcMessage) (sentKeys []string, err error) { + ctx, sent := withAIMLAPISentKeys(ctx) + defer func() { + sentKeys = sent.list() + err = scrubAIMLAPIKeyError(err, append(append([]string(nil), client.streamKeys...), sentKeys...)...) + }() endpointURL, err := client.currentEndpoint() if err != nil { - return "", err + return nil, err } body, err := json.Marshal(message) if err != nil { - return "", err + return nil, err } request, err := http.NewRequestWithContext(ctx, http.MethodPost, endpointURL, bytes.NewReader(body)) if err != nil { - return "", fmt.Errorf("create MCP SSE post for %s: %w", client.server.Name, err) + return nil, fmt.Errorf("create MCP SSE post for %s: %w", client.server.Name, err) } client.applyPostHeaders(request) response, err := client.client.Do(request) if err != nil { - return "", fmt.Errorf("post MCP SSE message to %s: %w", client.server.Name, err) + return nil, fmt.Errorf("post MCP SSE message to %s: %w", client.server.Name, err) } defer closeResponseBody(&err, client.server, response.Body) if response.StatusCode < http.StatusOK || response.StatusCode >= http.StatusMultipleChoices { - return "", httpStatusError(client.server, response) + return nil, httpStatusError(client.server, response) } _, _ = io.Copy(io.Discard, response.Body) - return injectedAIMLAPIKey(client.server, response.Request), nil + return nil, nil } func (client *networkClient) applyHeaders(request *http.Request) { @@ -542,7 +557,7 @@ func (client *remoteSSEClient) setEndpoint(value string) error { endpointURL, err := resolveSSEEndpointURL(client.server.URL, strings.TrimSpace(value)) if err != nil { // A URL parse error quotes the endpoint, which is the server's text. - return scrubAIMLAPIKeyError(err, client.streamKey) + return scrubAIMLAPIKeyError(err, client.streamKeys...) } client.mu.Lock() client.endpointURL = endpointURL