security: harden MCP server against 6 pentest findings

Four-model adversarial pentest (Codex, Sonnet, Claude, Qwen) found
input validation, policy bypass, and replay race issues in the MCP
server. This commit addresses all 6 findings plus 3 additional issues
caught during Sonnet's post-fix review.

Fixes:
- Add pre-dispatch _validate_arguments() with type checks and size
  limits (1MB plaintext, 100KB keys) to prevent server crashes from
  malformed MCP arguments (envelope=[], plaintext=None, name=42)
- Broaden exception handling: catch TypeError, AttributeError, KeyError,
  RuntimeError, OSError in hybrid dispatch; catch RuntimeError/OSError
  in import guards for graceful degradation
- Close PQC_REQUIRE_KEY_HANDLES bypass: key_store_save now rejects
  key_data containing secret_key fields when policy is active
- Add 100KB size limit on key_data in key_store_save
- Fix replay cache TOCTOU: replace split check()/mark() with atomic
  check_and_mark() after verify+decrypt
- Reject bool values for integer fields (bool is subclass of int)

25 new tests covering all fixes. 301 total tests passing.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: scottdhughes

pqc_mcp_server/__init__.py

@@ -29,15 +29,15 @@
     import oqs  # noqa: F401
 
     HAS_LIBOQS = True
-except ImportError:
+except (ImportError, RuntimeError, OSError):
     HAS_LIBOQS = False
 
 try:
     from cryptography.exceptions import InvalidTag
     from pqc_mcp_server.hybrid import SenderVerificationError
 
     HAS_HYBRID = True
-except ImportError:
+except (ImportError, RuntimeError, OSError):
     HAS_HYBRID = False
 
     class InvalidTag(Exception):  # type: ignore[no-redef]
@@ -119,6 +119,99 @@ class SenderVerificationError(Exception):  # type: ignore[no-redef]
         "pqc_key_store_delete": handle_key_store_delete,
     }
 
