From 2fe2f8aeacfbd7bd3bed0af6317b3caecd11612f Mon Sep 17 00:00:00 2001 From: Yatsuiii Date: Tue, 25 Aug 2026 00:03:05 +0530 Subject: [PATCH 1/2] fix: close #518's gateway-side gap in _handle_mcp qubeena07's review on #556 flagged that _handle_mcp in mcp/server.py has the same gap #556 closed in the mock upstream: no strict jsonrpc/id validation, no argument depth/key-count cap, and NaN/Infinity/-Infinity pass through json.loads silently. This applies the same three checks here, matching scripts/mock_upstream.py's behavior and _MAX_ARG_DEPTH / _MAX_ARG_KEYS values, kept in sync rather than unified per qubeena07's explicit request not to merge the two validation paths in one PR. Also extracts _parse_mcp_envelope (body size/parse) and _deny_response (policy-deny error mapping) out of _handle_mcp and _handle_tool_call, which were pre-existing at or over the repo's complexity threshold and would have gone further over it with the new checks inline. Both extractions are pure code motion, no behavior change. 14 new tests (6 jsonrpc/id, 4 depth/key-count, 3 NaN/Infinity, matching #556's coverage) plus all 64 existing tests across the affected test files pass. --- src/cmcp_runtime/mcp/server.py | 218 ++++++++++++++++++++--------- tests/unit/test_mcp_server_auth.py | 152 +++++++++++++++++++- 2 files changed, 299 insertions(+), 71 deletions(-) diff --git a/src/cmcp_runtime/mcp/server.py b/src/cmcp_runtime/mcp/server.py index a0c4d8bc..bb11beda 100644 --- a/src/cmcp_runtime/mcp/server.py +++ b/src/cmcp_runtime/mcp/server.py @@ -41,6 +41,47 @@ # Endpoints exempt from bearer-token auth (Kubernetes liveness / readiness probes) _AUTH_EXEMPT_PATHS = {"/health", "/readyz"} +# #518: DOS-001's byte cap bounds total size, not shape. A payload well under +# the limit can still push toward Python's recursion limit through deep +# nesting, or cost real time to iterate through a flat object with thousands +# of short keys. These values mirror `_MAX_ARG_DEPTH` / `_MAX_ARG_KEYS` in +# scripts/mock_upstream.py - same judgment call, same margin above any real +# `arguments` payload this repo ships examples for, kept in sync rather than +# unified per qubeena07's review on #518. +_MAX_ARG_DEPTH = 20 +_MAX_ARG_KEYS = 256 + + +def _reject_nan_and_infinity(text: str) -> float: + raise ValueError(f"non-standard JSON value not allowed: {text}") + + +def _valid_rpc_id(value: Any) -> bool: + # JSON-RPC 2.0 id must be a string, number, or null -- not a bool, even + # though bool is an int subclass in Python. + return value is None or (isinstance(value, (str, int, float)) and not isinstance(value, bool)) + + +def _arg_shape_violation(value: Any, *, depth: int = 0) -> str | None: + """Return a message describing the first depth/key-count violation found + under `value`, or None if it fits within `_MAX_ARG_DEPTH` / `_MAX_ARG_KEYS`.""" + if depth > _MAX_ARG_DEPTH: + return f"arguments nested past the depth cap of {_MAX_ARG_DEPTH}" + if isinstance(value, dict): + if len(value) > _MAX_ARG_KEYS: + return f"object has more than {_MAX_ARG_KEYS} keys" + for child in value.values(): + violation = _arg_shape_violation(child, depth=depth + 1) + if violation is not None: + return violation + elif isinstance(value, list): + for child in value: + violation = _arg_shape_violation(child, depth=depth + 1) + if violation is not None: + return violation + return None + + # Revisions the gateway can negotiate at `initialize`, newest first. # # `initialize` belongs to the handshake era only. `PROTOCOL_VERSION` (#509) is @@ -271,8 +312,12 @@ def __init__( exception_handlers={Exception: _unhandled_error_handler}, ) - async def _handle_mcp(self, request: Request) -> Response: - """Handle MCP JSON-RPC 2.0 calls.""" + async def _parse_mcp_envelope(self, request: Request) -> dict[str, Any] | Response: + """Read, size-check, and parse the request body. + + Returns the parsed JSON-RPC message dict, or an error Response if the + body is oversized, unparsable, or not a JSON object. + """ # DOS-001: reject oversized requests before parsing to prevent OOM content_length = request.headers.get("content-length") if content_length: @@ -300,8 +345,10 @@ async def _handle_mcp(self, request: Request) -> Response: }, status_code=413, ) - msg = json.loads(body) - except (json.JSONDecodeError, UnicodeDecodeError) as exc: + # #518: parse_constant intercepts NaN/Infinity/-Infinity before + # json.loads would otherwise accept them silently. + msg = json.loads(body, parse_constant=_reject_nan_and_infinity) + except (json.JSONDecodeError, UnicodeDecodeError, ValueError) as exc: import hashlib payload_hash = f"sha256:{hashlib.sha256(body).hexdigest()}" logger.warning( @@ -325,8 +372,22 @@ async def _handle_mcp(self, request: Request) -> Response: if not isinstance(msg, dict): return _invalid_request() + return msg + async def _handle_mcp(self, request: Request) -> Response: + """Handle MCP JSON-RPC 2.0 calls.""" + parsed = await self._parse_mcp_envelope(request) + if isinstance(parsed, Response): + return parsed + msg = parsed + + # #518: strict jsonrpc/id validation, matching scripts/mock_upstream.py. rpc_id = msg.get("id") + if "id" in msg and not _valid_rpc_id(rpc_id): + return _invalid_request() + if msg.get("jsonrpc") != "2.0": + return _invalid_request(rpc_id) + method = msg.get("method", "") if not isinstance(method, str): return _invalid_request(rpc_id) @@ -369,12 +430,94 @@ async def _handle_mcp(self, request: Request) -> Response: status_code=404, ) + def _deny_response(self, rpc_id: Any, call_id: str, result: Any) -> JSONResponse: + """Build the JSON-RPC error response for a policy-denied tool call.""" + deny_reason = result.deny_reason or "" + # Upstream transport/tool failure is a 502, not a policy deny. + if deny_reason.startswith("upstream_error:"): + return JSONResponse( + { + "jsonrpc": "2.0", + "error": { + "code": -32000, + "message": "Upstream MCP server error", + "data": { + "error_code": deny_reason.removeprefix("upstream_error:"), + "call_id": call_id, + }, + }, + "id": rpc_id, + }, + status_code=502, + ) + _HEALTH_REASONS = {"attestation_stale", "catalog_drift"} + if result.deny_reason in _HEALTH_REASONS: + return JSONResponse( + { + "jsonrpc": "2.0", + "error": { + "code": -32000, + "message": result.deny_reason, + "data": { + "error_code": result.deny_reason.upper(), + "call_id": call_id, + }, + }, + "id": rpc_id, + }, + status_code=503, + ) + # INJECT-003: log deny_reason internally; do not reflect internal detail to caller + error_code = ( + "TOOL_NOT_IN_CATALOG" + if "catalog" in (result.deny_reason or "") + else "POLICY_DENY" + ) + logger.info( + "POLICY_DENY: call_id=%s error_code=%s reason=%s", + call_id, error_code, result.deny_reason, + ) + error_data: dict[str, Any] = { + "error_code": error_code, + "call_id": call_id, + } + # Advice annotations come from the hash-pinned policy bundle + # (operator-authored, not caller input), so reflecting them does + # not violate INJECT-003. They carry e.g. HITL escalation payloads. + if result.advice: + error_data["advice"] = result.advice + return JSONResponse( + { + "jsonrpc": "2.0", + "error": { + "code": -32000, + "message": "Request denied by policy", + "data": error_data, + }, + "id": rpc_id, + }, + status_code=403, + ) + async def _handle_tool_call(self, rpc_id: Any, params: dict[str, Any]) -> Response: """Route a tools/call request through the proxy.""" # POLICY-002: canonicalize tool names at ingress so Cedar policy, catalog, and # request all use the same case - prevents case-variant bypass of deny rules. tool_name: str = params.get("name", "").lower() arguments: dict[str, Any] = params.get("arguments", {}) + + # #518: depth/key-count cap, matching scripts/mock_upstream.py. + violation = _arg_shape_violation(arguments) + if violation is not None: + return JSONResponse( + { + "jsonrpc": "2.0", + "error": {"code": -32602, "message": f"Invalid params: {violation}"}, + "id": rpc_id, + }, + status_code=400, + ) + call_id = str(uuid.uuid4()) # A malformed _cmcp (string, list, number) must not 500 the call path. cmcp_params = params.get("_cmcp") @@ -412,72 +555,7 @@ async def _handle_tool_call(self, rpc_id: Any, params: dict[str, Any]) -> Respon ) if not result.allowed: - deny_reason = result.deny_reason or "" - # Upstream transport/tool failure is a 502, not a policy deny. - if deny_reason.startswith("upstream_error:"): - return JSONResponse( - { - "jsonrpc": "2.0", - "error": { - "code": -32000, - "message": "Upstream MCP server error", - "data": { - "error_code": deny_reason.removeprefix("upstream_error:"), - "call_id": call_id, - }, - }, - "id": rpc_id, - }, - status_code=502, - ) - _HEALTH_REASONS = {"attestation_stale", "catalog_drift"} - if result.deny_reason in _HEALTH_REASONS: - return JSONResponse( - { - "jsonrpc": "2.0", - "error": { - "code": -32000, - "message": result.deny_reason, - "data": { - "error_code": result.deny_reason.upper(), - "call_id": call_id, - }, - }, - "id": rpc_id, - }, - status_code=503, - ) - # INJECT-003: log deny_reason internally; do not reflect internal detail to caller - error_code = ( - "TOOL_NOT_IN_CATALOG" - if "catalog" in (result.deny_reason or "") - else "POLICY_DENY" - ) - logger.info( - "POLICY_DENY: call_id=%s error_code=%s reason=%s", - call_id, error_code, result.deny_reason, - ) - error_data: dict[str, Any] = { - "error_code": error_code, - "call_id": call_id, - } - # Advice annotations come from the hash-pinned policy bundle - # (operator-authored, not caller input), so reflecting them does - # not violate INJECT-003. They carry e.g. HITL escalation payloads. - if result.advice: - error_data["advice"] = result.advice - return JSONResponse( - { - "jsonrpc": "2.0", - "error": { - "code": -32000, - "message": "Request denied by policy", - "data": error_data, - }, - "id": rpc_id, - }, - status_code=403, - ) + return self._deny_response(rpc_id, call_id, result) cmcp_meta: dict[str, Any] = { "call_id": call_id, diff --git a/tests/unit/test_mcp_server_auth.py b/tests/unit/test_mcp_server_auth.py index 16fd8d4e..0ab96ed5 100644 --- a/tests/unit/test_mcp_server_auth.py +++ b/tests/unit/test_mcp_server_auth.py @@ -2,12 +2,13 @@ from __future__ import annotations +from typing import Any from unittest.mock import AsyncMock, MagicMock, patch import pytest from starlette.testclient import TestClient -from cmcp_runtime.mcp.server import MCPServer +from cmcp_runtime.mcp.server import _MAX_ARG_DEPTH, _MAX_ARG_KEYS, MCPServer def _make_server(bearer_token: str | None = None) -> MCPServer: @@ -471,3 +472,152 @@ def test_deny_response_does_not_include_internal_reason(): assert "Cedar eval error" not in str(body) assert "AttributeAccessError" not in str(body) assert body["error"]["message"] == "Request denied by policy" + + +# ── #518: strict jsonrpc/id validation, matching scripts/mock_upstream.py ──── + +def test_missing_jsonrpc_field_returns_invalid_request(): + server = _make_server() + client = TestClient(server.app, raise_server_exceptions=False) + resp = client.post("/mcp", json={"method": "initialize", "id": 1}) + assert resp.status_code == 400 + assert resp.json()["error"]["code"] == -32600 + + +def test_wrong_jsonrpc_version_returns_invalid_request(): + server = _make_server() + client = TestClient(server.app, raise_server_exceptions=False) + resp = client.post("/mcp", json={"jsonrpc": "1.0", "method": "initialize", "id": 1}) + assert resp.status_code == 400 + assert resp.json()["error"]["code"] == -32600 + + +def test_jsonrpc_version_as_number_is_rejected(): + """jsonrpc must equal the string "2.0" exactly, not the numeric value 2.0.""" + server = _make_server() + client = TestClient(server.app, raise_server_exceptions=False) + resp = client.post("/mcp", json={"jsonrpc": 2.0, "method": "initialize", "id": 1}) + assert resp.status_code == 400 + assert resp.json()["error"]["code"] == -32600 + + +def test_bool_id_returns_invalid_request(): + """id must be a string, number, or null -- bool is an int subclass but not valid.""" + server = _make_server() + client = TestClient(server.app, raise_server_exceptions=False) + resp = client.post("/mcp", json={"jsonrpc": "2.0", "method": "initialize", "id": True}) + assert resp.status_code == 400 + assert resp.json()["error"]["code"] == -32600 + + +def test_object_id_returns_invalid_request(): + server = _make_server() + client = TestClient(server.app, raise_server_exceptions=False) + resp = client.post("/mcp", json={"jsonrpc": "2.0", "method": "initialize", "id": {}}) + assert resp.status_code == 400 + assert resp.json()["error"]["code"] == -32600 + + +def test_null_id_is_valid_notification(): + server = _make_server() + client = TestClient(server.app, raise_server_exceptions=False) + resp = client.post("/mcp", json={"jsonrpc": "2.0", "method": "initialize", "id": None}) + assert resp.status_code == 200 + + +def test_string_id_is_valid(): + server = _make_server() + client = TestClient(server.app, raise_server_exceptions=False) + resp = client.post("/mcp", json={"jsonrpc": "2.0", "method": "initialize", "id": "req-1"}) + assert resp.status_code == 200 + assert resp.json()["id"] == "req-1" + + +# ── #518: argument depth/key-count caps, matching scripts/mock_upstream.py ── + +def _nested(depth: int) -> Any: + value: Any = "leaf" + for _ in range(depth): + value = {"a": value} + return value + + +def test_arguments_within_depth_cap_is_allowed(): + server = _make_server() + client = TestClient(server.app, raise_server_exceptions=False) + resp = client.post( + "/mcp", + json={ + "jsonrpc": "2.0", + "method": "tools/call", + "params": {"name": "t", "arguments": {"nested": _nested(_MAX_ARG_DEPTH - 1)}}, + "id": 1, + }, + ) + assert resp.status_code == 200 + + +def test_arguments_exceeding_depth_cap_returns_invalid_params(): + server = _make_server() + client = TestClient(server.app, raise_server_exceptions=False) + resp = client.post( + "/mcp", + json={ + "jsonrpc": "2.0", + "method": "tools/call", + "params": {"name": "t", "arguments": {"nested": _nested(_MAX_ARG_DEPTH + 10)}}, + "id": 1, + }, + ) + assert resp.status_code == 400 + assert resp.json()["error"]["code"] == -32602 + + +def test_arguments_within_key_cap_is_allowed(): + server = _make_server() + client = TestClient(server.app, raise_server_exceptions=False) + arguments = {f"k{i}": i for i in range(_MAX_ARG_KEYS)} + resp = client.post( + "/mcp", + json={ + "jsonrpc": "2.0", + "method": "tools/call", + "params": {"name": "t", "arguments": arguments}, + "id": 1, + }, + ) + assert resp.status_code == 200 + + +def test_arguments_exceeding_key_cap_returns_invalid_params(): + server = _make_server() + client = TestClient(server.app, raise_server_exceptions=False) + arguments = {f"k{i}": i for i in range(_MAX_ARG_KEYS + 1)} + resp = client.post( + "/mcp", + json={ + "jsonrpc": "2.0", + "method": "tools/call", + "params": {"name": "t", "arguments": arguments}, + "id": 1, + }, + ) + assert resp.status_code == 400 + assert resp.json()["error"]["code"] == -32602 + + +# ── #518: non-standard JSON values (NaN, Infinity, -Infinity) ─────────────── + +@pytest.mark.parametrize("literal", ["NaN", "Infinity", "-Infinity"]) +def test_non_standard_json_constant_returns_parse_error(literal): + server = _make_server() + client = TestClient(server.app, raise_server_exceptions=False) + body = ( + '{"jsonrpc": "2.0", "method": "tools/call", ' + f'"params": {{"name": "t", "arguments": {{"x": {literal}}}}}, "id": 1}}' + ) + resp = client.post( + "/mcp", content=body.encode(), headers={"Content-Type": "application/json"} + ) + assert resp.status_code == 400 + assert resp.json()["error"]["code"] == -32700 From 68be9662f9e5cbfbbb8218b442a0cc2f491d6b27 Mon Sep 17 00:00:00 2001 From: Yatsuiii Date: Tue, 25 Aug 2026 00:21:09 +0530 Subject: [PATCH 2/2] test: cover the deny-response and list-recursion branches Codecov flagged on #561 Codecov's patch coverage check on PR #561 flagged 7 lines as untested in src/cmcp_runtime/mcp/server.py. All 7 are pre-existing behavior that got attributed to this PR's diff because the code moved during the _deny_response extraction, plus one branch in the new _arg_shape_violation that no existing test exercised. Adds tests for the upstream_error 502 branch, the attestation_stale and catalog_drift 503 branches, the advice-included branch in the deny response, and a depth-cap violation nested inside a list rather than a dict, since _arg_shape_violation recurses into lists too. No production code changed. All 7 previously-missing lines now show as covered; 49 tests pass, ruff clean. --- tests/unit/test_mcp_server_auth.py | 93 ++++++++++++++++++++++++++++++ 1 file changed, 93 insertions(+) diff --git a/tests/unit/test_mcp_server_auth.py b/tests/unit/test_mcp_server_auth.py index 0ab96ed5..a2166d67 100644 --- a/tests/unit/test_mcp_server_auth.py +++ b/tests/unit/test_mcp_server_auth.py @@ -474,6 +474,82 @@ def test_deny_response_does_not_include_internal_reason(): assert body["error"]["message"] == "Request denied by policy" +def test_upstream_error_deny_reason_returns_502(): + proxy = MagicMock() + proxy._catalog = MagicMock() + proxy._catalog.entries = {} + proxy.call_tool = AsyncMock(return_value=MagicMock( + allowed=False, + deny_reason="upstream_error:CONNECTION_REFUSED", + audit_entry_hash=None, + would_have_denied=False, + latency_us=0, + advice=None, + )) + with patch("cmcp_runtime.mcp.server.StatelessKernel"): + server = MCPServer(proxy) + client = TestClient(server.app, raise_server_exceptions=False) + resp = client.post( + "/mcp", + json={"jsonrpc": "2.0", "method": "tools/call", "params": {"name": "t", "arguments": {}}, "id": 1}, + ) + assert resp.status_code == 502 + body = resp.json() + assert body["error"]["message"] == "Upstream MCP server error" + assert body["error"]["data"]["error_code"] == "CONNECTION_REFUSED" + + +@pytest.mark.parametrize("reason", ["attestation_stale", "catalog_drift"]) +def test_health_deny_reason_returns_503(reason): + proxy = MagicMock() + proxy._catalog = MagicMock() + proxy._catalog.entries = {} + proxy.call_tool = AsyncMock(return_value=MagicMock( + allowed=False, + deny_reason=reason, + audit_entry_hash=None, + would_have_denied=False, + latency_us=0, + advice=None, + )) + with patch("cmcp_runtime.mcp.server.StatelessKernel"): + server = MCPServer(proxy) + client = TestClient(server.app, raise_server_exceptions=False) + resp = client.post( + "/mcp", + json={"jsonrpc": "2.0", "method": "tools/call", "params": {"name": "t", "arguments": {}}, "id": 1}, + ) + assert resp.status_code == 503 + body = resp.json() + assert body["error"]["message"] == reason + assert body["error"]["data"]["error_code"] == reason.upper() + + +def test_deny_response_includes_advice_when_present(): + """Advice annotations come from the operator-authored policy bundle, so + reflecting them (unlike deny_reason) does not violate INJECT-003.""" + proxy = MagicMock() + proxy._catalog = MagicMock() + proxy._catalog.entries = {} + proxy.call_tool = AsyncMock(return_value=MagicMock( + allowed=False, + deny_reason="policy denied", + audit_entry_hash=None, + would_have_denied=False, + latency_us=0, + advice={"escalation": "HITL"}, + )) + with patch("cmcp_runtime.mcp.server.StatelessKernel"): + server = MCPServer(proxy) + client = TestClient(server.app, raise_server_exceptions=False) + resp = client.post( + "/mcp", + json={"jsonrpc": "2.0", "method": "tools/call", "params": {"name": "t", "arguments": {}}, "id": 1}, + ) + assert resp.status_code == 403 + assert resp.json()["error"]["data"]["advice"] == {"escalation": "HITL"} + + # ── #518: strict jsonrpc/id validation, matching scripts/mock_upstream.py ──── def test_missing_jsonrpc_field_returns_invalid_request(): @@ -573,6 +649,23 @@ def test_arguments_exceeding_depth_cap_returns_invalid_params(): assert resp.json()["error"]["code"] == -32602 +def test_depth_cap_violation_inside_a_list_is_caught(): + """The depth walk recurses into list items, not only dict values.""" + server = _make_server() + client = TestClient(server.app, raise_server_exceptions=False) + resp = client.post( + "/mcp", + json={ + "jsonrpc": "2.0", + "method": "tools/call", + "params": {"name": "t", "arguments": {"items": [_nested(_MAX_ARG_DEPTH + 10)]}}, + "id": 1, + }, + ) + assert resp.status_code == 400 + assert resp.json()["error"]["code"] == -32602 + + def test_arguments_within_key_cap_is_allowed(): server = _make_server() client = TestClient(server.app, raise_server_exceptions=False)