"""Unit tests for `mcp.server.request_state`: codec, security policy, and default principal binding."""

import base64

import string

from collections.abc import Callable

from typing import Any, cast

import pytest

from inline_snapshot import snapshot

from mcp.server.auth.middleware.auth_context import auth_context_var

from mcp.server.auth.middleware.bearer_auth import AuthenticatedUser, authorization_context

from mcp.server.auth.provider import AccessToken, principal_components

from mcp.server.context import ServerRequestContext

from mcp.server.request_state import (

AESGCMRequestStateCodec,

InvalidRequestState,

RequestStateSecurity,

authenticated_principal,

)

_TOKEN_PREFIX = "v1."

_KID_LEN = 4

_NONCE_LEN = 12

_GCM_TAG_LEN = 16

_BODY_FLOOR = _KID_LEN + _NONCE_LEN + _GCM_TAG_LEN

_B64URL_ALPHABET = set(string.ascii_letters + string.digits + "-_")

_KEY_A = b"a" * 32

_KEY_B = b"b" * 32

_KEY_OLD = b"o" * 32

_KEY_NEW = b"n" * 32

# Distinctive plaintext: opacity and log-secrecy assertions search for it.

_PAYLOAD = b"sentinel-plaintext-3f9c"

# `InvalidRequestState` messages are short log-only reason codes, never payload.

_REASON_CODE_MAX_LEN = 40

def _b64u_nopad(data: bytes) -> str:

return base64.urlsafe_b64encode(data).decode().rstrip("=")

def _decode_body(token: str) -> bytes:

body = token.removeprefix(_TOKEN_PREFIX)

return base64.urlsafe_b64decode(body + "=" * (-len(body) % 4))

def _flip_body_byte(token: str, index: int) -> str:

raw = bytearray(_decode_body(token))

raw[index] ^= 0xFF

return _TOKEN_PREFIX + _b64u_nopad(bytes(raw))

def _flip_prefix_char(token: str) -> str:

return "x" + token[1:]

def _flip_kid_byte(token: str) -> str:

return _flip_body_byte(token, 0)

def _flip_nonce_byte(token: str) -> str:

return _flip_body_byte(token, _KID_LEN)

def _flip_ciphertext_byte(token: str) -> str:

return _flip_body_byte(token, _KID_LEN + _NONCE_LEN)

def _flip_tag_byte(token: str) -> str:

return _flip_body_byte(token, -1)

def _inject_junk_chars(body: str) -> str:

return body[:10] + "!@\n*" + body[10:]

def _append_newline(body: str) -> str:

return body + "\n"

def _append_padding(body: str) -> str:

return body + "=" * (-len(body) % 4 or 4)

def _bare_context() -> ServerRequestContext[Any, Any]:

return ServerRequestContext(

session=cast("Any", None),

lifespan_context={},

protocol_version="2026-07-28",

method="tools/call",

)

class _StaticCodec:

"""Minimal `RequestStateCodec` stand-in for policy tests; no real crypto."""

def seal(self, payload: bytes) -> str:

return payload.hex()

def unseal(self, token: str) -> bytes:

return bytes.fromhex(token)

# -- AESGCMRequestStateCodec --------------------------------------------------

@pytest.mark.parametrize(

"payload",

[

pytest.param(b"", id="empty"),

pytest.param(b"plain ascii state", id="ascii"),

pytest.param("ünïcødé – 状態".encode(), id="multi-byte-utf8"),

pytest.param(bytes(range(256)), id="raw-binary"),

pytest.param(bytes(range(256)) * 256, id="64KiB"),

],

)

def test_seal_unseal_round_trips_any_payload(payload: bytes) -> None:

"""SDK-defined: the codec is byte-transparent, so any payload survives seal/unseal unchanged."""

codec = AESGCMRequestStateCodec([_KEY_A])

assert codec.unseal(codec.seal(payload)) == payload

def test_a_sealed_token_is_v1_plus_unpadded_b64url_over_kid_nonce_and_ciphertext() -> None:

"""SDK-defined token format: "v1." plus unpadded base64url over kid(4) || nonce(12) || ciphertext+tag."""

token = AESGCMRequestStateCodec([_KEY_A]).seal(_PAYLOAD)

assert token.startswith(_TOKEN_PREFIX)

body = token.removeprefix(_TOKEN_PREFIX)

assert "=" not in body

