From 6e94710384ee2be0cf8d3e47f073d41e344702ce Mon Sep 17 00:00:00 2001 From: Vasanthdev2004 Date: Sat, 7 Nov 2026 14:47:58 +0530 Subject: [PATCH] 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,