diff --git a/scripts/mock_upstream.py b/scripts/mock_upstream.py index 787761a..ec8b667 100644 --- a/scripts/mock_upstream.py +++ b/scripts/mock_upstream.py @@ -32,9 +32,50 @@ # to prevent -- past this ceiling we give up on a clean response and close. DRAIN_CEILING_BYTES = 10 * MAX_REQUEST_BYTES +# #518: the byte cap above bounds total size, not shape. A payload well +# under 1MB can still be expensive to walk -- deep nesting pushes toward +# Python's recursion limit, and a flat object with thousands of short keys +# costs little in bytes but is not free to iterate. These are judgment +# calls, not values derived from an observed cMCP tool call: no real +# `arguments` payload we ship examples for goes past 3-4 levels or a +# handful of keys, so both caps sit an order of magnitude above that, +# leaving room for a legitimate tool with a genuinely nested schema. +_MAX_ARG_DEPTH = 20 +_MAX_ARG_KEYS = 256 + logger = logging.getLogger("mock_upstream") +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 + + class MockMCPHandler(BaseHTTPRequestHandler): def log_message(self, *_args) -> None: # silence per-request access logging pass @@ -85,7 +126,7 @@ def do_POST(self) -> None: raw = self.rfile.read(length) try: - msg = json.loads(raw) + msg = json.loads(raw, parse_constant=_reject_nan_and_infinity) except (ValueError, TypeError): self._reject("parse_error", -32700, "Parse error", status=400) return @@ -95,6 +136,14 @@ def do_POST(self) -> None: return rpc_id = msg.get("id") + if "id" in msg and not _valid_rpc_id(rpc_id): + self._reject("invalid_request", -32600, "Invalid Request", status=400) + return + + if msg.get("jsonrpc") != "2.0": + self._reject("invalid_request", -32600, "Invalid Request", status=400, rpc_id=rpc_id) + return + method = msg.get("method", "") if not isinstance(method, str) or not method: self._reject("invalid_request", -32600, "Invalid Request", status=400, rpc_id=rpc_id) @@ -105,6 +154,9 @@ def do_POST(self) -> None: self._reject("invalid_request", -32600, "Invalid Request", status=400, rpc_id=rpc_id) return + self._handle_tool_call(rpc_id, params) + + def _handle_tool_call(self, rpc_id: Any, params: dict[str, Any]) -> None: tool_name = params.get("name") if not isinstance(tool_name, str) or not tool_name: self._reject( @@ -121,6 +173,14 @@ def do_POST(self) -> None: ) return + violation = _arg_shape_violation(arguments) + if violation is not None: + self._reject( + "invalid_params", -32602, f"Invalid params: {violation}", + status=400, rpc_id=rpc_id, + ) + return + text = f"mock upstream: {tool_name} called with {json.dumps(arguments, sort_keys=True)}" self._send_json( { diff --git a/tests/unit/test_mock_upstream_gate.py b/tests/unit/test_mock_upstream_gate.py index 936ba3d..98b037e 100644 --- a/tests/unit/test_mock_upstream_gate.py +++ b/tests/unit/test_mock_upstream_gate.py @@ -15,6 +15,7 @@ import threading from http.server import HTTPServer from pathlib import Path +from typing import Any import pytest @@ -286,3 +287,166 @@ def test_valid_request_is_not_logged(upstream, caplog): with caplog.at_level("WARNING", logger="mock_upstream"): _post(upstream, VALID_REQUEST) assert caplog.records == [] + + +# --------------------------------------------------------------------------- +# Strict jsonrpc/id validation (#518) +# --------------------------------------------------------------------------- + + +def test_wrong_jsonrpc_version_returns_invalid_request(upstream): + req = json.dumps( + {"jsonrpc": "1.0", "id": 20, "method": "tools/call", "params": {"name": "echo"}} + ).encode() + status, body = _post(upstream, req) + assert status == 400 + assert body["error"]["code"] == -32600 + assert body["id"] == 20 + + +def test_missing_jsonrpc_returns_invalid_request(upstream): + req = json.dumps({"id": 21, "method": "tools/call", "params": {"name": "echo"}}).encode() + status, body = _post(upstream, req) + assert status == 400 + assert body["error"]["code"] == -32600 + + +def test_object_id_returns_invalid_request_with_null_id(upstream): + req = json.dumps( + {"jsonrpc": "2.0", "id": {"nope": True}, "method": "tools/call", "params": {}} + ).encode() + status, body = _post(upstream, req) + assert status == 400 + assert body["error"]["code"] == -32600 + assert body["id"] is None + + +def test_bool_id_returns_invalid_request_with_null_id(upstream): + req = json.dumps({"jsonrpc": "2.0", "id": True, "method": "tools/call", "params": {}}).encode() + status, body = _post(upstream, req) + assert status == 400 + assert body["error"]["code"] == -32600 + assert body["id"] is None + + +def test_null_id_is_valid(upstream): + req = json.dumps( + {"jsonrpc": "2.0", "id": None, "method": "tools/call", "params": {"name": "echo"}} + ).encode() + status, body = _post(upstream, req) + assert status == 200 + assert body["id"] is None + + +def test_absent_id_is_valid(upstream): + req = json.dumps( + {"jsonrpc": "2.0", "method": "tools/call", "params": {"name": "echo"}} + ).encode() + status, body = _post(upstream, req) + assert status == 200 + assert body["id"] is None + + +# --------------------------------------------------------------------------- +# Argument depth/key-count caps (#518) +# --------------------------------------------------------------------------- + + +def _nested(depth: int) -> dict: + value: Any = "leaf" + for _ in range(depth): + value = {"child": value} + return value + + +def test_arguments_within_depth_cap_are_accepted(upstream): + req = json.dumps( + { + "jsonrpc": "2.0", + "id": 30, + "method": "tools/call", + "params": {"name": "echo", "arguments": _nested(3)}, + } + ).encode() + status, body = _post(upstream, req) + assert status == 200 + + +def test_arguments_past_depth_cap_returns_invalid_params(upstream): + req = json.dumps( + { + "jsonrpc": "2.0", + "id": 31, + "method": "tools/call", + "params": {"name": "echo", "arguments": _nested(25)}, + } + ).encode() + status, body = _post(upstream, req) + assert status == 400 + assert body["error"]["code"] == -32602 + assert body["id"] == 31 + + +def test_arguments_past_key_count_cap_returns_invalid_params(upstream): + huge_flat = {f"k{i}": i for i in range(300)} + req = json.dumps( + { + "jsonrpc": "2.0", + "id": 32, + "method": "tools/call", + "params": {"name": "echo", "arguments": huge_flat}, + } + ).encode() + status, body = _post(upstream, req) + assert status == 400 + assert body["error"]["code"] == -32602 + assert body["id"] == 32 + + +def test_arguments_within_key_count_cap_are_accepted(upstream): + small_flat = {f"k{i}": i for i in range(10)} + req = json.dumps( + { + "jsonrpc": "2.0", + "id": 33, + "method": "tools/call", + "params": {"name": "echo", "arguments": small_flat}, + } + ).encode() + status, body = _post(upstream, req) + assert status == 200 + + +# --------------------------------------------------------------------------- +# Non-standard JSON values (NaN, Infinity, -Infinity) (#518) +# --------------------------------------------------------------------------- + + +def test_nan_in_arguments_is_rejected(upstream): + body = ( + b'{"jsonrpc": "2.0", "id": 40, "method": "tools/call", ' + b'"params": {"name": "echo", "arguments": {"x": NaN}}}' + ) + status, resp = _post(upstream, body) + assert status == 400 + assert resp["error"]["code"] == -32700 + + +def test_infinity_in_arguments_is_rejected(upstream): + body = ( + b'{"jsonrpc": "2.0", "id": 41, "method": "tools/call", ' + b'"params": {"name": "echo", "arguments": {"x": Infinity}}}' + ) + status, resp = _post(upstream, body) + assert status == 400 + assert resp["error"]["code"] == -32700 + + +def test_negative_infinity_in_arguments_is_rejected(upstream): + body = ( + b'{"jsonrpc": "2.0", "id": 42, "method": "tools/call", ' + b'"params": {"name": "echo", "arguments": {"x": -Infinity}}}' + ) + status, resp = _post(upstream, body) + assert status == 400 + assert resp["error"]["code"] == -32700