assert set(body) <= _B64URL_ALPHABET

assert len(_decode_body(token)) == _KID_LEN + _NONCE_LEN + len(_PAYLOAD) + _GCM_TAG_LEN

def test_two_seals_of_the_same_payload_produce_distinct_tokens_that_both_unseal() -> None:

"""SDK-defined: every seal draws a fresh nonce, so identical payloads yield distinct tokens that both verify."""

codec = AESGCMRequestStateCodec([_KEY_A])

first = codec.seal(_PAYLOAD)

second = codec.seal(_PAYLOAD)

assert first != second

assert codec.unseal(first) == _PAYLOAD

assert codec.unseal(second) == _PAYLOAD

@pytest.mark.parametrize(

"corrupt",

[

pytest.param(_flip_prefix_char, id="prefix-char"),

pytest.param(_flip_kid_byte, id="kid-byte"),

pytest.param(_flip_nonce_byte, id="nonce-byte"),

pytest.param(_flip_ciphertext_byte, id="ciphertext-byte"),

pytest.param(_flip_tag_byte, id="tag-byte"),

],

)

def test_a_token_corrupted_in_any_region_is_rejected_without_echoing_the_payload(

corrupt: Callable[[str], str],

) -> None:

"""Spec-mandated (basic/patterns/mrtr, server requirement 4): any corrupted token region is rejected."""

codec = AESGCMRequestStateCodec([_KEY_A])

token = codec.seal(_PAYLOAD)

with pytest.raises(InvalidRequestState) as exc:

codec.unseal(corrupt(token))

message = str(exc.value)

assert len(message) <= _REASON_CODE_MAX_LEN

assert _PAYLOAD.decode() not in message

@pytest.mark.parametrize(

"token",

[

pytest.param("", id="empty-string"),

pytest.param(_b64u_nopad(b"\x00" * 64), id="missing-prefix"),

pytest.param(_TOKEN_PREFIX + "!!!not-base64!!!", id="garbage-after-prefix"),

pytest.param(_TOKEN_PREFIX + _b64u_nopad(b"\x00" * (_BODY_FLOOR - 1)), id="below-floor"),

],

)

def test_a_structurally_malformed_token_is_rejected(token: str) -> None:

"""Spec-mandated (basic/patterns/mrtr, server requirement 4): tokens this codec never minted fail."""

with pytest.raises(InvalidRequestState):

AESGCMRequestStateCodec([_KEY_A]).unseal(token)

def test_a_token_minted_under_a_key_outside_the_ring_is_rejected_as_unknown_key() -> None:

"""Spec-mandated (basic/patterns/mrtr, server requirement 4): a foreign-key token fails as "unknown key"."""

token = AESGCMRequestStateCodec([_KEY_A]).seal(_PAYLOAD)

with pytest.raises(InvalidRequestState) as exc:

AESGCMRequestStateCodec([_KEY_B]).unseal(token)

assert str(exc.value) == "unknown key"

@pytest.mark.parametrize(

"ring",

[

pytest.param([_KEY_OLD, _KEY_NEW], id="rotation-phase-1"),

pytest.param([_KEY_NEW, _KEY_OLD], id="rotation-phase-2"),

],

)

def test_a_token_minted_under_the_old_key_unseals_under_any_ring_containing_it(ring: list[bytes]) -> None:

"""SDK-defined rotation: every ring key verifies, so old-key state survives both rollout phases."""

token = AESGCMRequestStateCodec([_KEY_OLD]).seal(_PAYLOAD)

assert AESGCMRequestStateCodec(ring).unseal(token) == _PAYLOAD

def test_the_first_ring_key_mints_and_later_ring_keys_only_verify() -> None:

"""SDK-defined rotation: keys[0] is the minter, so [new, old] state verifies under [new] but not [old]."""

token = AESGCMRequestStateCodec([_KEY_NEW, _KEY_OLD]).seal(_PAYLOAD)

assert AESGCMRequestStateCodec([_KEY_NEW]).unseal(token) == _PAYLOAD

with pytest.raises(InvalidRequestState):

AESGCMRequestStateCodec([_KEY_OLD]).unseal(token)

def test_a_token_minted_under_a_retired_key_is_rejected() -> None:

"""Spec-mandated (basic/patterns/mrtr, server requirement 4): retired-key state fails verification."""

token = AESGCMRequestStateCodec([_KEY_OLD]).seal(_PAYLOAD)

