From 5b508cd5599779649594981e9176dbe31820adc0 Mon Sep 17 00:00:00 2001 From: breken-ai <312387581+breken-ai@users.noreply.github.com> Date: Thu, 8 Oct 2026 10:04:23 -0700 Subject: [PATCH] fix(audit): detect JSON-RPC errors in responses larger than the detection buffer The audit middleware parsed a 512-byte prefix of the response with json.Unmarshal. A cut-off JSON object does not unmarshal, so any error body over 512 bytes was recorded as outcome=success with no error code. When the prefix fills the buffer and the normal parse finds no error, walk the top-level object token by token and stop at the "error" key. Only the bytes before the cut matter, nested "error" keys are ignored, and a truncated "result" is still a success. The error code is kept when it was read before the cut; the message is kept only when it was read completely. Fixes #6720 Signed-off-by: breken-ai <312387581+breken-ai@users.noreply.github.com> --- pkg/audit/auditor.go | 8 +++- pkg/audit/auditor_test.go | 57 ++++++++++++++++++++++++++ pkg/mcp/response.go | 70 +++++++++++++++++++++++++++++++ pkg/mcp/response_test.go | 86 +++++++++++++++++++++++++++++++++++++++ 4 files changed, 220 insertions(+), 1 deletion(-) diff --git a/pkg/audit/auditor.go b/pkg/audit/auditor.go index 458ebfccb0..03827195fc 100644 --- a/pkg/audit/auditor.go +++ b/pkg/audit/auditor.go @@ -464,7 +464,13 @@ func (*Auditor) detectApplicationError(rw *responseWriter) *mcp.ParsedMCPRespons prefix = rw.errorDetectionBody.Bytes() } if len(prefix) > 0 && prefix[0] == '{' { - return mcp.ParseMCPResponse(prefix) + parsed := mcp.ParseMCPResponse(prefix) + // A full buffer may be a cut-off response, which does not unmarshal. + // Fall back to a token walk that only needs the bytes before the cut. + if !parsed.HasError && len(prefix) >= errorDetectionBufferSize { + return mcp.ParseTruncatedMCPResponse(prefix) + } + return parsed } return nil } diff --git a/pkg/audit/auditor_test.go b/pkg/audit/auditor_test.go index 04e587df57..2a3501ec4d 100644 --- a/pkg/audit/auditor_test.go +++ b/pkg/audit/auditor_test.go @@ -1253,6 +1253,63 @@ func TestMiddlewareDetectsJSONRPCErrors(t *testing.T) { assert.Contains(t, logOutput, "jsonrpc_error_code") }) + t.Run("detects JSON-RPC error in a body larger than the detection buffer", func(t *testing.T) { + t.Parallel() + // Bodies of 511, 512 and 513+ bytes: the old prefix parse missed + // everything past the buffer size. + for _, size := range []int{errorDetectionBufferSize - 1, errorDetectionBufferSize, errorDetectionBufferSize + 1, 4 * errorDetectionBufferSize} { + var logBuf bytes.Buffer + detectErrors := true + config := &Config{DetectApplicationErrors: &detectErrors} + auditor, err := NewAuditorWithTransport(config, "streamable-http") + require.NoError(t, err) + auditor.auditLogger = NewAuditLogger(&logBuf) + + const envelope = `{"jsonrpc":"2.0","id":1,"error":{"code":-32000,"message":""}}` + errorResponse := `{"jsonrpc":"2.0","id":1,"error":{"code":-32000,"message":"` + + strings.Repeat("a", size-len(envelope)) + `"}}` + require.Len(t, errorResponse, size) + handler := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusOK) + _, err := w.Write([]byte(errorResponse)) + require.NoError(t, err) + }) + + req := httptest.NewRequest("POST", "/mcp", strings.NewReader(`{"jsonrpc":"2.0","id":"1","method":"tools/call","params":{"name":"test"}}`)) + req.Header.Set("Content-Type", "application/json") + rr := httptest.NewRecorder() + auditor.Middleware(handler).ServeHTTP(rr, req) + + assert.Equal(t, errorResponse, rr.Body.String(), "size %d", size) + assert.Contains(t, logBuf.String(), OutcomeApplicationError, "size %d", size) + assert.Contains(t, logBuf.String(), "-32000", "size %d", size) + } + }) + + t.Run("keeps outcome=success for a large JSON-RPC result", func(t *testing.T) { + t.Parallel() + var logBuf bytes.Buffer + detectErrors := true + config := &Config{DetectApplicationErrors: &detectErrors} + auditor, err := NewAuditorWithTransport(config, "streamable-http") + require.NoError(t, err) + auditor.auditLogger = NewAuditLogger(&logBuf) + + response := `{"jsonrpc":"2.0","id":1,"result":{"error":{"code":1,"message":"x"},"text":"` + + strings.Repeat("a", 4*errorDetectionBufferSize) + `"}}` + handler := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusOK) + _, err := w.Write([]byte(response)) + require.NoError(t, err) + }) + + req := httptest.NewRequest("POST", "/mcp", strings.NewReader(`{"jsonrpc":"2.0","id":"1","method":"tools/call","params":{"name":"test"}}`)) + req.Header.Set("Content-Type", "application/json") + auditor.Middleware(handler).ServeHTTP(httptest.NewRecorder(), req) + + assert.NotContains(t, logBuf.String(), OutcomeApplicationError) + }) + t.Run("keeps outcome=success for valid JSON-RPC result", func(t *testing.T) { t.Parallel() var logBuf bytes.Buffer diff --git a/pkg/mcp/response.go b/pkg/mcp/response.go index 9a50b75161..36ede425dd 100644 --- a/pkg/mcp/response.go +++ b/pkg/mcp/response.go @@ -4,6 +4,7 @@ package mcp import ( + "bytes" "encoding/json" ) @@ -60,3 +61,72 @@ func ParseMCPResponse(body []byte) *ParsedMCPResponse { ErrorMessage: envelope.Error.Message, } } + +// ParseTruncatedMCPResponse detects a top-level JSON-RPC "error" object in a +// response body that may have been cut off, such as a fixed-size prefix of a +// larger response. ParseMCPResponse cannot do this because a truncated JSON +// document does not unmarshal. +// +// It walks the top-level object token by token and stops at the "error" key, +// so only the bytes before the cut matter. "code" and "message" are filled in +// when they were read completely; otherwise they keep their zero values and +// HasError is still true. Values of other top-level keys are skipped, and +// nested "error" keys are ignored. Like ParseMCPResponse it is lenient and +// returns HasError=false when the body is not a JSON object or has no +// top-level "error" object before the cut. +func ParseTruncatedMCPResponse(body []byte) *ParsedMCPResponse { + dec := json.NewDecoder(bytes.NewReader(body)) + if tok, err := dec.Token(); err != nil || tok != json.Delim('{') { + return &ParsedMCPResponse{} + } + for dec.More() { + keyTok, err := dec.Token() + if err != nil { + return &ParsedMCPResponse{} + } + key, ok := keyTok.(string) + if !ok { + return &ParsedMCPResponse{} + } + if key != "error" { + var skip json.RawMessage + if err := dec.Decode(&skip); err != nil { + return &ParsedMCPResponse{} + } + continue + } + if tok, err := dec.Token(); err != nil || tok != json.Delim('{') { + // "error": null or a non-object value is not a JSON-RPC error. + return &ParsedMCPResponse{} + } + res := &ParsedMCPResponse{HasError: true} + for dec.More() { + fieldTok, err := dec.Token() + if err != nil { + return res + } + field, _ := fieldTok.(string) + switch field { + case "code": + var code int + if err := dec.Decode(&code); err != nil { + return res + } + res.ErrorCode = code + case "message": + var msg string + if err := dec.Decode(&msg); err != nil { + return res + } + res.ErrorMessage = msg + default: + var skip json.RawMessage + if err := dec.Decode(&skip); err != nil { + return res + } + } + } + return res + } + return &ParsedMCPResponse{} +} diff --git a/pkg/mcp/response_test.go b/pkg/mcp/response_test.go index aa963b9ed4..5538661f74 100644 --- a/pkg/mcp/response_test.go +++ b/pkg/mcp/response_test.go @@ -90,3 +90,89 @@ func TestParseMCPResponse(t *testing.T) { }) } } + +func TestParseTruncatedMCPResponse(t *testing.T) { + t.Parallel() + + full := `{"jsonrpc":"2.0","id":1,"error":{"code":-32000,"message":"` + strings.Repeat("a", 600) + `"}}` + + tests := []struct { + name string + body string + wantHasError bool + wantCode int + wantMessage string + }{ + { + name: "complete error", + body: `{"jsonrpc":"2.0","id":1,"error":{"code":-32603,"message":"boom"}}`, + wantHasError: true, + wantCode: -32603, + wantMessage: "boom", + }, + { + name: "cut inside message string keeps code", + body: full[:200], + wantHasError: true, + wantCode: -32000, + }, + { + name: "cut right after code", + body: `{"jsonrpc":"2.0","id":1,"error":{"code":-32000,`, + wantHasError: true, + wantCode: -32000, + }, + { + name: "cut inside the error key name", + body: `{"jsonrpc":"2.0","id":1,"err`, + wantHasError: false, + }, + { + name: "cut right after the error key opens", + body: `{"jsonrpc":"2.0","id":1,"error":{`, + wantHasError: true, + }, + { + name: "data before message is skipped", + body: `{"jsonrpc":"2.0","id":1,"error":{"code":-32000,"data":{"a":[1,2]},"message":"late"}}`, + wantHasError: true, + wantCode: -32000, + wantMessage: "late", + }, + { + name: "truncated result is not an error", + body: `{"jsonrpc":"2.0","id":1,"result":{"content":[{"type":"text","text":"` + strings.Repeat("a", 600), + wantHasError: false, + }, + { + name: "nested error key is ignored", + body: `{"jsonrpc":"2.0","id":1,"result":{"error":{"code":1,"message":"x"},"pad":"` + strings.Repeat("a", 600), + wantHasError: false, + }, + { + name: "null error", + body: `{"jsonrpc":"2.0","id":1,"error":null}`, + wantHasError: false, + }, + { + name: "not an object", + body: `[{"error":{"code":1}}`, + wantHasError: false, + }, + { + name: "empty", + body: "", + wantHasError: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + got := ParseTruncatedMCPResponse([]byte(tt.body)) + assert.Equal(t, tt.wantHasError, got.HasError) + assert.Equal(t, tt.wantCode, got.ErrorCode) + assert.Equal(t, tt.wantMessage, got.ErrorMessage) + }) + } +}