fix: close the remaining #518 gaps in the mock upstream's JSON-RPC gate

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(
             {

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