"""Integrity protection for the multi-round-trip `requestState` (MCP 2026-07-28). The spec requires servers to treat the client-echoed `requestState` as attacker-controlled: `RequestStateBoundary` seals every outgoing value and verifies every inbound echo, so handlers only ever see plaintext they minted. """ from __future__ import annotations import base64 import hashlib import hmac import json import logging import math import os import time from collections.abc import Callable, Mapping, Sequence from dataclasses import replace from typing import Any, NoReturn, Protocol, cast from cryptography.exceptions import InvalidTag from cryptography.hazmat.primitives.ciphers.aead import AESGCM from cryptography.hazmat.primitives.hashes import SHA256 from cryptography.hazmat.primitives.kdf.hkdf import HKDF from mcp_types import INTERNAL_ERROR, INVALID_PARAMS from mcp_types.methods import INPUT_REQUIRED_METHODS, is_input_required from mcp.server.auth.middleware.auth_context import get_access_token from mcp.server.auth.provider import principal_components from mcp.server.context import CallNext, HandlerResult, ServerRequestContext from mcp.shared.exceptions import MCPError __all__ = [ "AESGCMRequestStateCodec", "InvalidRequestState", "RequestStateBoundary", "RequestStateCodec", "RequestStateSecurity", "authenticated_principal", ] logger = logging.getLogger(__name__) class InvalidRequestState(Exception): """A sealed `requestState` token failed verification. The message is a log-only reason code; the boundary never puts it on the wire. """ class RequestStateCodec(Protocol): """Authenticated crypto over the framework's request-state envelope. The framework stamps and re-verifies every envelope claim (expiry, request binding, principal); a codec only provides integrity and, ideally, confidentiality (a sign-only codec leaves the payload client-readable). Requirements: `unseal(seal(payload))` round-trips, and `unseal` raises `InvalidRequestState` for any token it did not mint unmodified; tokens never name their algorithm (version with a format prefix bound under the authentication tag, RFC 8725); comparisons are constant-time. Both methods are synchronous, so cache key material rather than calling a KMS per token. """ def seal(self, payload: bytes) -> str: """Return an opaque URL-safe token protecting `payload`.""" ... def unseal(self, token: str) -> bytes: """Reverse `seal`. Raises: InvalidRequestState: Malformed, unauthentic, or unknown-key token. """ ... def authenticated_principal(ctx: ServerRequestContext[Any, Any]) -> str | None: """Default principal binding: the authenticated (client, issuer, subject) identity. Uses the same components session ownership uses, so two users of one OAuth client are distinct principals whenever the token verifier supplies a subject, and the binding degrades to the client identity when it does not. Returns `None` (state not principal-bound) on unauthenticated transports. """ token = get_access_token() if token is None: return None return compact_json(principal_components(token)) class RequestStateSecurity: """Policy for protecting `requestState`: codec, TTL, principal, audience. Exactly one of `keys` or `codec`: RequestStateSecurity(keys=[secret]) # built-in AES-256-GCM RequestStateSecurity(codec=MyKmsCodec()) # bring your own crypto RequestStateSecurity.ephemeral() # process-local key `keys` is the rotation ring: `keys[0]` seals, every key unseals. Zero-downtime rotation, each phase fully rolled out before the next: `keys=[old, new]`, then `keys=[new, old]`, then `keys=[new]` after one TTL. The boundary enforces expiry, request binding, audience, and principal for every codec, fail-closed in both directions. `audience=None` defers to the boundary's `default_audience` (`MCPServer` passes its server name). """ codec: RequestStateCodec ttl: float bind_principal: Callable[[ServerRequestContext[Any, Any]], str | None] | None audience: str | None def __init__( self, *, keys: Sequence[bytes | bytearray | str] | None = None, codec: RequestStateCodec | None = None, ttl: float = 600.0, bind_principal: Callable[[ServerRequestContext[Any, Any]], str | None] | None = authenticated_principal, audience: str | None = None, ) -> None: if (keys is None) == (codec is None): raise ValueError("RequestStateSecurity takes exactly one of keys= or codec=") if not (math.isfinite(ttl) and ttl > 0): raise ValueError(f"request-state ttl must be a positive finite number, got {ttl!r}") if keys is not None: self.codec = AESGCMRequestStateCodec(keys) else: assert codec is not None self.codec = codec self.ttl = ttl self.bind_principal = bind_principal self.audience = audience @classmethod def ephemeral(cls, *, ttl: float = 600.0, audience: str | None = None) -> RequestStateSecurity: """Protection under a key generated now and held only by this process. This is the policy `MCPServer` installs when `request_state_security=` is omitted; call it yourself on the lowlevel tier or to set `ttl`/ `audience`. Suits single-process deployments (stdio, one HTTP worker): state minted before a restart or by another worker is rejected. Multi-instance deployments must share a key via `keys=[...]`. """ return cls(keys=[os.urandom(32)], ttl=ttl, audience=audience) _KDF_INFO = b"mcp/request-state/v1/aes-256-gcm" _KID_INFO = b"mcp/request-state/v1/kid:" _TOKEN_PREFIX = "v1." _KID_LEN = 4 _NONCE_LEN = 12 def compact_json(value: Any, *, sort_keys: bool = False) -> str: """Canonical JSON for everything the state path digests or seals. ASCII output keeps the encode total: a lone surrogate in client-supplied text escapes instead of raising. Anything consuming this must parse with stdlib `json.loads`, which accepts those escapes (pydantic's JSON parser does not). """ return json.dumps(value, sort_keys=sort_keys, separators=(",", ":")) def _b64u(data: bytes) -> str: return base64.urlsafe_b64encode(data).decode().rstrip("=") def _b64u_decode(text: str) -> bytes: """Strict inverse of `_b64u`: only the canonical unpadded encoding decodes.""" raw = base64.urlsafe_b64decode(text + "=" * (-len(text) % 4)) if _b64u(raw) != text: raise ValueError("non-canonical base64url") return raw def _derive_key(secret: bytes) -> bytes: """Stretch an operator secret (>= 32 bytes, any format) into the AES-256 key.""" return HKDF(algorithm=SHA256(), length=32, salt=None, info=_KDF_INFO).derive(secret) class AESGCMRequestStateCodec: """Built-in codec: AES-256-GCM under key(s) derived with HKDF-SHA256. Tokens are encrypted, not merely signed, so clients cannot read the state. `keys[0]` seals; all keys unseal (rotation, see `RequestStateSecurity`). Each token carries a 4-byte non-secret key fingerprint for an O(1) ring lookup, and the "v1." prefix and fingerprint are bound into the GCM associated data, so a token cannot be replayed into another format version or ring slot. Key bytes are copied at construction. """ def __init__(self, keys: Sequence[bytes | bytearray | str]) -> None: for i, key in enumerate(cast("Sequence[object]", keys)): if not isinstance(key, bytes | bytearray | str): # Never coerce: bytes(32) would silently build an all-zero key. raise TypeError( f"request-state keys must be bytes, bytearray, or str; keys[{i}] is {type(key).__name__}" ) material = [k.encode() if isinstance(k, str) else bytes(k) for k in keys] if not material: raise ValueError("AESGCMRequestStateCodec requires at least one key") for i, k in enumerate(material): if len(k) < 32: raise ValueError( f"request-state keys must be at least 32 bytes of secret randomness; " f"keys[{i}] is {len(k)} bytes. " 'Generate one with: python -c "import secrets; print(secrets.token_hex(32))"' ) self._ring: dict[bytes, AESGCM] = {} self._mint_kid = b"" for i, secret in enumerate(material): key = _derive_key(secret) kid = hashlib.sha256(_KID_INFO + key).digest()[:_KID_LEN] if kid in self._ring: raise ValueError(f"keys[{i}] duplicates an earlier ring key") self._ring[kid] = AESGCM(key) if i == 0: self._mint_kid = kid def seal(self, payload: bytes) -> str: kid = self._mint_kid nonce = os.urandom(_NONCE_LEN) sealed = self._ring[kid].encrypt(nonce, payload, _TOKEN_PREFIX.encode() + kid) return _TOKEN_PREFIX + _b64u(kid + nonce + sealed) def unseal(self, token: str) -> bytes: if not token.startswith(_TOKEN_PREFIX): raise InvalidRequestState("malformed") try: raw = _b64u_decode(token[len(_TOKEN_PREFIX) :]) except ValueError as exc: raise InvalidRequestState("malformed") from exc if len(raw) < _KID_LEN + _NONCE_LEN + 16: raise InvalidRequestState("malformed") kid, nonce, sealed = raw[:_KID_LEN], raw[_KID_LEN : _KID_LEN + _NONCE_LEN], raw[_KID_LEN + _NONCE_LEN :] aead = self._ring.get(kid) if aead is None: raise InvalidRequestState("unknown key") try: return aead.decrypt(nonce, sealed, _TOKEN_PREFIX.encode() + kid) except InvalidTag: raise InvalidRequestState("seal") from None # The multi-round-trip carriers: the only methods whose results may carry `requestState`. _MRTR_METHODS = INPUT_REQUIRED_METHODS _ENVELOPE_VERSION = 1 _FUTURE_SKEW = 60.0 _PRINCIPAL_LABEL = b"mcp/request-state/principal:" _RoundBinding = tuple[str, str, str | None] """The (target, args-digest, principal) one round's envelope binds, computed once per round.""" def _reject(method: str, reason: str) -> NoReturn: """Refuse a round: frozen wire error, real reason to the server log only.""" logger.warning("requestState rejected on %s: %s", method, reason) raise MCPError( code=INVALID_PARAMS, message="Invalid or expired requestState", data={"reason": "invalid_request_state"}, ) def _request_identity(method: str, params: Mapping[str, Any] | None) -> tuple[str, str]: """Salient (target, args-digest) for the request a token binds to. Per-method allowlist, never a denylist: a future wire field cannot silently join the digest. """ p: Mapping[str, Any] = params or {} args: dict[str, Any] = {} if method == "resources/read": target = str(p.get("uri", "")) else: target, args = str(p.get("name", "")), p.get("arguments") or args return target, _b64u(hashlib.sha256(compact_json(args, sort_keys=True).encode()).digest()[:16]) def _principal_claim(principal: str) -> str: salt = os.urandom(8) tag = hashlib.sha256(_PRINCIPAL_LABEL + salt + _principal_bytes(principal)).digest()[:16] return _b64u(salt + tag) def _principal_matches(claim: str, principal: str) -> bool: try: raw = _b64u_decode(claim) except ValueError: return False # A wrong-length claim never matches: compare_digest handles mismatched sizes. expected = hashlib.sha256(_PRINCIPAL_LABEL + raw[:8] + _principal_bytes(principal)).digest()[:16] return hmac.compare_digest(raw[8:], expected) def _principal_bytes(principal: str) -> bytes: # The digest input is one-way and never decoded, so surrogatepass keeps it total. return principal.encode("utf-8", "surrogatepass") def _bound_principal( security: RequestStateSecurity, ctx: ServerRequestContext[Any, Any], fail: Callable[[str], NoReturn], ) -> str | None: """Run `bind_principal` under the deny-on-error discipline, in one place for both directions. `fail` converts a failure into the calling direction's wire shape: the frozen rejection when verifying, the sanitized internal error when sealing. """ try: principal = security.bind_principal(ctx) if security.bind_principal is not None else None except Exception: # deny-on-error: a raising principal binding must fail closed logger.exception("bind_principal raised while processing requestState on %s", ctx.method) fail("principal binding error") # The declared return type is str | None, but a user callback can ignore it. if principal is not None and not isinstance(cast("object", principal), str): fail(f"bind_principal returned {type(principal).__name__}, expected str or None") return principal class RequestStateBoundary: """Server middleware sealing/unsealing `requestState` at the wire boundary. Acts only on the multi-round-trip carriers (tools/call, prompts/get, resources/read); every other method passes through untouched. Inbound state is verified (codec unseal plus claims check) and replaced with the plaintext the server minted before any interceptor or handler runs; failure answers -32602 with the frozen message "Invalid or expired requestState", the real reason going to the server log only. Outbound, an `input_required` result carrying `requestState` is sealed in a fresh claims envelope; handlers and resolvers never call the codec. `default_audience` seeds the audience claim when the policy sets none, and must be stated explicitly: it is the service identity that stops state minted by another service sharing the same keys. `MCPServer` installs this middleware with its server name by default (under an ephemeral policy unless `request_state_security=` supplies one); lowlevel `Server` users append one to `server.middleware`, passing their server's name (or `None` to deliberately leave tokens audience-free). """ def __init__(self, security: RequestStateSecurity, *, default_audience: str | None) -> None: self._security = security self._audience = security.audience if security.audience is not None else default_audience async def __call__(self, ctx: ServerRequestContext[Any, Any], call_next: CallNext) -> HandlerResult: if ctx.method not in _MRTR_METHODS: return await call_next(ctx) binding: _RoundBinding | None = None if ctx.params is not None and ctx.params.get("requestState") is not None: # An explicit JSON null counts as absent: stripping the field is already in any client's power. plaintext, binding = self._unseal(ctx) ctx = replace(ctx, params={**ctx.params, "requestState": plaintext}) result = await call_next(ctx) return self._seal_result(ctx, result, binding) def _unseal(self, ctx: ServerRequestContext[Any, Any]) -> tuple[str, _RoundBinding]: assert ctx.params is not None wire = ctx.params["requestState"] if not isinstance(wire, str): _reject(ctx.method, "non-string requestState") security = self._security try: payload = security.codec.unseal(wire) except InvalidRequestState as exc: _reject(ctx.method, str(exc)) except Exception: # deny-on-error: a buggy custom codec must fail closed logger.exception("requestState codec raised during unseal on %s", ctx.method) _reject(ctx.method, "codec error") try: claims = json.loads(payload) version, iat, exp, inner = claims["v"], claims["iat"], claims["exp"], claims["s"] except (ValueError, KeyError, TypeError): _reject(ctx.method, "malformed") if version != _ENVELOPE_VERSION or not isinstance(inner, str): _reject(ctx.method, "malformed") now = time.time() # Accept-conditions are stated positively so a NaN claim fails the comparison and rejects. if not isinstance(iat, int | float) or not (iat <= now + _FUTURE_SKEW): _reject(ctx.method, "minted in the future") if not isinstance(exp, int | float) or not (now < exp): _reject(ctx.method, "expired") target, args_digest = _request_identity(ctx.method, ctx.params) if claims.get("m") != ctx.method or claims.get("t") != target or claims.get("a") != args_digest: _reject(ctx.method, "request binding") if claims.get("aud") != self._audience: _reject(ctx.method, "audience") def fail_verify(reason: str) -> NoReturn: _reject(ctx.method, reason) principal = _bound_principal(security, ctx, fail_verify) claim = claims.get("p") if (claim is None) != (principal is None): _reject(ctx.method, "principal drift") if claim is not None and principal is not None: if not isinstance(claim, str) or not _principal_matches(claim, principal): _reject(ctx.method, "principal") return inner, (target, args_digest, principal) def _seal_result( self, ctx: ServerRequestContext[Any, Any], result: HandlerResult, binding: _RoundBinding | None ) -> HandlerResult: # Spec-path results arrive as wire mappings; a short-circuiting middleware may return a model. if not is_input_required(result): return result state = result.get("requestState") if isinstance(result, Mapping) else result.request_state if state is None: return result if isinstance(result, Mapping): if not isinstance(state, str): # Only a short-circuiting middleware can put a non-string here; nothing to seal. return result return {**result, "requestState": self._seal(ctx, state, binding)} return result.model_copy(update={"request_state": self._seal(ctx, state, binding)}) def _seal(self, ctx: ServerRequestContext[Any, Any], state: str, binding: _RoundBinding | None = None) -> str: security = self._security if binding is None: def fail_seal(reason: str) -> NoReturn: logger.error("refusing to seal requestState on %s: %s", ctx.method, reason) raise MCPError(code=INTERNAL_ERROR, message="Internal error") target, args_digest = _request_identity(ctx.method, ctx.params) binding = (target, args_digest, _bound_principal(security, ctx, fail_seal)) target, args_digest, principal = binding now = time.time() claims: dict[str, Any] = { "v": _ENVELOPE_VERSION, "iat": now, "exp": now + security.ttl, "m": ctx.method, "t": target, "a": args_digest, "s": state, } if self._audience is not None: claims["aud"] = self._audience if principal is not None: claims["p"] = _principal_claim(principal) payload = compact_json(claims).encode() try: return security.codec.seal(payload) except Exception: # deny-on-error: a raising custom codec must not leak its failure logger.exception("requestState codec raised during seal on %s", ctx.method) raise MCPError(code=INTERNAL_ERROR, message="Internal error") from None