with pytest.raises(InvalidRequestState):

AESGCMRequestStateCodec([_KEY_NEW]).unseal(token)

def test_an_empty_key_ring_is_rejected_at_construction() -> None:

"""SDK-defined: an empty ring is a configuration error caught at construction."""

with pytest.raises(ValueError) as exc:

AESGCMRequestStateCodec([])

assert str(exc.value) == snapshot("AESGCMRequestStateCodec requires at least one key")

def test_a_key_shorter_than_32_bytes_is_rejected_with_generation_guidance() -> None:

"""SDK-defined: keys must carry at least 32 bytes; the error includes generation guidance."""

with pytest.raises(ValueError) as exc:

AESGCMRequestStateCodec([b"k" * 31])

assert str(exc.value) == snapshot(

"request-state keys must be at least 32 bytes of secret randomness; keys[0] is 31 bytes. "

'Generate one with: python -c "import secrets; print(secrets.token_hex(32))"'

)

def test_a_duplicate_key_in_the_ring_is_rejected_at_construction() -> None:

"""SDK-defined: duplicate ring keys are a rotation mistake caught at construction."""

with pytest.raises(ValueError) as exc:

AESGCMRequestStateCodec([_KEY_A, _KEY_A])

assert str(exc.value) == snapshot("keys[1] duplicates an earlier ring key")

def test_a_non_key_typed_ring_entry_is_rejected_naming_its_index_and_type() -> None:

"""SDK-defined: a non-key ring entry raises a TypeError naming its index and type, in codec and policy."""

with pytest.raises(TypeError) as exc:

AESGCMRequestStateCodec([_KEY_A, cast("Any", 32)])

assert str(exc.value) == snapshot("request-state keys must be bytes, bytearray, or str; keys[1] is int")

with pytest.raises(TypeError) as exc:

RequestStateSecurity(keys=[cast("Any", 32)])

assert str(exc.value) == snapshot("request-state keys must be bytes, bytearray, or str; keys[0] is int")

def test_a_mixed_ring_of_bytes_bytearray_and_str_entries_still_works() -> None:

"""SDK-defined: bytes, bytearray, and str keys interoperate in one ring."""

codec = AESGCMRequestStateCodec([_KEY_A, bytearray(_KEY_B), "c" * 32])

assert codec.unseal(codec.seal(_PAYLOAD)) == _PAYLOAD

assert codec.unseal(AESGCMRequestStateCodec([bytearray(_KEY_B)]).seal(_PAYLOAD)) == _PAYLOAD

assert codec.unseal(AESGCMRequestStateCodec(["c" * 32]).seal(_PAYLOAD)) == _PAYLOAD

def test_a_str_key_is_equivalent_to_its_utf8_bytes_form() -> None:

"""SDK-defined: a str key is utf-8 encoded, so it is the same ring key as its bytes spelling."""

token = AESGCMRequestStateCodec(["k" * 32]).seal(_PAYLOAD)

assert AESGCMRequestStateCodec([b"k" * 32]).unseal(token) == _PAYLOAD

def test_bytearray_key_material_is_copied_at_construction() -> None:

"""SDK-defined: key bytes are copied at construction; mutating the caller's bytearray later has no effect."""

material = bytearray(b"m" * 32)

codec = AESGCMRequestStateCodec([cast("Any", material)])

minted_before_mutation = codec.seal(_PAYLOAD)

material[:] = b"X" * 32

assert codec.unseal(minted_before_mutation) == _PAYLOAD

assert AESGCMRequestStateCodec([b"m" * 32]).unseal(codec.seal(_PAYLOAD)) == _PAYLOAD

def test_the_token_reveals_the_payload_neither_in_its_text_nor_its_decoded_bytes() -> None:

"""SDK-defined: the token is encrypted, not merely signed, so the plaintext appears nowhere in it."""

token = AESGCMRequestStateCodec([_KEY_A]).seal(_PAYLOAD)

assert _PAYLOAD.decode() not in token

assert _b64u_nopad(_PAYLOAD) not in token

assert _PAYLOAD.hex() not in token

assert _PAYLOAD not in _decode_body(token)

def test_every_substitution_of_the_final_token_character_is_rejected() -> None:

"""Spec-mandated (basic/patterns/mrtr, server requirement 4): canonical decoding

rejects every final-character substitution despite base64 don't-care padding bits."""

