diff --git a/pkg/mcp/classification_response.go b/pkg/mcp/classification_response.go new file mode 100644 index 0000000000..4268dd826e --- /dev/null +++ b/pkg/mcp/classification_response.go @@ -0,0 +1,89 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package mcp + +import ( + "bytes" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" +) + +// WriteClassificationError writes an HTTP 400 response with a JSON-RPC error +// body for an mcp.ClassifyRevision failure. Use this with http.ResponseWriter +// in the streamable HTTP proxy. +func WriteClassificationError(w http.ResponseWriter, requestID any, err error) { + body := classificationErrorBody(requestID, err) + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusBadRequest) + //nolint:gosec // G104: writing a JSON-RPC error response to an HTTP client + _, _ = w.Write(body) +} + +// ClassificationErrorResponse constructs an *http.Response with HTTP 400 and +// a JSON-RPC error body for an mcp.ClassifyRevision failure. Use this in +// httputil.ReverseProxy.ModifyResponse/RoundTrip (transparent proxy) where no +// http.ResponseWriter is available. +// +// req is attached as resp.Request, matching session.NotFoundResponse: a +// RoundTripper-produced *http.Response is generally expected to carry the +// request it answers, and httputil.ReverseProxy relies on that field. +func ClassificationErrorResponse(req *http.Request, requestID any, err error) *http.Response { + body := classificationErrorBody(requestID, err) + hdr := make(http.Header) + hdr.Set("Content-Type", "application/json") + return &http.Response{ + StatusCode: http.StatusBadRequest, + Status: fmt.Sprintf("%d %s", http.StatusBadRequest, http.StatusText(http.StatusBadRequest)), + Proto: "HTTP/1.1", + ProtoMajor: 1, + ProtoMinor: 1, + Header: hdr, + ContentLength: int64(len(body)), + Body: io.NopCloser(bytes.NewReader(body)), + Request: req, + } +} + +// classificationErrorBody renders an mcp.ClassifyRevision error as a JSON-RPC +// error body, modeled on session.NotFoundBody: the body is marshaled first +// (with a hand-crafted fallback on marshal failure) so callers only write +// headers/status once a valid body is ready. +// +// It uses the error's Code(), Error() message, and Data() (when non-empty) if +// the error implements CodedError, falling back to the standard JSON-RPC +// Invalid Params code otherwise -- a fallback that is currently unreachable, +// since every error ClassifyRevision returns implements CodedError. +func classificationErrorBody(requestID any, err error) []byte { + code := CodeInvalidParams + var coded CodedError + var data map[string]any + if errors.As(err, &coded) { + code = coded.Code() + data = coded.Data() + } + + errBody := map[string]any{ + "code": code, + "message": err.Error(), + } + if len(data) > 0 { + errBody["data"] = data + } + resp := map[string]any{ + "jsonrpc": "2.0", + "error": errBody, + "id": requestID, + } + + body, marshalErr := json.Marshal(resp) + if marshalErr != nil { + // This should never happen with simple map types, but return a + // hand-crafted fallback to guarantee a valid JSON-RPC error. + return []byte(`{"jsonrpc":"2.0","error":{"code":-32602,"message":"Invalid params"},"id":null}`) + } + return body +} diff --git a/pkg/mcp/classification_response_test.go b/pkg/mcp/classification_response_test.go new file mode 100644 index 0000000000..631af936c5 --- /dev/null +++ b/pkg/mcp/classification_response_test.go @@ -0,0 +1,165 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package mcp + +import ( + "encoding/json" + "errors" + "io" + "net/http" + "net/http/httptest" + "testing" +) + +func TestClassificationError(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + requestID any + err error + wantCode int64 + wantData bool + }{ + { + name: "header mismatch", + requestID: "req-1", + err: &HeaderMismatchError{Header: "2026-07-28", Body: "2025-11-25"}, + wantCode: CodeHeaderMismatch, + wantData: true, + }, + { + name: "unsupported version", + requestID: float64(42), + err: &UnsupportedVersionError{Requested: "1999-01-01", Supported: []string{MCPVersionModern}}, + wantCode: CodeUnsupportedProtocolVersion, + wantData: true, + }, + { + name: "missing client capability", + requestID: "req-3", + err: &MissingClientCapabilityError{RequiredCapabilities: map[string]any{}}, + wantCode: CodeMissingClientCapability, + wantData: true, + }, + { + name: "missing modern metadata", + requestID: "req-4", + err: &MissingModernMetadataError{}, + wantCode: CodeInvalidParams, + wantData: false, // Data() returns an empty (non-nil) map, so len(data) == 0 + }, + { + name: "nil request id", + requestID: nil, + err: &HeaderMismatchError{Header: "a", Body: "b"}, + wantCode: CodeHeaderMismatch, + wantData: true, + }, + { + name: "plain non-coded error falls back to invalid params", + requestID: "req-6", + err: errors.New("boom"), + wantCode: CodeInvalidParams, + wantData: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + rec := httptest.NewRecorder() + WriteClassificationError(rec, tt.requestID, tt.err) + + req := httptest.NewRequest(http.MethodPost, "/mcp", nil) + resp := ClassificationErrorResponse(req, tt.requestID, tt.err) + + if rec.Code != http.StatusBadRequest { + t.Fatalf("WriteClassificationError: status = %d, want %d", rec.Code, http.StatusBadRequest) + } + if resp.StatusCode != http.StatusBadRequest { + t.Fatalf("ClassificationErrorResponse: status = %d, want %d", resp.StatusCode, http.StatusBadRequest) + } + if ct := rec.Header().Get("Content-Type"); ct != "application/json" { + t.Fatalf("WriteClassificationError: Content-Type = %q, want application/json", ct) + } + if ct := resp.Header.Get("Content-Type"); ct != "application/json" { + t.Fatalf("ClassificationErrorResponse: Content-Type = %q, want application/json", ct) + } + if resp.Request != req { + t.Fatalf("ClassificationErrorResponse: Request not attached") + } + if resp.Body == nil { + t.Fatalf("ClassificationErrorResponse: Body is nil") + } + + respBody, err := io.ReadAll(resp.Body) + if err != nil { + t.Fatalf("reading response body: %v", err) + } + if resp.ContentLength != int64(len(respBody)) { + t.Fatalf("ContentLength = %d, want %d", resp.ContentLength, len(respBody)) + } + + wireBody := rec.Body.Bytes() + if string(wireBody) != string(respBody) { + t.Fatalf("WriteClassificationError and ClassificationErrorResponse bodies differ:\n%s\nvs\n%s", wireBody, respBody) + } + + var decoded struct { + JSONRPC string `json:"jsonrpc"` + Error struct { + Code int64 `json:"code"` + Message string `json:"message"` + Data map[string]any `json:"data"` + } `json:"error"` + ID any `json:"id"` + } + if err := json.Unmarshal(wireBody, &decoded); err != nil { + t.Fatalf("unmarshaling response body: %v", err) + } + + if decoded.JSONRPC != "2.0" { + t.Errorf("jsonrpc = %q, want \"2.0\"", decoded.JSONRPC) + } + if decoded.Error.Code != tt.wantCode { + t.Errorf("code = %d, want %d", decoded.Error.Code, tt.wantCode) + } + if decoded.Error.Message != tt.err.Error() { + t.Errorf("message = %q, want %q", decoded.Error.Message, tt.err.Error()) + } + if tt.wantData && len(decoded.Error.Data) == 0 { + t.Errorf("expected non-empty data, got %v", decoded.Error.Data) + } + if !tt.wantData && len(decoded.Error.Data) != 0 { + t.Errorf("expected no data, got %v", decoded.Error.Data) + } + + gotID, wantID := decoded.ID, tt.requestID + if wantID == nil { + if gotID != nil { + t.Errorf("id = %v, want nil", gotID) + } + } else if gotID != wantID { + t.Errorf("id = %v (%T), want %v (%T)", gotID, gotID, wantID, wantID) + } + }) + } +} + +// TestClassificationErrorBodyMarshalFallback verifies the defensive fallback: +// if the response fails to marshal (unreachable in production, where requestID +// is always a json.RawMessage), classificationErrorBody still returns a valid +// JSON-RPC error rather than an empty body. A channel is not JSON-marshalable, +// so it forces json.Marshal to fail. +func TestClassificationErrorBodyMarshalFallback(t *testing.T) { + t.Parallel() + + const want = `{"jsonrpc":"2.0","error":{"code":-32602,"message":"Invalid params"},"id":null}` + got := classificationErrorBody(make(chan int), errors.New("boom")) + if string(got) != want { + t.Fatalf("fallback body = %s, want %s", got, want) + } +} diff --git a/pkg/transport/proxy/streamable/streamable_proxy.go b/pkg/transport/proxy/streamable/streamable_proxy.go index 3877473e80..482a5ef606 100644 --- a/pkg/transport/proxy/streamable/streamable_proxy.go +++ b/pkg/transport/proxy/streamable/streamable_proxy.go @@ -784,7 +784,7 @@ func (p *HTTPProxy) resolveSessionForRequest( protoHeader := r.Header.Get("MCP-Protocol-Version") rev, err := mcp.ClassifyRevision(req.Method, meta, protoHeader) if err != nil { - writeClassificationError(w, req.ID.Raw(), err) + mcp.WriteClassificationError(w, req.ID.Raw(), err) return "", false, err } @@ -841,49 +841,6 @@ func (p *HTTPProxy) resolveSessionForRequest( return sessID, false, nil } -// writeClassificationError renders an mcp.ClassifyRevision error as an HTTP 400 -// JSON-RPC error response, modeled on session.NotFoundBody/WriteNotFound: the -// body is marshaled first (with a hand-crafted fallback on marshal failure) so -// headers and status are only written once a valid body is ready. -// It uses the error's Code(), Error() message, and Data() (when non-empty) if -// the error implements mcp.CodedError, falling back to the standard JSON-RPC -// Invalid Params code otherwise -- a fallback that is currently unreachable, -// since every error ClassifyRevision returns implements mcp.CodedError. -func writeClassificationError(w http.ResponseWriter, requestID any, err error) { - code := mcp.CodeInvalidParams - var coded mcp.CodedError - var data map[string]any - if errors.As(err, &coded) { - code = coded.Code() - data = coded.Data() - } - - errBody := map[string]any{ - "code": code, - "message": err.Error(), - } - if len(data) > 0 { - errBody["data"] = data - } - resp := map[string]any{ - "jsonrpc": "2.0", - "error": errBody, - "id": requestID, - } - - body, marshalErr := json.Marshal(resp) - if marshalErr != nil { - // This should never happen with simple map types, but return a - // hand-crafted fallback to guarantee a valid JSON-RPC error. - body = []byte(`{"jsonrpc":"2.0","error":{"code":-32602,"message":"Invalid params"},"id":null}`) - } - - w.Header().Set("Content-Type", "application/json") - w.WriteHeader(http.StatusBadRequest) - //nolint:gosec // G104: writing a JSON-RPC error response to an HTTP client - _, _ = w.Write(body) -} - func isBatch(body []byte) bool { t := bytes.TrimSpace(body) return len(t) > 0 && t[0] == '[' diff --git a/pkg/transport/proxy/transparent/method_gate_test.go b/pkg/transport/proxy/transparent/method_gate_test.go index 9bfc2e35c5..4f761e9f51 100644 --- a/pkg/transport/proxy/transparent/method_gate_test.go +++ b/pkg/transport/proxy/transparent/method_gate_test.go @@ -4,69 +4,83 @@ package transparent import ( + "io" "net/http" "net/http/httptest" + "strings" "testing" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive/pkg/mcp" ) -func TestStatelessMethodGate(t *testing.T) { +func TestMethodGate(t *testing.T) { t.Parallel() - inner := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { - w.WriteHeader(http.StatusOK) - }) - tests := []struct { name string + stateless bool + protocolHeader string method string - expectedStatus int - expectAllow bool + expectGated bool }{ - { - name: "GET returns 405 with Allow header", - method: http.MethodGet, - expectedStatus: http.StatusMethodNotAllowed, - expectAllow: true, - }, - { - name: "HEAD returns 405 with Allow header", - method: http.MethodHead, - expectedStatus: http.StatusMethodNotAllowed, - expectAllow: true, - }, - { - name: "DELETE returns 405 with Allow header", - method: http.MethodDelete, - expectedStatus: http.StatusMethodNotAllowed, - expectAllow: true, - }, - { - name: "POST is forwarded", - method: http.MethodPost, - expectedStatus: http.StatusOK, - }, - { - name: "PUT is forwarded", - method: http.MethodPut, - expectedStatus: http.StatusOK, - }, + // stateless=true gates GET/HEAD/DELETE regardless of header + {"stateless GET gated", true, "", http.MethodGet, true}, + {"stateless HEAD gated", true, "", http.MethodHead, true}, + {"stateless DELETE gated", true, "", http.MethodDelete, true}, + {"stateless POST passes", true, "", http.MethodPost, false}, + {"stateless with Modern header GET gated", true, mcp.MCPVersionModern, http.MethodGet, true}, + + // stateless=false, Modern header gates GET/HEAD/DELETE + {"Modern header GET gated", false, mcp.MCPVersionModern, http.MethodGet, true}, + {"Modern header HEAD gated", false, mcp.MCPVersionModern, http.MethodHead, true}, + {"Modern header DELETE gated", false, mcp.MCPVersionModern, http.MethodDelete, true}, + {"Modern header POST passes", false, mcp.MCPVersionModern, http.MethodPost, false}, + + // stateless=false, no/other header: legacy behavior, nothing gated + {"legacy GET passes", false, "", http.MethodGet, false}, + {"legacy HEAD passes", false, "", http.MethodHead, false}, + {"legacy DELETE passes", false, "", http.MethodDelete, false}, + {"legacy POST passes", false, "", http.MethodPost, false}, + {"other header GET passes", false, "2025-11-25", http.MethodGet, false}, } for _, tc := range tests { t.Run(tc.name, func(t *testing.T) { t.Parallel() - handler := statelessMethodGate(inner) + var gotBody string + inner := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + gotBody = string(body) + w.WriteHeader(http.StatusOK) + }) + + p := &TransparentProxy{stateless: tc.stateless} + handler := p.methodGate(inner) rec := httptest.NewRecorder() - req := httptest.NewRequest(tc.method, "/", nil) + + var body string + if tc.method == http.MethodPost { + body = "request-body-payload" + } + req := httptest.NewRequest(tc.method, "/", strings.NewReader(body)) + if tc.protocolHeader != "" { + req.Header.Set("MCP-Protocol-Version", tc.protocolHeader) + } handler.ServeHTTP(rec, req) - assert.Equal(t, tc.expectedStatus, rec.Code) - if tc.expectAllow { + if tc.expectGated { + assert.Equal(t, http.StatusMethodNotAllowed, rec.Code) assert.Equal(t, "POST, OPTIONS", rec.Header().Get("Allow")) + } else { + assert.Equal(t, http.StatusOK, rec.Code) + if tc.method == http.MethodPost { + require.Equal(t, body, gotBody, "POST body must reach next handler intact") + } } }) } diff --git a/pkg/transport/proxy/transparent/revision_classification_test.go b/pkg/transport/proxy/transparent/revision_classification_test.go new file mode 100644 index 0000000000..7551c74508 --- /dev/null +++ b/pkg/transport/proxy/transparent/revision_classification_test.go @@ -0,0 +1,288 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package transparent + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive/pkg/mcp" +) + +// roundTripFunc adapts a function to http.RoundTripper, letting tests spy on +// whether the backend transport was invoked without standing up a real server. +type roundTripFunc func(*http.Request) (*http.Response, error) + +func (f roundTripFunc) RoundTrip(r *http.Request) (*http.Response, error) { return f(r) } + +func TestParseRPCRequest(t *testing.T) { + t.Parallel() + + tp := &tracingTransport{p: &TransparentProxy{targetURI: "http://backend"}} + + tests := []struct { + name string + body string + wantMethod string + wantID string // "" means nil/absent + wantSingleRequest bool + wantSawInitialize bool + }{ + { + name: "single request with id", + body: `{"jsonrpc":"2.0","id":1,"method":"tools/list"}`, + wantMethod: "tools/list", + wantID: "1", + wantSingleRequest: true, + }, + { + name: "single initialize", + body: `{"jsonrpc":"2.0","id":1,"method":"initialize"}`, + wantMethod: "initialize", + wantID: "1", + wantSingleRequest: true, + wantSawInitialize: true, + }, + { + name: "notification has no id", + body: `{"jsonrpc":"2.0","method":"notifications/progress"}`, + wantMethod: "notifications/progress", + wantSingleRequest: false, + }, + { + name: "explicit null id is not a valid request id", + body: `{"jsonrpc":"2.0","id":null,"method":"tools/list"}`, + wantMethod: "tools/list", + wantID: "null", + wantSingleRequest: false, + }, + { + name: "response-shaped body has no method", + body: `{"jsonrpc":"2.0","id":1,"result":{}}`, + wantMethod: "", + wantID: "1", + wantSingleRequest: false, + }, + { + name: "batch with initialize", + body: `[{"jsonrpc":"2.0","id":1,"method":"initialize"},{"jsonrpc":"2.0","id":2,"method":"tools/list"}]`, + wantSingleRequest: false, + wantSawInitialize: true, + }, + { + name: "batch without initialize", + body: `[{"jsonrpc":"2.0","id":1,"method":"tools/list"},{"jsonrpc":"2.0","id":2,"method":"tools/call"}]`, + wantSingleRequest: false, + }, + { + name: "malformed JSON", + body: `{not valid json`, + wantSingleRequest: false, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + method, _, id, singleRequest, sawInitialize := tp.parseRPCRequest([]byte(tc.body)) + + assert.Equal(t, tc.wantMethod, method) + assert.Equal(t, tc.wantSingleRequest, singleRequest) + assert.Equal(t, tc.wantSawInitialize, sawInitialize) + if tc.wantID == "" { + assert.Empty(t, string(id)) + } else { + assert.Equal(t, tc.wantID, string(id)) + } + }) + } +} + +// TestRoundTripClassifiesModernRequests drives tracingTransport.RoundTrip +// directly (bypassing httputil.ReverseProxy) with a spy backend RoundTripper, +// so tests can assert both the returned response and whether the backend was +// ever contacted. +func TestRoundTripClassifiesModernRequests(t *testing.T) { + t.Parallel() + + newProxy := func(spy http.RoundTripper) (*tracingTransport, *TransparentProxy) { + p := NewTransparentProxy("127.0.0.1", 0, "", nil, nil, nil, false, false, + "streamable-http", nil, nil, "", false) + return newTracingTransport(spy, p), p + } + + t.Run("malformed Modern single-request is rejected before the backend is called", func(t *testing.T) { + t.Parallel() + + var backendCalled atomic.Bool + spy := roundTripFunc(func(*http.Request) (*http.Response, error) { + backendCalled.Store(true) + return httptest.NewRecorder().Result(), nil + }) + tt, _ := newProxy(spy) + + // Header claims Modern but the body carries no _meta at all: a + // HeaderMismatchError (bad/absent body version, non-empty header). + req := httptest.NewRequest(http.MethodPost, "/mcp", + strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"tools/call"}`)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("MCP-Protocol-Version", mcp.MCPVersionModern) + + resp, err := tt.RoundTrip(req) + require.NoError(t, err) + require.NotNil(t, resp) + assert.Equal(t, http.StatusBadRequest, resp.StatusCode) + assert.False(t, backendCalled.Load(), "backend must not be contacted for a rejected request") + }) + + t.Run("well-formed Modern single-request falls through to the backend", func(t *testing.T) { + t.Parallel() + + var backendCalled atomic.Bool + spy := roundTripFunc(func(*http.Request) (*http.Response, error) { + backendCalled.Store(true) + return &http.Response{StatusCode: http.StatusOK, Header: make(http.Header), Body: http.NoBody}, nil + }) + tt, _ := newProxy(spy) + + // Not "initialize", has a valid id, and _meta carries a matching + // protocolVersion plus clientCapabilities: ClassifyRevision returns + // (RevisionModern, nil), so the request must reach the backend. + body := `{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"_meta":{` + + `"io.modelcontextprotocol/protocolVersion":"` + mcp.MCPVersionModern + `",` + + `"io.modelcontextprotocol/clientCapabilities":{}}}}` + req := httptest.NewRequest(http.MethodPost, "/mcp", strings.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("MCP-Protocol-Version", mcp.MCPVersionModern) + + resp, err := tt.RoundTrip(req) + require.NoError(t, err) + require.NotNil(t, resp) + assert.Equal(t, http.StatusOK, resp.StatusCode, "a well-formed Modern request must not be rejected") + assert.True(t, backendCalled.Load(), "backend must be contacted for a well-formed Modern request") + }) + + t.Run("batch with Modern header is forwarded, not rejected", func(t *testing.T) { + t.Parallel() + + var backendCalled atomic.Bool + spy := roundTripFunc(func(*http.Request) (*http.Response, error) { + backendCalled.Store(true) + return &http.Response{StatusCode: http.StatusOK, Header: make(http.Header), Body: http.NoBody}, nil + }) + tt, _ := newProxy(spy) + + req := httptest.NewRequest(http.MethodPost, "/mcp", + strings.NewReader(`[{"jsonrpc":"2.0","id":1,"method":"tools/call"}]`)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("MCP-Protocol-Version", mcp.MCPVersionModern) + + resp, err := tt.RoundTrip(req) + require.NoError(t, err) + require.NotNil(t, resp) + assert.Equal(t, http.StatusOK, resp.StatusCode, "batches are never classified, so they must be forwarded") + assert.True(t, backendCalled.Load(), "backend must be contacted for a forwarded batch") + }) + + t.Run("large-integer id is preserved verbatim in the 400 body", func(t *testing.T) { + t.Parallel() + + const largeID = "9007199254740993" // 2^53 + 1: loses precision if round-tripped through float64 + spy := roundTripFunc(func(*http.Request) (*http.Response, error) { + t.Fatal("backend must not be contacted for a rejected request") + return nil, nil + }) + tt, _ := newProxy(spy) + + req := httptest.NewRequest(http.MethodPost, "/mcp", + strings.NewReader(`{"jsonrpc":"2.0","id":`+largeID+`,"method":"tools/call"}`)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("MCP-Protocol-Version", mcp.MCPVersionModern) + + resp, err := tt.RoundTrip(req) + require.NoError(t, err) + require.NotNil(t, resp) + require.Equal(t, http.StatusBadRequest, resp.StatusCode) + + var decoded struct { + ID json.RawMessage `json:"id"` + } + require.NoError(t, json.NewDecoder(resp.Body).Decode(&decoded)) + assert.Equal(t, largeID, string(decoded.ID), "large integer id must not lose precision") + }) + + t.Run("Modern 200 response flips serverInitialized", func(t *testing.T) { + t.Parallel() + + spy := roundTripFunc(func(*http.Request) (*http.Response, error) { + return &http.Response{StatusCode: http.StatusOK, Header: make(http.Header), Body: http.NoBody}, nil + }) + tt, p := newProxy(spy) + require.False(t, p.serverInitialized(), "precondition: latch starts unset") + + body := `{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"_meta":{` + + `"io.modelcontextprotocol/protocolVersion":"` + mcp.MCPVersionModern + `",` + + `"io.modelcontextprotocol/clientCapabilities":{}}}}` + req := httptest.NewRequest(http.MethodPost, "/mcp", strings.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("MCP-Protocol-Version", mcp.MCPVersionModern) + + resp, err := tt.RoundTrip(req) + require.NoError(t, err) + require.Equal(t, http.StatusOK, resp.StatusCode) + assert.True(t, p.serverInitialized(), "a 200 to a Modern request must flip the readiness latch") + }) + + t.Run("Modern non-200 response does not flip serverInitialized", func(t *testing.T) { + t.Parallel() + + spy := roundTripFunc(func(*http.Request) (*http.Response, error) { + return &http.Response{StatusCode: http.StatusInternalServerError, Header: make(http.Header), Body: http.NoBody}, nil + }) + tt, p := newProxy(spy) + require.False(t, p.serverInitialized(), "precondition: latch starts unset") + + body := `{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"_meta":{` + + `"io.modelcontextprotocol/protocolVersion":"` + mcp.MCPVersionModern + `",` + + `"io.modelcontextprotocol/clientCapabilities":{}}}}` + req := httptest.NewRequest(http.MethodPost, "/mcp", strings.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("MCP-Protocol-Version", mcp.MCPVersionModern) + + resp, err := tt.RoundTrip(req) + require.NoError(t, err) + require.Equal(t, http.StatusInternalServerError, resp.StatusCode) + assert.False(t, p.serverInitialized(), "a non-200 response must not flip the readiness latch") + }) + + t.Run("Legacy 200 response without session header or initialize does not flip serverInitialized", func(t *testing.T) { + t.Parallel() + + spy := roundTripFunc(func(*http.Request) (*http.Response, error) { + return &http.Response{StatusCode: http.StatusOK, Header: make(http.Header), Body: http.NoBody}, nil + }) + tt, p := newProxy(spy) + require.False(t, p.serverInitialized(), "precondition: latch starts unset") + + // Legacy request (no MCP-Protocol-Version header, not initialize) whose + // 200 response carries no Mcp-Session-Id: neither existing branch fires, + // and the new Modern branch must not broaden to cover this case either. + req := httptest.NewRequest(http.MethodPost, "/mcp", + strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"tools/list"}`)) + req.Header.Set("Content-Type", "application/json") + + resp, err := tt.RoundTrip(req) + require.NoError(t, err) + require.Equal(t, http.StatusOK, resp.StatusCode) + assert.False(t, p.serverInitialized(), "a Legacy 200 with no session header must not flip the readiness latch") + }) +} diff --git a/pkg/transport/proxy/transparent/revision_guard_regression_test.go b/pkg/transport/proxy/transparent/revision_guard_regression_test.go new file mode 100644 index 0000000000..698cadd1a5 --- /dev/null +++ b/pkg/transport/proxy/transparent/revision_guard_regression_test.go @@ -0,0 +1,249 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package transparent + +import ( + "io" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + + "github.com/google/uuid" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive/pkg/mcp" + "github.com/stacklok/toolhive/pkg/transport/session" +) + +// This file is a SECURITY regression suite: it proves that a client-forgeable +// Modern revision signal (MCP-Protocol-Version header and/or params._meta) +// cannot bypass the Legacy session machinery in tracingTransport.RoundTrip. +// That machinery is deliberately keyed on Mcp-Session-Id PRESENCE, never on +// the classified revision, because revision is derived entirely from +// client-controlled input while Mcp-Session-Id is validated against +// server-side session state. See transparent_proxy.go's RoundTrip guard +// comment and revision_classification_test.go for the classification-only +// coverage this file complements. +// +// Scope: this suite covers single-request Modern-signal forgery only. Batch +// payloads are never revision-classified (batching was removed from MCP in +// 2025-06-18), so they cannot carry a Modern signal and are out of scope +// here. The separate, pre-existing question of whether a batch containing +// "initialize" should be exempt from the unknown-session guard (see +// TestRoundTripAllowsBatchInitializeWithUnknownSession) is Legacy behavior +// tracked for follow-up, not a Modern-spoofing vector. + +// modernToolsCallBody is a well-formed Modern (2026-07-28) JSON-RPC request: +// not "initialize", a valid non-null id, and params._meta carrying both +// reserved keys required for a nil-error classification. Reused verbatim (and +// with a missing header in one case) across every test below so each test +// starts from a request that ClassifyRevision genuinely accepts as Modern. +const modernToolsCallBody = `{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"_meta":{` + + `"io.modelcontextprotocol/protocolVersion":"` + mcp.MCPVersionModern + `",` + + `"io.modelcontextprotocol/clientCapabilities":{}}}}` + +// newGuardTransport builds a tracingTransport backed by a spy RoundTripper, +// mirroring the harness in revision_classification_test.go. targetURI is a +// syntactically valid placeholder: backendRecovery.reinitializeAndReplay +// parses it with url.Parse but the spy intercepts every call, so it is never +// actually dialed. +func newGuardTransport(spy http.RoundTripper) (*tracingTransport, *TransparentProxy) { + p := NewTransparentProxy("127.0.0.1", 0, "http://backend", nil, nil, nil, false, false, + "streamable-http", nil, nil, "", false) + return newTracingTransport(spy, p), p +} + +// assertClassifiesModernNil is a precondition check, not the regression +// itself: it confirms the given body/header combination really does +// classify as Modern with a nil error. If this fails, the test below it is +// not exercising a forged-Modern scenario at all. +func assertClassifiesModernNil(t *testing.T, tt *tracingTransport, body []byte, protoHeader string) { + t.Helper() + method, params, _, singleRequest, _ := tt.parseRPCRequest(body) + require.True(t, singleRequest, "precondition: body must parse as a single JSON-RPC request") + meta := mcp.ExtractMeta(params) + rev, err := mcp.ClassifyRevision(method, meta, protoHeader) + require.NoError(t, err, "precondition: body must classify Modern with a nil error") + require.Equal(t, mcp.RevisionModern, rev, "precondition: body must classify as Modern") +} + +// TestGuardUnknownSessionFiresDespiteForgedModernRevision is the core +// mutation-check case: a well-formed Modern request with an unknown/stale +// Mcp-Session-Id must still be rejected with the session-not-found response, +// exactly as it would be for a Legacy request (see +// TestRoundTripReturns404ForUnknownSession in backend_routing_test.go). +// +// This test FAILS if the guard in tracingTransport.RoundTrip is ever +// re-keyed on `revision == mcp.RevisionModern` (or any check derived from +// the classified revision) instead of Mcp-Session-Id presence — that would +// let a forged Modern signal skip session validation entirely. +func TestGuardUnknownSessionFiresDespiteForgedModernRevision(t *testing.T) { + t.Parallel() + + var backendCalled atomic.Bool + spy := roundTripFunc(func(*http.Request) (*http.Response, error) { + backendCalled.Store(true) + return &http.Response{StatusCode: http.StatusOK, Header: make(http.Header), Body: http.NoBody}, nil + }) + tt, _ := newGuardTransport(spy) + + assertClassifiesModernNil(t, tt, []byte(modernToolsCallBody), mcp.MCPVersionModern) + + req := httptest.NewRequest(http.MethodPost, "/mcp", strings.NewReader(modernToolsCallBody)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("MCP-Protocol-Version", mcp.MCPVersionModern) + req.Header.Set("Mcp-Session-Id", uuid.New().String()) // unknown/stale, never added to the session manager + + resp, err := tt.RoundTrip(req) + require.NoError(t, err) + require.NotNil(t, resp) + require.Equal(t, http.StatusNotFound, resp.StatusCode) + + body, err := io.ReadAll(resp.Body) + require.NoError(t, err) + assert.Contains(t, string(body), `"code":-32001`) + assert.False(t, backendCalled.Load(), + "backend must not be contacted: the unknown-session guard must fire regardless of the classified revision") +} + +// TestGuardBackendSIDRewriteStillHappensWithForgedModernRevision verifies +// that a well-formed Modern request against a KNOWN session whose metadata +// carries a backend_sid still has its outbound Mcp-Session-Id rewritten to +// that backend SID — the rewrite (like the guard above) is driven by session +// metadata, not by the classified revision. Mirrors the assertion style of +// backend_routing_test.go's TestRoundTripReinitializesOnBackend404 (which +// checks sessionMetadataBackendSID) but observes the header directly via a +// spy, as in revision_classification_test.go. +func TestGuardBackendSIDRewriteStillHappensWithForgedModernRevision(t *testing.T) { + t.Parallel() + + const backendSID = "backend-assigned-opaque-sid" + var gotSID atomic.Value + spy := roundTripFunc(func(r *http.Request) (*http.Response, error) { + gotSID.Store(r.Header.Get("Mcp-Session-Id")) + return &http.Response{StatusCode: http.StatusOK, Header: make(http.Header), Body: http.NoBody}, nil + }) + tt, p := newGuardTransport(spy) + + assertClassifiesModernNil(t, tt, []byte(modernToolsCallBody), mcp.MCPVersionModern) + + clientSID := uuid.New().String() + sess := session.NewProxySession(clientSID) + sess.SetMetadata(sessionMetadataBackendSID, backendSID) + require.NoError(t, p.sessionManager.AddSession(sess)) + + req := httptest.NewRequest(http.MethodPost, "/mcp", strings.NewReader(modernToolsCallBody)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("MCP-Protocol-Version", mcp.MCPVersionModern) + req.Header.Set("Mcp-Session-Id", clientSID) + + resp, err := tt.RoundTrip(req) + require.NoError(t, err) + require.Equal(t, http.StatusOK, resp.StatusCode) + + assert.Equal(t, backendSID, gotSID.Load(), + "outbound Mcp-Session-Id must be rewritten to backend_sid even though the request forges a Modern revision signal") +} + +// TestGuardReinitRecoveryStillTriggersWithForgedModernRevision verifies that +// the transparent re-initialize-and-replay recovery path (triggered by a 404 +// from a known session with a stored init body) still fires for a +// well-formed Modern request. Mirrors backend_routing_test.go's +// TestRoundTripReinitializesOnBackend404, but the request itself carries a +// forged Modern signal to prove recovery is unaffected by classified +// revision. +func TestGuardReinitRecoveryStillTriggersWithForgedModernRevision(t *testing.T) { + t.Parallel() + + freshSessionID := uuid.New().String() + var initCalls, otherCalls atomic.Int32 + spy := roundTripFunc(func(r *http.Request) (*http.Response, error) { + body, _ := io.ReadAll(r.Body) + if strings.Contains(string(body), `"initialize"`) { + initCalls.Add(1) + return &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Mcp-Session-Id": []string{freshSessionID}}, + Body: http.NoBody, + }, nil + } + n := otherCalls.Add(1) + if n == 1 { + // First non-initialize forward: simulate the backend pod having + // lost its in-memory session state. + return &http.Response{StatusCode: http.StatusNotFound, Header: make(http.Header), Body: http.NoBody}, nil + } + // Second non-initialize forward is the replay after re-init. + return &http.Response{StatusCode: http.StatusOK, Header: make(http.Header), Body: http.NoBody}, nil + }) + tt, p := newGuardTransport(spy) + + assertClassifiesModernNil(t, tt, []byte(modernToolsCallBody), mcp.MCPVersionModern) + + clientSID := uuid.New().String() + sess := session.NewProxySession(clientSID) + sess.SetMetadata(sessionMetadataInitBody, `{"jsonrpc":"2.0","id":1,"method":"initialize"}`) + require.NoError(t, p.sessionManager.AddSession(sess)) + + req := httptest.NewRequest(http.MethodPost, "/mcp", strings.NewReader(modernToolsCallBody)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("MCP-Protocol-Version", mcp.MCPVersionModern) + req.Header.Set("Mcp-Session-Id", clientSID) + + resp, err := tt.RoundTrip(req) + require.NoError(t, err) + require.Equal(t, http.StatusOK, resp.StatusCode, + "client must see 200 after transparent re-init, even for a forged-Modern request") + assert.Equal(t, int32(1), initCalls.Load(), "exactly one re-initialize call expected") + assert.Equal(t, int32(2), otherCalls.Load(), "original forward + replay expected") + + updated, ok := p.sessionManager.Get(normalizeSessionID(clientSID)) + require.True(t, ok, "session should still exist after re-init") + backendSID, exists := updated.GetMetadataValue(sessionMetadataBackendSID) + require.True(t, exists) + assert.Equal(t, freshSessionID, backendSID, + "backend_sid must be updated by the recovery path regardless of the classified revision") +} + +// TestGuardDeleteCleanupStillWorksWithBodyMetaButNoHeader verifies that DELETE +// session cleanup (keyed on Mcp-Session-Id, req.Method and response status — +// see RoundTrip's DELETE cleanup block) still runs when the DELETE body +// carries Modern _meta but no MCP-Protocol-Version header. Note: a DELETE +// that also carried the MCP-Protocol-Version: 2026-07-28 header would be +// rejected with 405 at the header-only methodGate and would never reach +// RoundTrip at all (see method_gate_test.go's "Modern header DELETE gated" +// case) — this test exercises the one shape of forged-Modern DELETE that +// does reach RoundTrip. Mirrors delete_session_test.go's cleanup assertion. +func TestGuardDeleteCleanupStillWorksWithBodyMetaButNoHeader(t *testing.T) { + t.Parallel() + + spy := roundTripFunc(func(*http.Request) (*http.Response, error) { + return &http.Response{StatusCode: http.StatusOK, Header: make(http.Header), Body: http.NoBody}, nil + }) + tt, p := newGuardTransport(spy) + + // No MCP-Protocol-Version header: classification relies solely on the + // reserved _meta keys, per mcp.ClassifyRevision's documented signal rules. + assertClassifiesModernNil(t, tt, []byte(modernToolsCallBody), "") + + clientSID := uuid.New().String() + sess := session.NewProxySession(clientSID) + require.NoError(t, p.sessionManager.AddSession(sess)) + + req := httptest.NewRequest(http.MethodDelete, "/mcp", strings.NewReader(modernToolsCallBody)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Mcp-Session-Id", clientSID) + // Deliberately no MCP-Protocol-Version header. + + resp, err := tt.RoundTrip(req) + require.NoError(t, err) + require.Equal(t, http.StatusOK, resp.StatusCode) + + _, ok := p.sessionManager.Get(normalizeSessionID(clientSID)) + assert.False(t, ok, + "DELETE session cleanup must still key off Mcp-Session-Id presence even when the body carries Modern _meta") +} diff --git a/pkg/transport/proxy/transparent/transparent_proxy.go b/pkg/transport/proxy/transparent/transparent_proxy.go index 47e1266393..7863bac92b 100644 --- a/pkg/transport/proxy/transparent/transparent_proxy.go +++ b/pkg/transport/proxy/transparent/transparent_proxy.go @@ -32,6 +32,7 @@ import ( "github.com/stacklok/toolhive/pkg/auth" "github.com/stacklok/toolhive/pkg/bodylimit" "github.com/stacklok/toolhive/pkg/healthcheck" + "github.com/stacklok/toolhive/pkg/mcp" "github.com/stacklok/toolhive/pkg/transport/proxy/socket" "github.com/stacklok/toolhive/pkg/transport/session" "github.com/stacklok/toolhive/pkg/transport/types" @@ -259,7 +260,7 @@ func WithRemoteRawQuery(rawQuery string) Option { } // WithStateless configures the proxy for stateless streamable-HTTP servers. -// In stateless mode, incoming GET and DELETE requests receive 405 Method Not Allowed +// In stateless mode, incoming GET, HEAD, and DELETE requests receive 405 Method Not Allowed // instead of being forwarded, and health checks use POST ping instead of GET. func WithStateless() Option { return func(p *TransparentProxy) { @@ -586,12 +587,25 @@ func (t *tracingTransport) RoundTrip(req *http.Request) (*http.Response, error) isMCP := strings.HasPrefix(path, "/mcp") isJSON := strings.Contains(req.Header.Get("Content-Type"), "application/json") sawInitialize := false + revision := mcp.RevisionLegacy if len(reqBody) > 0 && ((isMCP && isJSON) || t.p.transportType == types.TransportTypeStreamableHTTP.String()) { - sawInitialize = t.detectInitialize(reqBody) + method, params, id, singleRequest, isInit := t.parseRPCRequest(reqBody) + sawInitialize = isInit + if singleRequest { + meta := mcp.ExtractMeta(params) + rev, cerr := mcp.ClassifyRevision(method, meta, req.Header.Get("MCP-Protocol-Version")) + if cerr != nil { + // Malformed Modern request: reject before the backend is ever contacted. + return mcp.ClassificationErrorResponse(req, id, cerr), nil + } + revision = rev + } } + //nolint:gosec // G706: logging target URI from config + slog.Debug("classified request revision", "modern", revision == mcp.RevisionModern, "target", t.p.targetURI) // Guard: reject non-initialize requests with unknown session IDs. // When multiple proxyrunner replicas share a Redis session store, @@ -737,6 +751,14 @@ func (t *tracingTransport) RoundTrip(req *http.Request) (*http.Response, error) t.p.setServerInitialized() return resp, nil } + + // Modern (2026-07-28, stateless) has no initialize handshake, so use the + // first successful (HTTP 200) Modern request as the readiness signal: it + // lets the background health monitor begin probing this backend, which + // otherwise never starts for a pure-Modern server. See #5831. + if revision == mcp.RevisionModern && !t.p.serverInitialized() { + t.p.setServerInitialized() + } } return resp, nil @@ -779,36 +801,56 @@ func readRequestBody(req *http.Request) ([]byte, error) { return reqBody, nil } -func (t *tracingTransport) detectInitialize(body []byte) bool { - type rpcMethod struct { - Method string `json:"method"` +// parseRPCRequest parses a POST body as a JSON-RPC request in a single pass. +// +// If the body is a single JSON-RPC object, method/params/id are populated +// from it and singleRequest reports whether it is a real request (has both a +// method and a valid, non-null id) as opposed to a notification (no id) or a +// response-shaped body (no method) — either of which is not eligible for +// mcp.ClassifyRevision. sawInitialize reports whether method == "initialize". +// +// If the body is not a single JSON object (e.g. a JSON-RPC batch), method, +// params and id are zero and singleRequest is false: batches are never +// classified. sawInitialize is still computed by scanning the batch for any +// member whose method is "initialize", preserving prior behavior. +func (t *tracingTransport) parseRPCRequest( + body []byte, +) (method string, params, id json.RawMessage, singleRequest, sawInitialize bool) { + type rpcRequest struct { + Method string `json:"method"` + Params json.RawMessage `json:"params"` + ID json.RawMessage `json:"id"` } - // Single JSON-RPC object. - var single rpcMethod + var single rpcRequest if err := json.Unmarshal(body, &single); err == nil { - if single.Method == "initialize" { + sawInitialize = single.Method == "initialize" + if sawInitialize { //nolint:gosec // G706: logging target URI from config slog.Debug("detected initialize method call", "target", t.p.targetURI) - return true } - return false + singleRequest = single.Method != "" && len(single.ID) > 0 && string(single.ID) != "null" + return single.Method, single.Params, single.ID, singleRequest, sawInitialize } - // JSON-RPC batch: array of objects. Return true if any member is initialize. + // JSON-RPC batch: array of objects. Only sawInitialize is computed; a + // batch is never eligible for revision classification. + type rpcMethod struct { + Method string `json:"method"` + } var batch []rpcMethod if err := json.Unmarshal(body, &batch); err != nil { slog.Debug("failed to parse JSON-RPC body", "error", err) - return false + return "", nil, nil, false, false } for _, rpc := range batch { if rpc.Method == "initialize" { //nolint:gosec // G706: logging target URI from config slog.Debug("detected initialize method call in batch", "target", t.p.targetURI) - return true + return "", nil, nil, false, true } } - return false + return "", nil, nil, false, false } // podBackendURL constructs a backend URL that targets the specific pod IP captured @@ -1225,10 +1267,7 @@ func (p *TransparentProxy) Start(ctx context.Context) error { // 5. Catch-all proxy handler (least specific - ServeMux routing handles precedence) // Note: No manual path checking needed - ServeMux longest-match routing ensures // more specific paths registered above take precedence over this catch-all. - // In stateless mode, wrap with a method gate that rejects GET/DELETE with 405. - if p.stateless { - finalHandler = statelessMethodGate(finalHandler) - } + finalHandler = p.methodGate(finalHandler) mux.Handle("/", finalHandler) // Use ListenConfig with SO_REUSEADDR to allow port reuse after unclean shutdown @@ -1485,15 +1524,27 @@ func (*TransparentProxy) ForwardResponseToClients(_ context.Context, _ jsonrpc2. return fmt.Errorf("ForwardResponseToClients not implemented for TransparentProxy") } -// statelessMethodGate wraps a handler to reject GET, HEAD, and DELETE requests with 405. -// Used in stateless mode where the server only supports POST. -// HEAD is blocked alongside GET because HEAD is semantically a GET without a response body; -// a server that cannot handle GET will not handle HEAD either. -func statelessMethodGate(next http.Handler) http.Handler { +// methodGate wraps a handler to reject GET, HEAD, and DELETE requests with 405 +// when the proxy is running stateless (--stateless) or the request is tagged +// with the Modern (stateless-only) MCP-Protocol-Version: neither has any +// protocol use for GET/HEAD/DELETE, since Modern is unary request/response +// over POST only. +// +// This check runs outermost, before auth and any other middleware: a method +// rejection needs no identity and reveals nothing auth-gated. It is +// deliberately header-only and never reads or otherwise touches r.Body, so it +// must not call mcp.ClassifyRevision (which requires a body) or consume the +// body in any way — doing so would break the body for downstream handlers. +// +// HEAD is blocked alongside GET because HEAD is semantically a GET without a +// response body; a server that cannot handle GET will not handle HEAD either. +func (p *TransparentProxy) methodGate(next http.Handler) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - if r.Method == http.MethodGet || r.Method == http.MethodHead || r.Method == http.MethodDelete { + isGateableMethod := r.Method == http.MethodGet || r.Method == http.MethodHead || r.Method == http.MethodDelete + isModern := r.Header.Get("MCP-Protocol-Version") == mcp.MCPVersionModern + if isGateableMethod && (p.stateless || isModern) { w.Header().Set("Allow", "POST, OPTIONS") - http.Error(w, "method not allowed: server is stateless (POST only)", http.StatusMethodNotAllowed) + http.Error(w, "method not allowed: server is stateless / stateless MCP revision (POST only)", http.StatusMethodNotAllowed) return } next.ServeHTTP(w, r)