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,