+# --- Pre-dispatch input validation ---
+
+# Size limits (bytes/chars)
+_MAX_PLAINTEXT_SIZE = 1_048_576  # 1 MB
+_MAX_MESSAGE_SIZE = 1_048_576  # 1 MB
+_MAX_KEY_FIELD_SIZE = 102_400  # 100 KB
+
+# Expected types by field name
+_STRING_FIELDS = frozenset(
+    {
+        "algorithm",
+        "public_key",
+        "secret_key",
+        "ciphertext",
+        "message",
+        "signature",
+        "plaintext",
+        "plaintext_base64",
+        "name",
+        "store_as",
+        "key_store_name",
+        "classical_public_key",
+        "pqc_public_key",
+        "classical_secret_key",
+        "pqc_secret_key",
+        "x25519_ephemeral_public_key",
+        "pqc_ciphertext",
+        "sender_secret_key",
+        "sender_public_key",
+        "expected_sender_public_key",
+        "expected_sender_fingerprint",
+        "sender_key_store_name",
+        "recipient_key_store_name",
+        "type",
+    }
+)
+_DICT_FIELDS = frozenset({"envelope", "key_data"})
+_BOOL_FIELDS = frozenset({"overwrite", "include_secret_key"})
+_INT_FIELDS = frozenset({"iterations", "max_age_seconds"})
+
+_SIZE_LIMITS: dict[str, int] = {
+    "plaintext": _MAX_PLAINTEXT_SIZE,
+    "plaintext_base64": _MAX_PLAINTEXT_SIZE,
+    "message": _MAX_MESSAGE_SIZE,
+    "public_key": _MAX_KEY_FIELD_SIZE,
+    "secret_key": _MAX_KEY_FIELD_SIZE,
+    "classical_public_key": _MAX_KEY_FIELD_SIZE,
+    "pqc_public_key": _MAX_KEY_FIELD_SIZE,
+    "classical_secret_key": _MAX_KEY_FIELD_SIZE,
+    "pqc_secret_key": _MAX_KEY_FIELD_SIZE,
+    "sender_secret_key": _MAX_KEY_FIELD_SIZE,
+    "sender_public_key": _MAX_KEY_FIELD_SIZE,
+    "ciphertext": _MAX_PLAINTEXT_SIZE,
+    "x25519_ephemeral_public_key": _MAX_KEY_FIELD_SIZE,
+    "pqc_ciphertext": _MAX_KEY_FIELD_SIZE,
+    "expected_sender_public_key": _MAX_KEY_FIELD_SIZE,
+    "signature": _MAX_PLAINTEXT_SIZE,
+}
+
+
+def _validate_arguments(arguments: dict[str, Any]) -> None:
+    """Validate argument types and sizes before handler dispatch.
+
+    Raises ValueError on type mismatch or size violation.
+    Unknown fields pass through silently (future-proof).
+    """
+    if not isinstance(arguments, dict):
+        raise ValueError("arguments must be a JSON object")
+
+    for key, value in arguments.items():
+        if value is None:
+            continue  # Optional fields may be absent/null
+
+        if key in _STRING_FIELDS:
+            if not isinstance(value, str):
+                raise ValueError(f"Parameter '{key}' must be a string, got {type(value).__name__}")
+            limit = _SIZE_LIMITS.get(key)
+            if limit is not None and len(value) > limit:
+                raise ValueError(f"Parameter '{key}' exceeds size limit ({limit} chars)")
+        elif key in _DICT_FIELDS:
+            if not isinstance(value, dict):
+                raise ValueError(
+                    f"Parameter '{key}' must be a JSON object, got {type(value).__name__}"
+                )
+        elif key in _BOOL_FIELDS:
+            if not isinstance(value, bool):
+                raise ValueError(f"Parameter '{key}' must be a boolean, got {type(value).__name__}")
+        elif key in _INT_FIELDS:
+            # bool is a subclass of int in Python — reject it explicitly
+            if isinstance(value, bool) or not isinstance(value, (int, float)):
+                raise ValueError(f"Parameter '{key}' must be a number, got {type(value).__name__}")
+
+
 # --- MCP Server ---
 
 server = Server("pqc-mcp-server")
@@ -135,6 +228,12 @@ async def list_tools() -> list[Tool]:
 
 @server.call_tool()
 async def call_tool(name: str, arguments: dict[str, Any]) -> list[TextContent]:
