Require integrity protection for MRTR requestState (#3032)
This commit is contained in:
Max
GitHub
parent
985652491a
commit
7322ca56f4
9 changed files with 3338 additions and 183 deletions
|
|
@ -13,7 +13,7 @@ from __future__ import annotations
|
|||
from collections.abc import Mapping
|
||||
from functools import cache
|
||||
from types import MappingProxyType, UnionType
|
||||
from typing import Any, Final, Literal, TypeVar, get_args
|
||||
from typing import Any, Final, Literal, TypeGuard, TypeVar, cast, get_args
|
||||
|
||||
from pydantic import BaseModel, TypeAdapter
|
||||
|
||||
|
|
@ -28,6 +28,7 @@ __all__ = [
|
|||
"CLIENT_REQUESTS",
|
||||
"CLIENT_RESULTS",
|
||||
"CacheableMethod",
|
||||
"INPUT_REQUIRED_METHODS",
|
||||
"MONOLITH_NOTIFICATIONS",
|
||||
"MONOLITH_REQUESTS",
|
||||
"MONOLITH_RESULTS",
|
||||
|
|
@ -36,6 +37,7 @@ __all__ = [
|
|||
"SERVER_RESULTS",
|
||||
"SPEC_CLIENT_METHODS",
|
||||
"SPEC_CLIENT_NOTIFICATION_METHODS",
|
||||
"is_input_required",
|
||||
"parse_client_notification",
|
||||
"parse_client_request",
|
||||
"parse_client_result",
|
||||
|
|
@ -423,6 +425,22 @@ CACHEABLE_METHODS: Final[frozenset[str]] = frozenset(
|
|||
)
|
||||
"""Runtime mirror of `CacheableMethod`, derived from `MONOLITH_RESULTS`."""
|
||||
|
||||
INPUT_REQUIRED_METHODS: Final[frozenset[str]] = frozenset(
|
||||
method
|
||||
for method, row in MONOLITH_RESULTS.items()
|
||||
if any(
|
||||
issubclass(arm, types.InputRequiredResult) for arm in (get_args(row) if isinstance(row, UnionType) else (row,))
|
||||
)
|
||||
)
|
||||
"""Methods whose results may be `InputRequiredResult`, derived from `MONOLITH_RESULTS`."""
|
||||
|
||||
|
||||
def is_input_required(result: object) -> TypeGuard[types.InputRequiredResult | dict[str, Any]]:
|
||||
"""True when `result` is an `input_required` interim result, typed or wire-shaped."""
|
||||
if isinstance(result, types.InputRequiredResult):
|
||||
return True
|
||||
return isinstance(result, Mapping) and cast("Mapping[str, Any]", result).get("resultType") == "input_required"
|
||||
|
||||
|
||||
# --- Parse functions ---
|
||||
|
||||
|
|
|
|||
|
|
@ -609,7 +609,9 @@ class Client:
|
|||
callbacks and the call is retried automatically (up to
|
||||
`input_required_max_rounds`). To drive the loop yourself — e.g. to
|
||||
persist `request_state` across process restarts — use
|
||||
`client.session.call_tool(..., allow_input_required=True)`.
|
||||
`client.session.call_tool(..., allow_input_required=True)`. Persisted
|
||||
state is still subject to the server's TTL, request binding, and key
|
||||
lifetime; a server on the default process-local key rejects it after a restart.
|
||||
|
||||
Args:
|
||||
name: The name of the tool to call.
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@ from starlette.authentication import AuthCredentials, AuthenticationBackend, Sim
|
|||
from starlette.requests import HTTPConnection
|
||||
from starlette.types import Receive, Scope, Send
|
||||
|
||||
from mcp.server.auth.provider import AccessToken, TokenVerifier
|
||||
from mcp.server.auth.provider import AccessToken, TokenVerifier, principal_components
|
||||
|
||||
|
||||
class AuthenticatedUser(SimpleUser):
|
||||
|
|
@ -34,13 +34,8 @@ def authorization_context(user: AuthenticatedUser) -> AuthorizationContext:
|
|||
See `examples/servers/simple-auth/mcp_simple_auth/token_verifier.py` for
|
||||
a verifier that populates `subject` and `claims` from an introspection
|
||||
response."""
|
||||
token = user.access_token
|
||||
issuer = (token.claims or {}).get("iss")
|
||||
return AuthorizationContext(
|
||||
client_id=token.client_id,
|
||||
issuer=str(issuer) if issuer is not None else None,
|
||||
subject=token.subject,
|
||||
)
|
||||
client_id, issuer, subject = principal_components(user.access_token)
|
||||
return AuthorizationContext(client_id=client_id, issuer=issuer, subject=subject)
|
||||
|
||||
|
||||
class BearerAuthBackend(AuthenticationBackend):
|
||||
|
|
|
|||
|
|
@ -59,6 +59,17 @@ class AccessToken(BaseModel):
|
|||
claims: dict[str, Any] | None = None # additional claims (e.g. `iss`, `act`)
|
||||
|
||||
|
||||
def principal_components(token: AccessToken) -> tuple[str, str | None, str | None]:
|
||||
"""The (client_id, issuer, subject) triple identifying the principal a token represents.
|
||||
|
||||
The single source for "who is this token's principal": session ownership and
|
||||
request-state binding both build on it. Components the token verifier does
|
||||
not supply are `None`, so comparisons degrade to the remaining components.
|
||||
"""
|
||||
issuer = (token.claims or {}).get("iss")
|
||||
return token.client_id, str(issuer) if issuer is not None else None, token.subject
|
||||
|
||||
|
||||
RegistrationErrorCode = Literal[
|
||||
"invalid_redirect_uri",
|
||||
"invalid_client_metadata",
|
||||
|
|
|
|||
|
|
@ -3,6 +3,14 @@
|
|||
from mcp_types import Icon
|
||||
|
||||
from mcp.server.extension import Extension, MethodBinding, ResourceBinding, ToolBinding
|
||||
from mcp.server.request_state import (
|
||||
AESGCMRequestStateCodec,
|
||||
InvalidRequestState,
|
||||
RequestStateBoundary,
|
||||
RequestStateCodec,
|
||||
RequestStateSecurity,
|
||||
authenticated_principal,
|
||||
)
|
||||
|
||||
from .context import Context
|
||||
from .resolve import (
|
||||
|
|
@ -36,4 +44,10 @@ __all__ = [
|
|||
"require_client_extension",
|
||||
"ResourceSecurity",
|
||||
"DEFAULT_RESOURCE_SECURITY",
|
||||
"RequestStateSecurity",
|
||||
"RequestStateCodec",
|
||||
"RequestStateBoundary",
|
||||
"AESGCMRequestStateCodec",
|
||||
"InvalidRequestState",
|
||||
"authenticated_principal",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -28,7 +28,11 @@ Whether the consumer receives the unwrapped model or the full
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import hashlib
|
||||
import inspect
|
||||
import json
|
||||
import logging
|
||||
import types
|
||||
import typing
|
||||
from collections.abc import Callable, Hashable, Mapping
|
||||
|
|
@ -43,6 +47,7 @@ from mcp_types import (
|
|||
ElicitRequestFormParams,
|
||||
ElicitResult,
|
||||
FormElicitationCapability,
|
||||
InputRequest,
|
||||
InputRequests,
|
||||
InputRequiredResult,
|
||||
InputResponses,
|
||||
|
|
@ -61,6 +66,7 @@ from mcp.server.elicitation import (
|
|||
)
|
||||
from mcp.server.mcpserver.context import Context
|
||||
from mcp.server.mcpserver.exceptions import InvalidSignature, ToolError
|
||||
from mcp.server.request_state import compact_json
|
||||
from mcp.shared._callable_inspection import is_async_callable
|
||||
from mcp.shared.exceptions import MCPError
|
||||
|
||||
|
|
@ -73,7 +79,9 @@ _ELICITATION_RESULT_MEMBERS = (AcceptedElicitation, DeclinedElicitation, Cancell
|
|||
# `InputRequiredResult` rather than as a standalone server-to-client request.
|
||||
# Pinned (not `LATEST_MODERN_VERSION`, which moves when newer revisions are added).
|
||||
_INPUT_REQUIRED_VERSION = "2026-07-28"
|
||||
_STATE_VERSION = 1
|
||||
_STATE_VERSION = 3 # v3: recorded and pended outcomes pinned to ASCII-canonical question renders
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class Resolve:
|
||||
|
|
@ -369,7 +377,11 @@ class _Resolution:
|
|||
self.context = context
|
||||
self.input_required = input_required
|
||||
self.answers: InputResponses = context.input_responses or {} if input_required else {}
|
||||
self.state = _decode_state(context.request_state) if input_required else {}
|
||||
decoded = _decode_state(context.request_state if input_required else None)
|
||||
self.state = decoded.outcomes
|
||||
# Digests of the questions asked last round: an answer is accepted only
|
||||
# for the exact rendering the client was shown.
|
||||
self.asked = decoded.asked
|
||||
# In-call dedup keyed by resolver identity (distinguishes two instances of
|
||||
# the same bound method); `persist` holds the wire-shaped record of each
|
||||
# elicited outcome, keyed by its wire key - exactly what the next round's
|
||||
|
|
@ -431,7 +443,8 @@ async def resolve_arguments(
|
|||
injected[name] = outcome if wants_union else _unwrap(outcome, name)
|
||||
|
||||
if res.pending:
|
||||
return InputRequiredResult(input_requests=res.pending, request_state=_encode_state(res.persist))
|
||||
asked = {key: _request_digest(request) for key, request in res.pending.items()}
|
||||
return InputRequiredResult(input_requests=res.pending, request_state=_encode_state(res.persist, asked))
|
||||
return injected
|
||||
|
||||
|
||||
|
|
@ -494,19 +507,25 @@ async def _elicit(elicit: Elicit[Any], key: str, res: _Resolution) -> Elicitatio
|
|||
if not res.input_required:
|
||||
return await res.context.elicit(elicit.message, elicit.schema)
|
||||
|
||||
request = _elicit_request(elicit)
|
||||
q = _request_digest(request)
|
||||
|
||||
# A recorded outcome from a prior round is consulted only here, after the body
|
||||
# decided to ask, so a `request_state` entry can never stand in for a resolver's
|
||||
# own computation. Re-validate it against the live `Elicit.schema`. A recorded
|
||||
# outcome wins over a re-sent answer; an invalid entry self-deletes and falls
|
||||
# through to the fresh answer (or to re-asking).
|
||||
outcome = _restore_outcome(res, key, elicit.schema)
|
||||
# own computation. A recorded outcome wins over a re-sent answer.
|
||||
outcome = _restore_outcome(res, key, elicit.schema, q)
|
||||
if outcome is not None:
|
||||
return outcome
|
||||
|
||||
answer = res.answers.get(key)
|
||||
# An answer counts only for the rendering recorded when it was asked; an answer to
|
||||
# an unrecorded or differently-worded question re-asks instead of being consumed.
|
||||
if answer is not None and res.asked.get(key) != q:
|
||||
logger.info("Discarding the answer for resolver %r: the question changed since it was asked", key)
|
||||
answer = None
|
||||
if answer is None:
|
||||
_require_form_elicitation(res.context, key)
|
||||
res.pending[key] = _elicit_request(elicit)
|
||||
res.pending[key] = request
|
||||
raise _Pending
|
||||
if not isinstance(answer, ElicitResult):
|
||||
raise ToolError(f"Resolver {key!r} received a non-elicitation response")
|
||||
|
|
@ -521,12 +540,12 @@ async def _elicit(elicit: Elicit[Any], key: str, res: _Resolution) -> Elicitatio
|
|||
) from e
|
||||
# Persist the exact wire content that just passed validation - never the
|
||||
# model - so restoring next round revalidates the same bytes the client sent.
|
||||
res.persist[key] = _StateEntry(action="accept", data=answer.content)
|
||||
res.persist[key] = _StateEntry(action="accept", data=answer.content, q=q)
|
||||
return AcceptedElicitation(data=data)
|
||||
if answer.action == "decline":
|
||||
res.persist[key] = _StateEntry(action="decline")
|
||||
res.persist[key] = _StateEntry(action="decline", q=q)
|
||||
return DeclinedElicitation()
|
||||
res.persist[key] = _StateEntry(action="cancel")
|
||||
res.persist[key] = _StateEntry(action="cancel", q=q)
|
||||
return CancelledElicitation()
|
||||
|
||||
|
||||
|
|
@ -595,37 +614,58 @@ class _StateEntry(BaseModel):
|
|||
|
||||
action: Literal["accept", "decline", "cancel"]
|
||||
data: Any = None
|
||||
q: str | None = None
|
||||
"""Digest of the exact rendered question this outcome answered."""
|
||||
|
||||
|
||||
def _request_digest(request: InputRequest) -> str:
|
||||
"""Pin an outcome to the exact rendered question the client was shown.
|
||||
|
||||
A redeploy that rewords or reshapes a question re-asks it instead of reusing the recorded answer.
|
||||
"""
|
||||
params = request.params
|
||||
rendered = compact_json(params.model_dump(mode="json", by_alias=True, exclude_none=True) if params else None)
|
||||
digest = hashlib.sha256(rendered.encode()).digest()[:16]
|
||||
return base64.urlsafe_b64encode(digest).decode().rstrip("=")
|
||||
|
||||
|
||||
class _State(BaseModel):
|
||||
"""The decoded `request_state`: resolver outcomes from earlier rounds."""
|
||||
"""The decoded `request_state`: resolver progress from earlier rounds."""
|
||||
|
||||
v: int
|
||||
outcomes: dict[str, _StateEntry] = {}
|
||||
asked: dict[str, str] = {}
|
||||
"""Question digest of each elicitation asked last round, keyed by wire key."""
|
||||
|
||||
|
||||
def _decode_state(request_state: str | None) -> dict[str, _StateEntry]:
|
||||
def _decode_state(request_state: str | None) -> _State:
|
||||
"""Decode the per-call resolution progress from `request_state`.
|
||||
|
||||
`request_state` is client-trusted (integrity sealing is a follow-up); validate
|
||||
it through `_State` and treat anything malformed as "no progress yet".
|
||||
Parsed with stdlib `json.loads` because `_encode_state` may emit escaped
|
||||
lone surrogates, which pydantic's JSON parser rejects. The string arrives
|
||||
boundary-authenticated, so malformed content or a version mismatch is
|
||||
drift within the operator's own fleet (e.g. a rolling upgrade) and is
|
||||
treated as "no progress yet".
|
||||
"""
|
||||
empty = _State(v=_STATE_VERSION)
|
||||
if not request_state:
|
||||
return {}
|
||||
return empty
|
||||
try:
|
||||
state = _State.model_validate_json(request_state)
|
||||
except ValidationError:
|
||||
return {}
|
||||
return state.outcomes if state.v == _STATE_VERSION else {}
|
||||
state = _State.model_validate(json.loads(request_state))
|
||||
except ValueError:
|
||||
return empty
|
||||
return state if state.v == _STATE_VERSION else empty
|
||||
|
||||
|
||||
def _encode_state(outcomes: Mapping[str, _StateEntry]) -> str:
|
||||
"""Encode recorded elicitation outcomes (keyed by wire key) for the next round.
|
||||
def _encode_state(outcomes: Mapping[str, _StateEntry], asked: Mapping[str, str]) -> str:
|
||||
"""Encode recorded outcomes and asked-question digests for the next round.
|
||||
|
||||
Entries already hold the client's wire-shaped data exactly as it was sent (and
|
||||
validated), so encoding is pure wrapping: encode-restore is the identity.
|
||||
Outcome entries already hold the client's wire-shaped data exactly as it was
|
||||
sent (and validated), so encoding is pure wrapping: encode-restore is the
|
||||
identity.
|
||||
"""
|
||||
return _State(v=_STATE_VERSION, outcomes=dict(outcomes)).model_dump_json()
|
||||
state = _State(v=_STATE_VERSION, outcomes=dict(outcomes), asked=dict(asked))
|
||||
return compact_json(state.model_dump(mode="json"))
|
||||
|
||||
|
||||
def _outcome_from_state(entry: _StateEntry, schema: type[BaseModel]) -> ElicitationResult[Any]:
|
||||
|
|
@ -642,12 +682,12 @@ def _outcome_from_state(entry: _StateEntry, schema: type[BaseModel]) -> Elicitat
|
|||
return _accepted(schema.model_validate(entry.data))
|
||||
|
||||
|
||||
def _restore_outcome(res: _Resolution, key: str, schema: type[BaseModel]) -> ElicitationResult[Any] | None:
|
||||
def _restore_outcome(res: _Resolution, key: str, schema: type[BaseModel], q: str) -> ElicitationResult[Any] | None:
|
||||
"""Restore `key`'s recorded outcome from a prior round, or `None` when absent.
|
||||
|
||||
`request_state` is client-trusted, so an entry whose data fails validation gets
|
||||
the `_decode_state` treatment - dropped as if no progress was recorded, so the
|
||||
question is asked again - rather than surfacing a validation error.
|
||||
An entry pinned to a question digest other than `q`, or whose accepted
|
||||
data fails validation against the live `schema`, is dropped as if no
|
||||
progress was recorded, so the question is asked again.
|
||||
|
||||
Carries the original decoded entry forward unchanged in `res.persist`: if a
|
||||
later resolver is still pending, the next round's `request_state` is built from
|
||||
|
|
@ -657,6 +697,9 @@ def _restore_outcome(res: _Resolution, key: str, schema: type[BaseModel]) -> Eli
|
|||
entry = res.state.get(key)
|
||||
if entry is None:
|
||||
return None
|
||||
if entry.q != q:
|
||||
del res.state[key]
|
||||
return None
|
||||
try:
|
||||
outcome = _outcome_from_state(entry, schema)
|
||||
except ValidationError:
|
||||
|
|
|
|||
|
|
@ -83,6 +83,7 @@ from mcp.server.mcpserver.resources import (
|
|||
from mcp.server.mcpserver.tools import Tool, ToolManager
|
||||
from mcp.server.mcpserver.utilities.context_injection import find_context_parameter
|
||||
from mcp.server.mcpserver.utilities.logging import configure_logging, get_logger
|
||||
from mcp.server.request_state import RequestStateBoundary, RequestStateSecurity
|
||||
from mcp.server.sse import SseServerTransport
|
||||
from mcp.server.stdio import stdio_server
|
||||
from mcp.server.streamable_http import EventStore
|
||||
|
|
@ -133,6 +134,15 @@ class Settings(BaseSettings, Generic[LifespanResultT]):
|
|||
auth: AuthSettings | None
|
||||
|
||||
|
||||
_MISSING_AUDIENCE = (
|
||||
"request_state_security is configured but this server has no name. Sealed\n"
|
||||
"requestState carries the server name as an audience claim, so state minted by\n"
|
||||
"another service that shares the same keys is rejected; unnamed servers would\n"
|
||||
"all stamp the same placeholder and the check would mean nothing. Name the\n"
|
||||
'server (MCPServer("my-service", ...)) or set RequestStateSecurity(audience=...).'
|
||||
)
|
||||
|
||||
|
||||
def lifespan_wrapper(
|
||||
app: MCPServer[LifespanResultT],
|
||||
lifespan: Callable[[MCPServer[LifespanResultT]], AbstractAsyncContextManager[LifespanResultT]],
|
||||
|
|
@ -170,6 +180,7 @@ class MCPServer(Generic[LifespanResultT]):
|
|||
lifespan: Callable[[MCPServer[LifespanResultT]], AbstractAsyncContextManager[LifespanResultT]] | None = None,
|
||||
auth: AuthSettings | None = None,
|
||||
resource_security: ResourceSecurity = DEFAULT_RESOURCE_SECURITY,
|
||||
request_state_security: RequestStateSecurity | None = None,
|
||||
cache_hints: Mapping[CacheableMethod, CacheHint] | None = None,
|
||||
):
|
||||
self._resource_security = resource_security
|
||||
|
|
@ -210,6 +221,17 @@ class MCPServer(Generic[LifespanResultT]):
|
|||
# We need to create a Lifespan type that is a generic on the server type, like Starlette does.
|
||||
lifespan=(lifespan_wrapper(self, self.settings.lifespan) if self.settings.lifespan else default_lifespan), # type: ignore
|
||||
)
|
||||
# Ordering: inside OpenTelemetry (spans record the sealed wire form),
|
||||
# outside extension interceptors (extensions see plaintext).
|
||||
if request_state_security is None:
|
||||
security = RequestStateSecurity.ephemeral()
|
||||
else:
|
||||
# A supplied policy usually means shared keys, where the audience claim is
|
||||
# what separates services; an unnamed server would stamp the placeholder.
|
||||
if not name and request_state_security.audience is None:
|
||||
raise ValueError(_MISSING_AUDIENCE)
|
||||
security = request_state_security
|
||||
self._lowlevel_server.middleware.append(RequestStateBoundary(security, default_audience=self.name))
|
||||
# Validate auth configuration
|
||||
if self.settings.auth is not None:
|
||||
if auth_server_provider and token_verifier: # pragma: no cover
|
||||
|
|
|
|||
454
src/mcp/server/request_state.py
Normal file
454
src/mcp/server/request_state.py
Normal file
|
|
@ -0,0 +1,454 @@
|
|||
"""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
|
||||
|
|
@ -204,7 +204,7 @@ class ServerRunner(Generic[LifespanT]):
|
|||
if (hint := self.server.cache_hints.get(method)) is not None:
|
||||
if isinstance(result, CacheableResult):
|
||||
result = apply_cache_hint(result, hint)
|
||||
elif isinstance(result, Mapping) and result.get("resultType") != "input_required":
|
||||
elif isinstance(result, Mapping) and not _methods.is_input_required(result):
|
||||
# Hint keys first so wire keys the handler set win, matching `apply_cache_hint` precedence.
|
||||
result = {"ttlMs": hint.ttl_ms, "cacheScope": hint.scope, **result}
|
||||
# Dump and serialize inside the chain so the OpenTelemetry span (the
|
||||
|
|
|
|||
Loading…
Reference in a new issue