Require integrity protection for MRTR requestState (#3032)
This commit is contained in:
Max
GitHub
parent
985652491a
commit
7322ca56f4
8 changed files with 3338 additions and 183 deletions
|
|
@ -18,9 +18,10 @@ from mcp_types import (
|
|||
TextContent,
|
||||
)
|
||||
|
||||
from docs_src.mrtr import tutorial001, tutorial002, tutorial003, tutorial004
|
||||
from docs_src.mrtr import tutorial001, tutorial002, tutorial003, tutorial004, tutorial005
|
||||
from mcp import Client, MCPError
|
||||
from mcp.client import ClientRequestContext
|
||||
from mcp.server.mcpserver import InvalidRequestState
|
||||
|
||||
# See test_index.py for why this is a per-module mark and not a conftest hook.
|
||||
pytestmark = [pytest.mark.anyio, pytest.mark.filterwarnings("error::mcp.MCPDeprecationWarning")]
|
||||
|
|
@ -161,3 +162,36 @@ async def test_the_prompt_auto_loop_returns_the_final_messages() -> None:
|
|||
],
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def test_a_custom_codec_round_trips_what_it_sealed() -> None:
|
||||
"""tutorial005: `unseal(seal(payload))` returns the payload; the token itself is opaque hex."""
|
||||
codec = tutorial005.EnvelopeCodec(tutorial005.unwrap_data_key())
|
||||
token = codec.seal(b"round-1")
|
||||
assert token.startswith(tutorial005.PREFIX)
|
||||
assert b"round-1" not in token.encode()
|
||||
assert codec.unseal(token) == b"round-1"
|
||||
|
||||
|
||||
def test_a_custom_codec_raises_invalid_request_state_for_any_bad_token() -> None:
|
||||
"""tutorial005: any token the codec did not mint intact raises `InvalidRequestState`."""
|
||||
codec = tutorial005.EnvelopeCodec(tutorial005.unwrap_data_key())
|
||||
token = codec.seal(b"round-1")
|
||||
with pytest.raises(InvalidRequestState):
|
||||
codec.unseal(token + "00")
|
||||
with pytest.raises(InvalidRequestState):
|
||||
codec.unseal("not-a-token")
|
||||
|
||||
|
||||
def test_a_custom_codec_rejects_every_alias_of_a_minted_token() -> None:
|
||||
"""tutorial005: only the exact minted string verifies; rewritten spellings of it do not."""
|
||||
codec = tutorial005.EnvelopeCodec(tutorial005.unwrap_data_key())
|
||||
token = codec.seal(b"round-1")
|
||||
body = token.removeprefix(tutorial005.PREFIX)
|
||||
for alias in (
|
||||
body, # prefix stripped
|
||||
tutorial005.PREFIX + body.upper(), # non-canonical hex case
|
||||
tutorial005.PREFIX + body[:8] + " " + body[8:], # whitespace bytes.fromhex would skip
|
||||
):
|
||||
with pytest.raises(InvalidRequestState):
|
||||
codec.unseal(alias)
|
||||
|
|
|
|||
|
|
@ -9,8 +9,18 @@ from starlette.datastructures import Headers
|
|||
from starlette.requests import Request
|
||||
from starlette.types import Message, Receive, Scope, Send
|
||||
|
||||
from mcp.server.auth.middleware.bearer_auth import AuthenticatedUser, BearerAuthBackend, RequireAuthMiddleware
|
||||
from mcp.server.auth.provider import AccessToken, OAuthAuthorizationServerProvider, ProviderTokenVerifier
|
||||
from mcp.server.auth.middleware.bearer_auth import (
|
||||
AuthenticatedUser,
|
||||
BearerAuthBackend,
|
||||
RequireAuthMiddleware,
|
||||
authorization_context,
|
||||
)
|
||||
from mcp.server.auth.provider import (
|
||||
AccessToken,
|
||||
OAuthAuthorizationServerProvider,
|
||||
ProviderTokenVerifier,
|
||||
principal_components,
|
||||
)
|
||||
|
||||
|
||||
class MockOAuthProvider:
|
||||
|
|
@ -446,3 +456,16 @@ class TestRequireAuthMiddleware:
|
|||
assert app.scope == scope
|
||||
assert app.receive == receive
|
||||
assert app.send == send
|
||||
|
||||
|
||||
def test_authorization_context_is_built_from_principal_components() -> None:
|
||||
"""Session ownership identifies the principal via the shared principal_components triple."""
|
||||
token = AccessToken(
|
||||
token="t", client_id="client-1", scopes=[], subject="alice", claims={"iss": "https://as.example"}
|
||||
)
|
||||
client_id, issuer, subject = principal_components(token)
|
||||
assert authorization_context(AuthenticatedUser(token)) == {
|
||||
"client_id": client_id,
|
||||
"issuer": issuer,
|
||||
"subject": subject,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
"""Tests for mcp.server.auth.provider module."""
|
||||
|
||||
from mcp.server.auth.provider import construct_redirect_uri
|
||||
from mcp.server.auth.provider import AccessToken, construct_redirect_uri, principal_components
|
||||
|
||||
|
||||
def test_construct_redirect_uri_no_existing_params():
|
||||
|
|
@ -77,3 +77,14 @@ def test_construct_redirect_uri_encoded_values():
|
|||
|
||||
# urlencode uses + for spaces by default
|
||||
assert "state=test+state+with+spaces" in result
|
||||
|
||||
|
||||
def test_principal_components_composes_client_issuer_subject():
|
||||
"""The triple identifying a token's principal, degrading per missing component."""
|
||||
bare = AccessToken(token="t", client_id="client-1", scopes=[])
|
||||
assert principal_components(bare) == ("client-1", None, None)
|
||||
|
||||
full = AccessToken(
|
||||
token="t", client_id="client-1", scopes=[], subject="alice", claims={"iss": "https://as.example"}
|
||||
)
|
||||
assert principal_components(full) == ("client-1", "https://as.example", "alice")
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -1866,7 +1866,8 @@ async def test_read_resource_template_not_found():
|
|||
assert exc_info.value.error.data == {"uri": "resource://users/999"}
|
||||
|
||||
|
||||
async def test_tool_returning_input_required_result_reaches_client_unchanged():
|
||||
async def test_tool_returning_input_required_result_reaches_client_sealed():
|
||||
# Default posture: the wire carries an opaque sealed token, never the handler's plaintext.
|
||||
mcp = MCPServer()
|
||||
|
||||
@mcp.tool()
|
||||
|
|
@ -1878,7 +1879,7 @@ async def test_tool_returning_input_required_result_reaches_client_unchanged():
|
|||
result = await client.session.call_tool("ask", allow_input_required=True)
|
||||
|
||||
assert isinstance(result, InputRequiredResult)
|
||||
assert result.request_state == "round-1"
|
||||
_assert_sealed(result.request_state, "round-1")
|
||||
assert result.input_requests is not None
|
||||
assert result.input_requests["roots"].method == "roots/list"
|
||||
|
||||
|
|
@ -1927,6 +1928,13 @@ async def test_tool_reads_input_responses_and_request_state_from_context_on_retr
|
|||
assert block.text == "Hello, Alice! (state=r1)"
|
||||
|
||||
|
||||
def _assert_sealed(state: str | None, plaintext: str) -> None:
|
||||
"""The wire form is an opaque sealed token, never the handler's plaintext."""
|
||||
assert state is not None
|
||||
assert state != plaintext
|
||||
assert state.startswith("v1.")
|
||||
|
||||
|
||||
def _ask_who() -> ElicitRequest:
|
||||
return ElicitRequest(
|
||||
params=ElicitRequestFormParams(
|
||||
|
|
@ -1940,9 +1948,9 @@ def _ask_who() -> ElicitRequest:
|
|||
)
|
||||
|
||||
|
||||
async def test_prompt_returning_input_required_result_reaches_client_unchanged():
|
||||
"""A prompt function may return an InputRequiredResult and the pipeline passes it
|
||||
through to the client (spec-mandated: SEP-2322 allows it on prompts/get)."""
|
||||
async def test_prompt_returning_input_required_result_reaches_client_sealed():
|
||||
"""A prompt function may return an InputRequiredResult and the pipeline delivers it
|
||||
to the client with the state sealed (spec-mandated: SEP-2322 allows it on prompts/get)."""
|
||||
mcp = MCPServer()
|
||||
|
||||
@mcp.prompt()
|
||||
|
|
@ -1954,7 +1962,7 @@ async def test_prompt_returning_input_required_result_reaches_client_unchanged()
|
|||
result = await client.session.get_prompt("briefing", allow_input_required=True)
|
||||
|
||||
assert isinstance(result, InputRequiredResult)
|
||||
assert result.request_state == "round-1"
|
||||
_assert_sealed(result.request_state, "round-1")
|
||||
assert result.input_requests is not None
|
||||
assert result.input_requests["who"].method == "elicitation/create"
|
||||
|
||||
|
|
@ -2023,9 +2031,9 @@ async def test_resource_template_input_required_result_on_legacy_session_is_a_se
|
|||
assert exc.value.error.message == "Handler returned an invalid result"
|
||||
|
||||
|
||||
async def test_resource_template_returning_input_required_result_reaches_client_unchanged():
|
||||
async def test_resource_template_returning_input_required_result_reaches_client_sealed():
|
||||
"""A resource template function may return an InputRequiredResult and the pipeline
|
||||
passes it through to the client (spec-mandated: SEP-2322 allows it on resources/read)."""
|
||||
delivers it with the state sealed (spec-mandated: SEP-2322 allows it on resources/read)."""
|
||||
mcp = MCPServer()
|
||||
|
||||
@mcp.resource("ask://{topic}")
|
||||
|
|
@ -2037,7 +2045,7 @@ async def test_resource_template_returning_input_required_result_reaches_client_
|
|||
result = await client.session.read_resource("ask://databases", allow_input_required=True)
|
||||
|
||||
assert isinstance(result, InputRequiredResult)
|
||||
assert result.request_state == "round-1"
|
||||
_assert_sealed(result.request_state, "round-1")
|
||||
assert result.input_requests is not None
|
||||
assert result.input_requests["who"].method == "elicitation/create"
|
||||
|
||||
|
|
@ -2121,22 +2129,26 @@ async def test_context_read_resource_keeps_outer_input_responses_from_the_nested
|
|||
return f"{topic} content"
|
||||
|
||||
@mcp.tool()
|
||||
async def outer(ctx: Context) -> str:
|
||||
async def outer(ctx: Context) -> str | InputRequiredResult:
|
||||
if ctx.input_responses is None:
|
||||
return InputRequiredResult(input_requests={"who": _ask_who()}, request_state="outer-state")
|
||||
contents = list(await ctx.read_resource("ask://databases"))
|
||||
assert isinstance(contents[0].content, str)
|
||||
return contents[0].content
|
||||
return f"{contents[0].content} (state={ctx.request_state})"
|
||||
|
||||
with anyio.fail_after(5):
|
||||
async with Client(mcp, mode="2026-07-28") as client:
|
||||
r1 = await client.session.call_tool("outer", allow_input_required=True)
|
||||
assert isinstance(r1, InputRequiredResult)
|
||||
result = await client.session.call_tool(
|
||||
"outer",
|
||||
input_responses={"who": ElicitResult(action="accept", content={"name": "Alice"})},
|
||||
request_state="outer-state",
|
||||
request_state=r1.request_state,
|
||||
)
|
||||
assert isinstance(result, CallToolResult)
|
||||
block = result.content[0]
|
||||
assert isinstance(block, TextContent)
|
||||
assert block.text == "databases content"
|
||||
assert block.text == "databases content (state=outer-state)"
|
||||
assert seen_responses == [None]
|
||||
assert seen_state == [None]
|
||||
|
||||
|
|
|
|||
479
tests/server/test_request_state.py
Normal file
479
tests/server/test_request_state.py
Normal file
|
|
@ -0,0 +1,479 @@
|
|||
"""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"]
|
||||
1297
tests/server/test_request_state_boundary.py
Normal file
1297
tests/server/test_request_state_boundary.py
Normal file
File diff suppressed because it is too large
Load diff
|
|
@ -553,6 +553,21 @@ def test_cacheable_methods_mirror_the_cacheable_method_literal():
|
|||
assert methods.CACHEABLE_METHODS == frozenset(get_args(methods.CacheableMethod))
|
||||
|
||||
|
||||
def test_input_required_methods_mirror_the_monolith_input_required_arms():
|
||||
"""MRTR weld: the spec's three multi-round-trip carriers are the only input_required methods."""
|
||||
assert methods.INPUT_REQUIRED_METHODS == frozenset({"prompts/get", "resources/read", "tools/call"})
|
||||
|
||||
|
||||
def test_is_input_required_matches_typed_and_wire_shapes():
|
||||
"""SDK-defined predicate: True only for the typed model and the tagged wire mapping."""
|
||||
assert methods.is_input_required(types.InputRequiredResult(request_state="s"))
|
||||
assert methods.is_input_required({"resultType": "input_required", "inputRequests": {}})
|
||||
assert not methods.is_input_required({"resultType": "complete", "content": []})
|
||||
assert not methods.is_input_required({})
|
||||
assert not methods.is_input_required(types.CallToolResult(content=[]))
|
||||
assert not methods.is_input_required(None)
|
||||
|
||||
|
||||
def test_minimal_request_bodies_parse_through_every_request_row():
|
||||
for (method, version), surface_type in methods.CLIENT_REQUESTS.items():
|
||||
parsed = methods.parse_client_request(method, version, REQUEST_PARAMS_FIXTURES[surface_type])
|
||||
|
|
|
|||
Loading…
Reference in a new issue