codec = AESGCMRequestStateCodec([_KEY_A])

body = codec.seal(_PAYLOAD).removeprefix(_TOKEN_PREFIX)

substitutions = [c for c in sorted(_B64URL_ALPHABET) if c != body[-1]]

assert len(substitutions) == 63

for c in substitutions:

with pytest.raises(InvalidRequestState):

codec.unseal(_TOKEN_PREFIX + body[:-1] + c)

@pytest.mark.parametrize(

"mangle",

[

pytest.param(_inject_junk_chars, id="junk-chars-injected"),

pytest.param(_append_newline, id="newline-appended"),

pytest.param(_append_padding, id="padding-appended"),

],

)

def test_a_non_canonical_token_body_is_rejected(mangle: Callable[[str], str]) -> None:

"""Spec-mandated (basic/patterns/mrtr, server requirement 4): lax-decoder aliases of a token are rejected."""

codec = AESGCMRequestStateCodec([_KEY_A])

body = codec.seal(_PAYLOAD).removeprefix(_TOKEN_PREFIX)

with pytest.raises(InvalidRequestState):

codec.unseal(_TOKEN_PREFIX + mangle(body))

def test_a_token_reprefixed_to_a_future_format_version_is_rejected() -> None:

"""Spec-mandated (basic/patterns/mrtr, server requirement 4): the prefix is tag-bound; "v2." replay fails."""

codec = AESGCMRequestStateCodec([_KEY_A])

token = codec.seal(_PAYLOAD)

with pytest.raises(InvalidRequestState):

codec.unseal("v2." + token.removeprefix(_TOKEN_PREFIX))

def test_a_kid_transplanted_onto_another_tokens_body_is_rejected() -> None:

"""Spec-mandated (basic/patterns/mrtr, server requirement 4): the kid is tag-bound; transplanting it fails."""

raw_a = _decode_body(AESGCMRequestStateCodec([_KEY_A]).seal(_PAYLOAD))

raw_b = _decode_body(AESGCMRequestStateCodec([_KEY_B]).seal(_PAYLOAD))

assert raw_a[:_KID_LEN] != raw_b[:_KID_LEN]

transplanted = _TOKEN_PREFIX + _b64u_nopad(raw_a[:_KID_LEN] + raw_b[_KID_LEN:])

with pytest.raises(InvalidRequestState):

AESGCMRequestStateCodec([_KEY_A, _KEY_B]).unseal(transplanted)

# -- RequestStateSecurity -----------------------------------------------------

def test_keys_and_codec_together_are_rejected_at_policy_construction() -> None:

"""SDK-defined: keys= and codec= are mutually exclusive."""

with pytest.raises(ValueError) as exc:

RequestStateSecurity(keys=[_KEY_A], codec=_StaticCodec())

assert str(exc.value) == snapshot("RequestStateSecurity takes exactly one of keys= or codec=")

def test_a_policy_with_neither_keys_nor_codec_is_rejected() -> None:

"""SDK-defined: a policy must name its codec; an empty policy is a mistake, not a posture."""

with pytest.raises(ValueError) as exc:

RequestStateSecurity()

assert str(exc.value) == snapshot("RequestStateSecurity takes exactly one of keys= or codec=")

@pytest.mark.parametrize(

"ttl",

[

pytest.param(0.0, id="zero"),

pytest.param(-600.0, id="negative"),

pytest.param(float("nan"), id="nan"),

pytest.param(float("inf"), id="inf"),

],

)

def test_a_non_positive_or_non_finite_ttl_is_rejected_at_policy_construction(ttl: float) -> None:

"""SDK-defined: zero, negative, NaN, and infinite ttl fail at construction for keys and ephemeral() alike."""

with pytest.raises(ValueError, match="positive finite"):

RequestStateSecurity(keys=[_KEY_A], ttl=ttl)

with pytest.raises(ValueError, match="positive finite"):

RequestStateSecurity.ephemeral(ttl=ttl)

def test_keys_produce_a_working_built_in_codec_on_the_policy() -> None:

"""SDK-defined: keys=[...] builds the built-in AES-GCM codec, exposed on .codec."""

security = RequestStateSecurity(keys=[_KEY_A])

assert isinstance(security.codec, AESGCMRequestStateCodec)

assert security.codec.unseal(security.codec.seal(_PAYLOAD)) == _PAYLOAD

