diff --git a/docs/EXTENDING.md b/docs/EXTENDING.md index 9b9f099d4..1ce8e8908 100644 --- a/docs/EXTENDING.md +++ b/docs/EXTENDING.md @@ -307,6 +307,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..8dfbc6f6f 100644 --- a/internal/mcp/aimlapi.go +++ b/internal/mcp/aimlapi.go @@ -1,13 +1,19 @@ 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" + "github.com/Twigpine/zero/internal/redaction" ) // The AI/ML API MCP server authenticates with the user's AI/ML API key. @@ -20,11 +26,17 @@ 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 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. +// 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 var aimlapiMCPHost = func() string { @@ -56,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 @@ -101,6 +116,159 @@ func hasUserAuthorizationHeader(server Server) bool { return false } +// 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 + } + } + 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 +// 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 (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 || len(keys) == 0 { + return err + } + text := err.Error() + scrubbed := redaction.RedactSecretValues(text, keys...) + if scrubbed == text { + return err + } + 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 new file mode 100644 index 000000000..b36705c87 --- /dev/null +++ b/internal/mcp/aimlapi_scrub_test.go @@ -0,0 +1,658 @@ +package mcp + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "net/url" + "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 + // 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 +} + +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 + } + if echo.postErr != nil { + return nil, echo.postErr + } + 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 { + 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() + 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" + 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}) + }) + 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()) +} + +// 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 deefc164e..f67b3b3bb 100644 --- a/internal/mcp/network_client.go +++ b/internal/mcp/network_client.go @@ -35,6 +35,11 @@ type remoteSSEClient struct { pending map[string]chan ssePendingResponse streamErr error closed bool + // 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 { @@ -259,12 +264,13 @@ func (client *remoteSSEClient) request(ctx context.Context, method string, param client.pending[key] = responseChannel client.mu.Unlock() - if err := client.post(ctx, rpcMessage{ + sentKeys, 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 +280,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, 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) } - 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 +327,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) { @@ -332,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) @@ -357,12 +370,18 @@ 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, 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 { @@ -380,6 +399,7 @@ func (client *remoteSSEClient) openStream(ctx context.Context) error { return err } client.streamBody = response.Body + client.streamKeys = sent.list() ready := make(chan error, 1) go client.readStream(response.Body, ready) @@ -391,31 +411,41 @@ 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 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 nil + return nil, nil } func (client *networkClient) applyHeaders(request *http.Request) { @@ -526,7 +556,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.streamKeys...) } 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,