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..a2166d67 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,245 @@ 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" + + +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(): + 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_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) + 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