+    # Pre-dispatch input validation (type checks + size limits)
+    try:
+        _validate_arguments(arguments)
+    except (ValueError, TypeError) as e:
+        return _json_response({"error": str(e)})
+
     if not HAS_LIBOQS:
         return _json_response(
             {
@@ -172,6 +271,12 @@ async def call_tool(name: str, arguments: dict[str, Any]) -> list[TextContent]:
             return _json_response(
                 {"error": "Decryption failed: ciphertext, key, or envelope metadata is invalid"}
             )
+        except (TypeError, AttributeError, KeyError) as e:
+            return _json_response({"error": f"Invalid argument type: {e}"})
+        except RuntimeError as e:
+            return _json_response({"error": f"Internal error: {e}"})
+        except OSError as e:
+            return _json_response({"error": f"I/O error: {e}"})
 
     return _json_response({"error": f"Unknown tool: {name}"})
 

pqc_mcp_server/handlers_hybrid.py

@@ -202,16 +202,16 @@ def handle_hybrid_auth_open(arguments: dict[str, Any]) -> dict[str, Any]:
 
     _validate_envelope_size(envelope)
 
-    # Replay dedup: check BEFORE verification, mark AFTER success.
-    # CONCURRENCY NOTE: Two concurrent opens of the same valid envelope can
-    # both pass check() before either calls mark(). This is a documented
-    # tradeoff for single-process async — avoids junk pre-blocking at the
-    # cost of not guaranteeing one-time acceptance under concurrency.
-    # For multi-process, add a lock or two-phase reserve/finalize pattern.
+    # Replay dedup: compute digest early (cheap SHA3-256), but check_and_mark
+    # runs AFTER verify+decrypt. This prevents junk pre-blocking (invalid
+    # envelopes never enter cache) while closing the TOCTOU window between
+    # the old separate check()/mark() calls. Safety guarantee: no `await`
+    # exists in this handler's call path (all crypto is synchronous liboqs),
+    # so the event loop cannot preempt between digest computation and
+    # check_and_mark. If an async operation is ever added here, this
+    # guarantee must be re-evaluated.
     cache = get_replay_cache()
     digest = signature_digest(envelope)
-    if cache.check(digest):
-        raise ValueError("Duplicate envelope (replay detected)")
 
     classical_sk, pqc_sk = _resolve_hybrid_secret(arguments)
     expected_pk = (
@@ -240,8 +240,12 @@ def handle_hybrid_auth_open(arguments: dict[str, Any]) -> dict[str, Any]:
         kwargs["max_age_seconds"] = max_age_int
     result = hybrid_auth_open(envelope, classical_sk, pqc_sk, **kwargs)
 
-    # Mark AFTER successful verify+decrypt — only verified envelopes enter cache
-    cache.mark(digest)
+    # Atomic check+mark AFTER successful verify+decrypt. Only verified
+    # envelopes enter the cache. Closes the TOCTOU window: under concurrent
+    # async, only the first coroutine to reach check_and_mark succeeds.
+    if cache.check_and_mark(digest):
+        raise ValueError("Duplicate envelope (replay detected)")
+
     return result
 
 

pqc_mcp_server/key_store.py

@@ -7,8 +7,11 @@
 
 import base64
 import hashlib
+import json as _json
 from typing import Any
 
+_MAX_KEY_DATA_SIZE = 102_400  # 100 KB
+
 # Module-level store (lives for the server's lifetime)
 _STORE: dict[str, dict[str, Any]] = {}
 
@@ -87,6 +90,27 @@ def _require_flat_kem(keys: dict[str, Any], name: str) -> None:
         )
 
 
+def _reject_secret_fields(key_data: dict[str, Any]) -> None:
+    """Reject key_data containing secret_key fields when handle-only policy is active.
+
+    Checks flat key structures and hybrid bundle sub-dicts.
+    """
+    if "secret_key" in key_data:
+        raise ValueError(
+            "key_store_save rejected: key_data contains 'secret_key'. "
+            "PQC_REQUIRE_KEY_HANDLES policy prohibits importing raw secrets. "
+            "Use pqc_generate_keypair with store_as or pqc_hybrid_keygen with store_as."
+        )
+    for component in ("classical", "pqc"):
+        sub = key_data.get(component)
+        if isinstance(sub, dict) and "secret_key" in sub:
+            raise ValueError(
+                f"key_store_save rejected: key_data['{component}'] contains 'secret_key'. "
+                "PQC_REQUIRE_KEY_HANDLES policy prohibits importing raw secrets. "
+                "Use pqc_hybrid_keygen with store_as to generate keys as opaque handles."
+            )
+
+
 def handle_key_store_save(arguments: dict[str, Any]) -> dict[str, Any]:
     """Save a keygen output by name."""
     name = arguments["name"]
@@ -95,6 +119,17 @@ def handle_key_store_save(arguments: dict[str, Any]) -> dict[str, Any]:
     if not isinstance(key_data, dict):
         raise ValueError("key_data must be a JSON object (e.g., output of pqc_hybrid_keygen)")
 
+    # Policy enforcement: reject raw secrets when handle-only mode is active
+    from pqc_mcp_server.security_policy import get_policy
+
+    if get_policy().require_key_handles:
+        _reject_secret_fields(key_data)
+
+    # Size limit on key_data (prevents memory exhaustion via oversized dicts)
+    key_data_size = len(_json.dumps(key_data))
+    if key_data_size > _MAX_KEY_DATA_SIZE:
+        raise ValueError(f"key_data is {key_data_size} bytes (max {_MAX_KEY_DATA_SIZE})")
+
     entry = {
         "name": name,
         "key_data": key_data,