455 lines
19 KiB
Python
455 lines
19 KiB
Python
"""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
|