From d6ecf1fc82dacbad094baadde0f5fdcf67bd0680 Mon Sep 17 00:00:00 2001 From: Juan Antonio Osorio Date: Thu, 16 Jul 2026 14:48:48 +0300 Subject: [PATCH] Add pre-dispatch denial gate to Streamable HTTP The Streamable HTTP transport (go-sdk under the shim) writes every JSON-RPC response at HTTP 200 and rejects an unknown tool with a hardcoded -32602, with no hook to signal an authorization denial. A host that enforces policy above the SDK therefore has no way to return an explicit denial for a request it blocks. Add WithCallGate: a generic, policy-free seam consulted on each POST before session validation and SDK dispatch. On a non-nil *Denial it writes the host-chosen JSON-RPC error code/message at a chosen HTTP status (403 by default), reusing the shim's mcp.JSONRPCError envelope for wire-shape parity. The body is read only on the deny path, to echo the request id best-effort. The shim carries no policy; what "denied" means is entirely the host's concern. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01UthZzKery5m7GBNcXYKALe --- mcpcompat/server/callgate.go | 145 ++++++++++++++ mcpcompat/server/callgate_test.go | 308 ++++++++++++++++++++++++++++++ mcpcompat/server/transports.go | 18 ++ 3 files changed, 471 insertions(+) create mode 100644 mcpcompat/server/callgate.go create mode 100644 mcpcompat/server/callgate_test.go diff --git a/mcpcompat/server/callgate.go b/mcpcompat/server/callgate.go new file mode 100644 index 0000000..1ed2795 --- /dev/null +++ b/mcpcompat/server/callgate.go @@ -0,0 +1,145 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "context" + "encoding/json" + "io" + "net/http" + + mcp "github.com/stacklok/toolhive-core/mcpcompat/mcp" +) + +// CallGate is consulted once per HTTP POST on the Streamable HTTP transport, +// BEFORE session-ID validation and BEFORE the message is dispatched to the +// underlying SDK. The gate decides from the request context (the host's own +// middleware has already run and populated it — identity, parsed request); +// the shim passes the request through untouched. +// +// A nil return admits the request. A non-nil *Denial short-circuits dispatch: +// the response is a JSON-RPC error envelope (the request's id echoed +// best-effort; null when unparsable) with Denial.Code/Message, written at +// Denial.HTTPStatus (403 when zero), Content-Type application/json. +// +// Ordering contract: the gate runs before session validation, so a denied +// call with an invalid/terminated session ID receives the denial (e.g. 403), +// not 404 — a denial is determinable without session state, and this matches +// hosts whose authorization middleware sits outside the SDK entirely. +// +// Mechanism only: the shim carries no policy. It never inspects the meaning of +// a Denial beyond its wire fields, and it reads context the host populated +// rather than stuffing context itself. What "denied" means is entirely the +// host's concern. +// +// Scope: the gate covers the Streamable HTTP transport only. The stdio bridge +// (request_handler.go's manual HandleMessage mirror) has no HTTP status to set +// and no per-request identity middleware — on the stdio path authorization is +// expected to deny upstream, before the bridge — and the legacy SSEServer +// transport is likewise not wired. Both are deliberately out of scope; the +// seam can generalize later if a host needs it there. +type CallGate func(ctx context.Context, r *http.Request) *Denial + +// Denial describes how a gated request is rejected on the wire. +type Denial struct { + // Code is the JSON-RPC error code placed in the response envelope. It is + // host-chosen and MUST NOT be a reserved JSON-RPC code (the -32000..-32768 + // range); it lives in application space so it never collides with SDK codes. + Code int + // Message is the JSON-RPC error message. + Message string + // HTTPStatus is the HTTP status the denial is written at. Zero means + // http.StatusForbidden (403). + HTTPStatus int +} + +// WithCallGate installs a per-call denial gate on the Streamable HTTP server. +// The gate is consulted on every POST (see CallGate for the contract). Passing +// a nil gate leaves the server ungated (identical to not calling this option). +func WithCallGate(gate CallGate) StreamableHTTPOption { + return func(s *StreamableHTTPServer) { s.callGate = gate } +} + +// denied consults the call gate for a POST request and, on a non-nil *Denial, +// writes the denial envelope and reports true (the caller must stop). It +// returns false — admitting the request — when no gate is installed, the method +// is not POST (GET/SSE and DELETE/terminate are transport lifecycle, not +// calls), or the gate returns nil. On the admit path the request body is never +// read, so the happy path carries zero overhead. +func (s *StreamableHTTPServer) denied(w http.ResponseWriter, r *http.Request) bool { + if s.callGate == nil || r.Method != http.MethodPost { + return false + } + d := s.callGate(r.Context(), r) + if d == nil { + return false + } + writeDenial(w, r, d) + return true +} + +// writeDenial emits the JSON-RPC error envelope for a denied request. It reads +// the POST body exactly once to echo the request's id best-effort (null on a +// parse failure or a batch), then writes the shim's standard mcp.JSONRPCError +// envelope at the chosen HTTP status with Content-Type application/json — so a +// denial is byte-identical in shape to every other JSON-RPC error the shim +// emits. +// +// This runs only on the deny path: on allow the body is never touched, so the +// happy path incurs zero read-and-restore overhead. Because the request is not +// forwarded once denied, there is no need to restore r.Body. +func writeDenial(w http.ResponseWriter, r *http.Request, d *Denial) { + status := d.HTTPStatus + if status == 0 { + status = http.StatusForbidden + } + + envelope := mcp.JSONRPCError{ + JSONRPC: "2.0", + ID: extractRequestID(r), + Error: mcp.NewJSONRPCErrorDetails(d.Code, d.Message, nil), + } + + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(status) + _ = json.NewEncoder(w).Encode(envelope) +} + +// extractRequestID reads the POST body once and returns the JSON-RPC id to echo +// in the error envelope. A zero-valued mcp.RequestId marshals to null, which is +// what is returned for anything that cannot be confidently attributed to a +// single request: an unreadable or unparsable body, a batch (JSON array), or a +// message with no id. The body is drained regardless so nothing is left +// half-read. +// +// The read is deliberately unbounded here: the host is expected to bound the +// request body outermost in its middleware chain (e.g. http.MaxBytesReader), so +// an over-cap body makes io.ReadAll error and falls through to a null id. A +// second limiter in this generic shim would redundantly hardcode a body-size +// policy that belongs to the host. +func extractRequestID(r *http.Request) mcp.RequestId { + var zero mcp.RequestId // marshals to null + if r.Body == nil { + return zero + } + raw, err := io.ReadAll(r.Body) + _ = r.Body.Close() + if err != nil || len(raw) == 0 { + return zero + } + var msg struct { + ID json.RawMessage `json:"id"` + } + // json.Unmarshal fails on a batch (top-level array) and on malformed JSON, + // and leaves msg.ID empty for a well-formed message with no id — all of which + // correctly fall through to the null id. + if err := json.Unmarshal(raw, &msg); err != nil || len(msg.ID) == 0 { + return zero + } + var id mcp.RequestId + if err := json.Unmarshal(msg.ID, &id); err != nil { + return zero + } + return id +} diff --git a/mcpcompat/server/callgate_test.go b/mcpcompat/server/callgate_test.go new file mode 100644 index 0000000..f931ca6 --- /dev/null +++ b/mcpcompat/server/callgate_test.go @@ -0,0 +1,308 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server_test + +import ( + "context" + "net/http" + "net/http/httptest" + "strings" + "sync" + "sync/atomic" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + mcp "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/server" +) + +// denialMsg is the stock denial message used across the gate tests. +const denialMsg = "denied" + +// addSpyTool registers a "greet" tool whose handler flips called when invoked, +// so a test can assert the handler never runs on the deny path (i.e. dispatch +// was short-circuited, not merely accompanied by a 403). +func addSpyTool(s *server.MCPServer, called *atomic.Bool) { + s.AddTool(mcp.NewTool("greet", mcp.WithDescription("greets")), + func(_ context.Context, _ mcp.CallToolRequest) (*mcp.CallToolResult, error) { + called.Store(true) + return mcp.NewToolResultText("hello"), nil + }) +} + +// addEchoTool registers a tool that echoes its "msg" string argument back in +// the result. It lets a test prove the request body reached the handler intact +// (i.e. the gate did not consume it) by asserting the argument round-trips. +func addEchoTool(s *server.MCPServer) { + s.AddTool(mcp.NewTool("echo", mcp.WithDescription("echoes msg")), + func(_ context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) { + return mcp.NewToolResultText(req.GetString("msg", "")), nil + }) +} + +// callToolBody builds a tools/call JSON-RPC request for the named tool. +func callToolBody(id, name, argsJSON string) string { + return `{"jsonrpc":"2.0","id":` + id + `,"method":"tools/call","params":{"name":"` + + name + `","arguments":` + argsJSON + `}}` +} + +// recordingGate wraps a decision function and records every (method) it was +// consulted for, so a test can assert the gate is (or is not) invoked. +type recordingGate struct { + mu sync.Mutex + methods []string + decide func(r *http.Request) *server.Denial +} + +func (g *recordingGate) gate() server.CallGate { + return func(_ context.Context, r *http.Request) *server.Denial { + g.mu.Lock() + g.methods = append(g.methods, r.Method) + g.mu.Unlock() + return g.decide(r) + } +} + +func (g *recordingGate) seen() []string { + g.mu.Lock() + defer g.mu.Unlock() + return append([]string(nil), g.methods...) +} + +// denyAll is a gate that denies every request with the given denial. +func denyAll(d *server.Denial) server.CallGate { + return func(_ context.Context, _ *http.Request) *server.Denial { return d } +} + +// doRequest issues an arbitrary HTTP request and returns the response. +func doRequest(ctx context.Context, t *testing.T, method, url, sid, body string) *http.Response { + t.Helper() + var rdr *strings.Reader + if body != "" { + rdr = strings.NewReader(body) + } else { + rdr = strings.NewReader("") + } + req, err := http.NewRequestWithContext(ctx, method, url, rdr) + require.NoError(t, err) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json, text/event-stream") + req.Header.Set("MCP-Protocol-Version", "2025-06-18") + if sid != "" { + req.Header.Set("Mcp-Session-Id", sid) + } + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + return resp +} + +// TestCallGate_DenyPath verifies a denying gate short-circuits dispatch with +// HTTP 403, a JSON content type, and a JSON-RPC error envelope that echoes the +// request id and carries the host-chosen code and message. +func TestCallGate_DenyPath(t *testing.T) { + t.Parallel() + var toolCalled atomic.Bool + srv := server.NewMCPServer("test", "1.0.0") + addSpyTool(srv, &toolCalled) + // A distinct application-space code (NOT 403) proves error.code is taken from + // the Denial, not derived from the HTTP status. + s := server.NewStreamableHTTPServer(srv, + server.WithCallGate(denyAll(&server.Denial{Code: 1001, Message: "denied by policy"}))) + ts := httptest.NewServer(s) + defer ts.Close() + + // No session established: the gate runs before session validation, so the + // denial does not depend on a live session. + resp := postRPC(t.Context(), t, ts.URL, "", callToolBody("42", "greet", "{}")) + require.Equal(t, http.StatusForbidden, resp.StatusCode) + assert.True(t, strings.HasPrefix(resp.Header.Get("Content-Type"), "application/json"), + "denial must be written as application/json") + + r := readFirstResult(t, resp) + require.NotNil(t, r.Error, "denial must carry a JSON-RPC error") + assert.Equal(t, 1001, r.Error.Code, "error.code must come from the Denial, not the HTTP status") + assert.Equal(t, "denied by policy", r.Error.Message) + assert.Equal(t, "42", string(r.ID), "the request id must be echoed") + assert.False(t, toolCalled.Load(), "tool handler must not run on the deny path") +} + +// TestCallGate_AllowPath verifies a nil-returning gate lets the request +// dispatch normally, and that the request body reaches the handler intact (the +// echoed argument round-trips, proving the gate did not consume the body). +func TestCallGate_AllowPath(t *testing.T) { + t.Parallel() + rg := &recordingGate{decide: func(_ *http.Request) *server.Denial { return nil }} + srv := server.NewMCPServer("test", "1.0.0") + addEchoTool(srv) + s := server.NewStreamableHTTPServer(srv, server.WithCallGate(rg.gate())) + ts := httptest.NewServer(s) + defer ts.Close() + + sid := initSession(t, ts.URL) + resp := postRPC(t.Context(), t, ts.URL, sid, callToolBody("7", "echo", `{"msg":"ping"}`)) + require.Equal(t, http.StatusOK, resp.StatusCode) + + r := readFirstResult(t, resp) + require.Nil(t, r.Error, "allowed call must not error") + assert.Contains(t, string(r.Result), "ping", + "echoed argument must round-trip, proving the gate left the body intact") + // The gate was consulted (initialize + notifications/initialized + this + // call are all POSTs), but only ever for POST methods. + seen := rg.seen() + require.NotEmpty(t, seen, "gate must have been consulted at least once") + for _, m := range seen { + assert.Equal(t, http.MethodPost, m, "gate must only be consulted for POST") + } +} + +// TestCallGate_NoGate is the additive-behavior regression guard: with no gate +// configured, a tools/call dispatches exactly as before. +func TestCallGate_NoGate(t *testing.T) { + t.Parallel() + srv := server.NewMCPServer("test", "1.0.0") + addEchoTool(srv) + s := server.NewStreamableHTTPServer(srv) // no WithCallGate + ts := httptest.NewServer(s) + defer ts.Close() + + sid := initSession(t, ts.URL) + resp := postRPC(t.Context(), t, ts.URL, sid, callToolBody("9", "echo", `{"msg":"pong"}`)) + require.Equal(t, http.StatusOK, resp.StatusCode) + r := readFirstResult(t, resp) + require.Nil(t, r.Error) + assert.Contains(t, string(r.Result), "pong") +} + +// TestCallGate_DenyBefore404 verifies the 403-before-404 ordering contract: a +// denied POST carrying a session ID that would otherwise 404 still gets the +// denial, both on the local-handler path (no manager) and on the cross-replica +// rehydration path (manager configured, session foreign). +func TestCallGate_DenyBefore404(t *testing.T) { + t.Parallel() + + t.Run("bogus session, no manager", func(t *testing.T) { + t.Parallel() + srv := server.NewMCPServer("test", "1.0.0") + addGreetTool(srv) + s := server.NewStreamableHTTPServer(srv, + server.WithCallGate(denyAll(&server.Denial{Code: 403, Message: denialMsg}))) + ts := httptest.NewServer(s) + defer ts.Close() + + resp := postRPC(t.Context(), t, ts.URL, "bogus-session-id", callToolBody("1", "greet", "{}")) + require.Equal(t, http.StatusForbidden, resp.StatusCode, "gate must win over the handler's 404") + r := readFirstResult(t, resp) + require.NotNil(t, r.Error) + assert.Equal(t, 403, r.Error.Code) + }) + + t.Run("foreign session, manager configured (rehydration path)", func(t *testing.T) { + t.Parallel() + mgr := newSharedSessionManager() + srv := server.NewMCPServer("test", "1.0.0") + addGreetTool(srv) + s := server.NewStreamableHTTPServer(srv, + server.WithSessionIdManager(mgr), + server.WithCallGate(denyAll(&server.Denial{Code: 403, Message: denialMsg}))) + ts := httptest.NewServer(s) + defer ts.Close() + + // "unknown-foreign" is not local and unknown to the manager, so without + // the gate serveRehydrated would 404. The gate runs first. + resp := postRPC(t.Context(), t, ts.URL, "unknown-foreign", callToolBody("2", "greet", "{}")) + require.Equal(t, http.StatusForbidden, resp.StatusCode, + "gate must win over the rehydration path's 404") + r := readFirstResult(t, resp) + require.NotNil(t, r.Error) + assert.Equal(t, 403, r.Error.Code) + }) +} + +// TestCallGate_UnparsableBodyDeny verifies the id is null when the request body +// cannot be attributed to a single request — malformed JSON, a batch, or a +// well-formed message with no id — while still emitting a 403 + envelope. +func TestCallGate_UnparsableBodyDeny(t *testing.T) { + t.Parallel() + + cases := map[string]string{ + "malformed json": `{not valid json`, + "batch array": `[{"jsonrpc":"2.0","id":1,"method":"tools/call"}]`, + "empty body": ``, + "notification, no id": `{"jsonrpc":"2.0","method":"notifications/cancelled"}`, + } + for name, body := range cases { + t.Run(name, func(t *testing.T) { + t.Parallel() + srv := server.NewMCPServer("test", "1.0.0") + s := server.NewStreamableHTTPServer(srv, + server.WithCallGate(denyAll(&server.Denial{Code: 403, Message: denialMsg}))) + ts := httptest.NewServer(s) + defer ts.Close() + + resp := postRPC(t.Context(), t, ts.URL, "", body) + require.Equal(t, http.StatusForbidden, resp.StatusCode) + r := readFirstResult(t, resp) + require.NotNil(t, r.Error) + assert.Equal(t, 403, r.Error.Code) + assert.Equal(t, "null", string(r.ID), "unattributable id must be null") + }) + } +} + +// TestCallGate_NonPOSTBypass verifies GET (SSE) and DELETE (terminate) never +// consult the gate — they are transport lifecycle, not calls. +func TestCallGate_NonPOSTBypass(t *testing.T) { + t.Parallel() + rg := &recordingGate{decide: func(_ *http.Request) *server.Denial { + return &server.Denial{Code: 403, Message: denialMsg} + }} + srv := server.NewMCPServer("test", "1.0.0") + addGreetTool(srv) + s := server.NewStreamableHTTPServer(srv, server.WithCallGate(rg.gate())) + ts := httptest.NewServer(s) + defer ts.Close() + + getResp := doRequest(t.Context(), t, http.MethodGet, ts.URL, "", "") + _ = getResp.Body.Close() + assert.NotEqual(t, http.StatusForbidden, getResp.StatusCode, + "GET must not be denied by the gate") + + delResp := doRequest(t.Context(), t, http.MethodDelete, ts.URL, "some-session", "") + _ = delResp.Body.Close() + assert.NotEqual(t, http.StatusForbidden, delResp.StatusCode, + "DELETE must not be denied by the gate") + + assert.Empty(t, rg.seen(), "gate must never be consulted for non-POST methods") +} + +// TestCallGate_HTTPStatus verifies HTTPStatus defaults to 403 when zero and is +// honored when set explicitly. +func TestCallGate_HTTPStatus(t *testing.T) { + t.Parallel() + cases := []struct { + name string + denial *server.Denial + wantStatus int + }{ + {"zero defaults to 403", &server.Denial{Code: 403, Message: denialMsg}, http.StatusForbidden}, + {"explicit status honored", &server.Denial{Code: 403, Message: denialMsg, HTTPStatus: http.StatusTooManyRequests}, http.StatusTooManyRequests}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + srv := server.NewMCPServer("test", "1.0.0") + s := server.NewStreamableHTTPServer(srv, server.WithCallGate(denyAll(tc.denial))) + ts := httptest.NewServer(s) + defer ts.Close() + + resp := postRPC(t.Context(), t, ts.URL, "", callToolBody("1", "greet", "{}")) + require.Equal(t, tc.wantStatus, resp.StatusCode) + r := readFirstResult(t, resp) + require.NotNil(t, r.Error) + assert.Equal(t, 403, r.Error.Code, "JSON-RPC code is independent of HTTP status") + }) + } +} diff --git a/mcpcompat/server/transports.go b/mcpcompat/server/transports.go index a18dcac..cd276c8 100644 --- a/mcpcompat/server/transports.go +++ b/mcpcompat/server/transports.go @@ -83,6 +83,7 @@ type StreamableHTTPServer struct { endpointPath string contextFunc HTTPContextFunc sessionIDMgr SessionIdManager + callGate CallGate heartbeat time.Duration // disableLocalhostProtection turns off go-sdk's DNS-rebinding/localhost // protection (which 403s requests on a loopback listener with a non-localhost @@ -267,6 +268,8 @@ func (s *StreamableHTTPServer) build() { } // ServeHTTP implements http.Handler. +// +//nolint:gocyclo // pre-dispatch gate adds one branch to an already at-limit dispatch function. func (s *StreamableHTTPServer) ServeHTTP(w http.ResponseWriter, r *http.Request) { s.build() if s.buildErr != nil { @@ -277,6 +280,21 @@ func (s *StreamableHTTPServer) ServeHTTP(w http.ResponseWriter, r *http.Request) if s.contextFunc != nil { r = r.WithContext(s.contextFunc(r.Context(), r)) } + // Pre-dispatch denial gate. Placement is deliberate: + // (a) AFTER contextFunc, so a gate that reads context injected via the + // shim option (identity, parsed request) sees it; + // (b) BEFORE the nonce bridge, so a denied request never registers a + // pending per-request context entry; + // (c) BEFORE both dispatch branches — the local go-sdk handler AND the + // cross-replica serveRehydrated path — so multi-replica deployments + // are gated identically; + // (d) BEFORE session-ID validation (403-before-404): a denial is + // determinable without session state, so a denied call with a stale, + // foreign, or terminated session ID receives the denial rather than a + // 404, matching hosts whose authorization sits outside the SDK. + if s.denied(w, r) { + return + } // Bridge per-request context values into the handler. go-sdk does not // propagate the per-POST request context into handlers for existing // sessions (it handles messages on the session's connection goroutine using