def test_a_custom_codec_is_stored_on_the_policy_as_is() -> None:

"""SDK-defined: codec=... stores the caller's object unwrapped."""

codec = _StaticCodec()

security = RequestStateSecurity(codec=codec)

assert security.codec is codec

assert codec.unseal(codec.seal(_PAYLOAD)) == _PAYLOAD

def test_ephemeral_policies_are_protected_and_mutually_unintelligible() -> None:

"""SDK-defined: ephemeral() protects under a process-local key, so a sibling instance rejects its tokens."""

first = RequestStateSecurity.ephemeral()

second = RequestStateSecurity.ephemeral()

token = first.codec.seal(_PAYLOAD)

assert first.codec.unseal(token) == _PAYLOAD

with pytest.raises(InvalidRequestState):

second.codec.unseal(token)

def test_the_policy_stores_an_explicit_audience_and_defaults_to_none() -> None:

"""SDK-defined: audience is stored as given; None defers to the server tier's `default_audience`."""

assert RequestStateSecurity(keys=[_KEY_A]).audience is None

assert RequestStateSecurity(keys=[_KEY_A], audience="svc").audience == "svc"

assert RequestStateSecurity.ephemeral(audience="svc").audience == "svc"

def test_the_default_principal_binding_is_authenticated_principal() -> None:

"""SDK-defined: an unconfigured policy binds state to the authenticated OAuth client by default."""

assert RequestStateSecurity(keys=[_KEY_A]).bind_principal is authenticated_principal

def test_an_explicit_principal_binding_callable_is_stored() -> None:

"""SDK-defined: a custom bind_principal callable is stored as given."""

def tenant_binding(ctx: ServerRequestContext[Any, Any]) -> str | None:

return "tenant-1"

security = RequestStateSecurity(keys=[_KEY_A], bind_principal=tenant_binding)

assert security.bind_principal is tenant_binding

assert tenant_binding(_bare_context()) == "tenant-1"

# -- authenticated_principal ----------------------------------------------------

def test_authenticated_principal_is_none_without_an_auth_context() -> None:

"""SDK-defined: without an auth context the default binding derives no principal."""

assert authenticated_principal(_bare_context()) is None

@pytest.mark.parametrize(

("token", "expected"),

[

pytest.param(

AccessToken(token="at-1", client_id="client-123", scopes=[]),

'["client-123",null,null]',

id="client-only",

),

pytest.param(

AccessToken(token="at-2", client_id="client-123", scopes=[], subject="alice"),

'["client-123",null,"alice"]',

id="with-subject",

),

pytest.param(

AccessToken(

token="at-3", client_id="client-123", scopes=[], subject="alice", claims={"iss": "https://as.example"}

),

'["client-123","https://as.example","alice"]',

id="with-issuer-and-subject",

),

],

)

def test_authenticated_principal_is_the_tokens_client_issuer_subject_identity(

token: AccessToken, expected: str

) -> None:

"""SDK-defined: the default binding composes (client_id, issuer, subject), degrading per component."""

reset = auth_context_var.set(AuthenticatedUser(token))

try:

assert authenticated_principal(_bare_context()) == expected

finally:

auth_context_var.reset(reset)

def test_authenticated_principal_distinguishes_two_subjects_of_one_client() -> None:

"""SDK-defined: two users of the same OAuth client are distinct principals when subjects are supplied."""

alice = AccessToken(token="at-a", client_id="https://agent.example/client.json", scopes=[], subject="alice")

bob = AccessToken(token="at-b", client_id="https://agent.example/client.json", scopes=[], subject="bob")

principals: list[str | None] = []

for token in (alice, bob):

reset = auth_context_var.set(AuthenticatedUser(token))

try:

principals.append(authenticated_principal(_bare_context()))

finally:

auth_context_var.reset(reset)

assert principals[0] != principals[1]

def test_authenticated_principal_uses_the_same_components_as_session_ownership() -> None:

"""SDK-defined: the binding and authorization_context derive from one principal_components source."""

token = AccessToken(

token="at-1", client_id="client-123", scopes=[], subject="alice", claims={"iss": "https://as.example"}

)

assert authorization_context(AuthenticatedUser(token)) == {

"client_id": "client-123",

"issuer": "https://as.example",

"subject": "alice",

}

assert list(principal_components(token)) == ["client-123", "https://as.example", "alice"]