From 31c456837af7065b3374bed23309855bd8b9f2b4 Mon Sep 17 00:00:00 2001 From: yucheng Date: Thu, 24 Sep 2026 08:17:45 +0000 Subject: [PATCH] refactor(mcp): single-assignment kdf fallback helpers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../mcp_server/gateway_dcr_flow.py | 20 +---- .../bridge_credentials.py | 37 ++++++-- .../outbound_credentials/key_derivation.py | 4 +- .../session_credentials.py | 34 +++++--- .../test_bridge_credentials.py | 87 +++++++++++-------- .../test_session_credentials.py | 85 ++++++++---------- 6 files changed, 142 insertions(+), 125 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py index ab922a67b2d..508c6f51593 100644 --- a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py +++ b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py @@ -68,6 +68,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credent SessionSigningConfigError, active_session_signing_keys, legacy_session_keys_from_master_key, + open_session_credential_with_legacy, open_session_refresh_bearer, ) from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import ( @@ -76,11 +77,9 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token i MintedSessionToken, OpenedSessionToken, SessionAudience, - SessionExpired, SessionKeys, SessionPrincipal, SessionSigningKeys, - SessionTokenOpenError, is_session_refresh_token, is_session_token, mint_session_refresh_token, @@ -1483,19 +1482,6 @@ async def revoke_refresh_token(token: str, client_id: str, master_key: str | Non return Response(content="{}", media_type="application/json", headers=TOKEN_NO_CACHE_HEADERS) -def _open_with_legacy( - open_token: Callable[[str, SessionSigningKeys, datetime], OpenedSessionToken | SessionTokenOpenError], - token: str, - keys: SessionSigningKeys, - legacy_keys: SessionKeys | None, - now: datetime, -) -> OpenedSessionToken | SessionTokenOpenError: - opened: Final = open_token(token, keys, now) - if isinstance(opened, OpenedSessionToken) or isinstance(opened, SessionExpired) or legacy_keys is None: - return opened - return open_token(token, legacy_keys, now) - - def _inactive_introspection_response() -> Response: """RFC 7662 section 2.2: any token the gateway cannot vouch for, whatever the reason (wrong family, bad signature, expired, revoked, or a deactivated user), answers 200 @@ -1556,9 +1542,9 @@ async def introspect_gateway_token( now: Final = datetime.now(timezone.utc) legacy_keys: Final = legacy_session_keys_from_master_key(master_key, keys) if is_session_token(token): - opened = _open_with_legacy(open_session_token, token, keys, legacy_keys, now) + opened = open_session_credential_with_legacy(open_session_token, token, keys, legacy_keys, now) elif is_session_refresh_token(token): - opened = _open_with_legacy(open_session_refresh_token, token, keys, legacy_keys, now) + opened = open_session_credential_with_legacy(open_session_refresh_token, token, keys, legacy_keys, now) else: return _inactive_introspection_response() if not isinstance(opened, OpenedSessionToken): diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/bridge_credentials.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/bridge_credentials.py index 749813df4d0..fb5a2e38236 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/bridge_credentials.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/bridge_credentials.py @@ -21,6 +21,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import EnvelopeIdentity, EnvelopeKeys, EnvelopeMintError, + EnvelopeOpenError, OpenedEnvelope, OpenedRefreshEnvelope, RefreshCredential, @@ -131,6 +132,34 @@ class BridgeRefreshInvalid(BaseModel): BridgeRefreshResult: TypeAlias = BridgeRefreshOpened | BridgeRefreshInvalid +def _open_envelope_with_fallback( + candidate: str, + keys: EnvelopeKeys, + legacy_keys: EnvelopeKeys | None, + now: datetime, +) -> OpenedEnvelope | EnvelopeOpenError: + """Open an access envelope under the active keys, then under the legacy keys when + grace supplied them.""" + primary: Final = open_envelope(candidate, keys, now) + if isinstance(primary, OpenedEnvelope) or legacy_keys is None: + return primary + return open_envelope(candidate, legacy_keys, now) + + +def _open_refresh_envelope_with_fallback( + candidate: str, + keys: EnvelopeKeys, + legacy_keys: EnvelopeKeys | None, + now: datetime, +) -> OpenedRefreshEnvelope | EnvelopeOpenError: + """Open a refresh envelope under the active keys, then under the legacy keys when + grace supplied them.""" + primary: Final = open_refresh_envelope(candidate, keys, now) + if isinstance(primary, OpenedRefreshEnvelope) or legacy_keys is None: + return primary + return open_refresh_envelope(candidate, legacy_keys, now) + + def open_bridge_refresh_envelope( refresh_value: str, keys: EnvelopeKeys, @@ -151,9 +180,7 @@ def open_bridge_refresh_envelope( candidate: Final = _strip_bearer(refresh_value) if not is_refresh_envelope(candidate): return BridgeRefreshInvalid() - opened = open_refresh_envelope(candidate, keys, now) - if not isinstance(opened, OpenedRefreshEnvelope) and legacy_keys is not None: - opened = open_refresh_envelope(candidate, legacy_keys, now) + opened: Final = _open_refresh_envelope_with_fallback(candidate, keys, legacy_keys, now) if not isinstance(opened, OpenedRefreshEnvelope): return BridgeRefreshInvalid() if opened.identity.server_id != expected_server_id: @@ -237,9 +264,7 @@ def resolve_bridge_envelope( return BridgeEnvelopeInvalid() if not is_envelope(candidate): return NotBridgeEnvelope() - opened = open_envelope(candidate, keys, now) - if not isinstance(opened, OpenedEnvelope) and legacy_keys is not None: - opened = open_envelope(candidate, legacy_keys, now) + opened: Final = _open_envelope_with_fallback(candidate, keys, legacy_keys, now) if not isinstance(opened, OpenedEnvelope): return BridgeEnvelopeInvalid() if opened.identity.server_id != expected_server_id: diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/key_derivation.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/key_derivation.py index 98c2218f2eb..9ea90386bc8 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/key_derivation.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/key_derivation.py @@ -19,12 +19,10 @@ from litellm.proxy.common_utils.fips import is_fips_mode LEGACY_KDF_GRACE_ENV_VAR: Final = "LITELLM_MCP_LEGACY_KDF_GRACE" DERIVED_KEY_BYTES: Final = 32 -# scrypt work factors (RFC 7914). n=2**15 with r=8/p=1 costs ~50ms and ~32MB per derivation. +# RFC 7914 work factors of the pre-rotation derivation; maxmem is twice the working set. _SCRYPT_N: Final = 2**15 _SCRYPT_R: Final = 8 _SCRYPT_P: Final = 1 -# scrypt's working-set is ~128 * N * r * p bytes; cap at twice that so the maxmem ceiling scales -# with every work factor and a future p or r bump does not trip "memory limit exceeded". _SCRYPT_MAXMEM: Final = 128 * _SCRYPT_N * _SCRYPT_R * _SCRYPT_P * 2 diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_credentials.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_credentials.py index 02b231fed27..396a821b0fa 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_credentials.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_credentials.py @@ -36,6 +36,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token i SessionPrincipal, SessionRotatedPublicKey, SessionSigningKeys, + SessionTokenOpenError, is_session_refresh_token, is_session_token, open_session_refresh_token, @@ -249,13 +250,7 @@ def resolve_session_bearer( return SessionBearerInvalid() if not is_session_token(candidate): return NotSessionBearer() - opened = open_session_token(candidate, keys, now) - if ( - not isinstance(opened, OpenedSessionToken) - and not isinstance(opened, SessionExpired) - and legacy_keys is not None - ): - opened = open_session_token(candidate, legacy_keys, now) + opened: Final = open_session_credential_with_legacy(open_session_token, candidate, keys, legacy_keys, now) if isinstance(opened, OpenedSessionToken): return SessionBearerAdmitted(principal=opened.principal) return SessionBearerInvalid(expired=isinstance(opened, SessionExpired)) @@ -283,6 +278,23 @@ class SessionRefreshInvalid(BaseModel): SessionRefreshResult: TypeAlias = SessionRefreshOpened | SessionRefreshInvalid +def open_session_credential_with_legacy( + open_token: Callable[[str, SessionSigningKeys, datetime], OpenedSessionToken | SessionTokenOpenError], + token: str, + keys: SessionSigningKeys, + legacy_keys: SessionKeys | None, + now: datetime, +) -> OpenedSessionToken | SessionTokenOpenError: + """Open a session credential under the active keys, then under the legacy keys when + grace supplied them. An :class:`OpenedSessionToken` or :class:`SessionExpired` under + the active keys is final: expiry must still report as expiry, so only a signature or + format failure falls through to the legacy derivation.""" + primary: Final = open_token(token, keys, now) + if isinstance(primary, (OpenedSessionToken, SessionExpired)) or legacy_keys is None: + return primary + return open_token(token, legacy_keys, now) + + def open_session_refresh_bearer( refresh_value: str, keys: SessionSigningKeys, @@ -303,13 +315,7 @@ def open_session_refresh_bearer( candidate: Final = _strip_bearer(refresh_value) if not is_session_refresh_token(candidate): return SessionRefreshInvalid() - opened = open_session_refresh_token(candidate, keys, now) - if ( - not isinstance(opened, OpenedSessionToken) - and not isinstance(opened, SessionExpired) - and legacy_keys is not None - ): - opened = open_session_refresh_token(candidate, legacy_keys, now) + opened: Final = open_session_credential_with_legacy(open_session_refresh_token, candidate, keys, legacy_keys, now) if not isinstance(opened, OpenedSessionToken): return SessionRefreshInvalid() if opened.principal.client_id != expected_client_id: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_bridge_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_bridge_credentials.py index 6718285ef0a..673c418041b 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_bridge_credentials.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_bridge_credentials.py @@ -10,7 +10,10 @@ through the consumer; and no path leaks the upstream token in a repr. import hashlib import hmac +from collections.abc import Callable, Mapping from datetime import datetime, timedelta, timezone +from types import MappingProxyType +from typing import Final import pytest from pydantic import SecretStr @@ -29,13 +32,6 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credenti open_bridge_refresh_envelope, resolve_bridge_envelope, ) -from litellm.proxy._experimental.mcp_server.outbound_credentials.key_derivation import ( - LEGACY_KDF_GRACE_ENV_VAR, -) -from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credentials import ( - session_keys_from_master_key, -) -from litellm.proxy.common_utils.fips import FIPS_MODE_ENV_VAR from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( ENVELOPE_PREFIX, EnvelopeIdentity, @@ -47,6 +43,13 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import key_hash_identity, mint_envelope, ) +from litellm.proxy._experimental.mcp_server.outbound_credentials.key_derivation import ( + LEGACY_KDF_GRACE_ENV_VAR, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credentials import ( + session_keys_from_master_key, +) +from litellm.proxy.common_utils.fips import FIPS_MODE_ENV_VAR _NOW = datetime(2026, 7, 9, 12, 0, 0, tzinfo=timezone.utc) _MASTER_KEY = "sk-master-key-for-derivation-tests-0123456789" @@ -332,11 +335,11 @@ def test_is_bridge_envelope_shaped_rejects_non_envelope_bearer(): def _rfc5869(ikm: bytes, info: bytes) -> bytes: """Independent RFC 5869 HKDF-SHA256, hand-written on hmac/hashlib so a wrong KDF construction in the product cannot agree with it.""" - prk = hmac.new(b"\x00" * 32, ikm, hashlib.sha256).digest() + prk: Final = hmac.new(b"\x00" * 32, ikm, hashlib.sha256).digest() return hmac.new(prk, info + b"\x01", hashlib.sha256).digest()[:32] -def _environ(values: dict[str, str]): +def _environ(values: Mapping[str, str]) -> Callable[[str], str | None]: return values.get @@ -353,7 +356,7 @@ def _legacy_keys(master_key: str) -> EnvelopeKeys: def test_envelope_key_derivation_matches_an_independent_rfc5869_vector(): - keys = envelope_keys_from_master_key("sk-1234") + keys: Final = envelope_keys_from_master_key("sk-1234") assert keys.signing_key.get_secret_value() == _rfc5869(b"sk-1234", b"litellm-mcp-bridge:envelope-signing:").hex() assert ( keys.encryption_key.get_secret_value() == _rfc5869(b"sk-1234", b"litellm-mcp-bridge:envelope-encryption:").hex() @@ -363,34 +366,38 @@ def test_envelope_key_derivation_matches_an_independent_rfc5869_vector(): def test_session_key_derivation_matches_an_independent_rfc5869_vector(): - keys = session_keys_from_master_key("sk-1234") + keys: Final = session_keys_from_master_key("sk-1234") assert keys.signing_key.get_secret_value() == _rfc5869(b"sk-1234", b"litellm-mcp-gateway:session-signing:").hex() assert keys.signing_key.get_secret_value() == "511f07b16ceeb5fba1660033785d4b9a4a6e80c749a51d639a688c287671a0e6" def test_the_three_derived_keys_are_pairwise_distinct(): - envelope = envelope_keys_from_master_key(_MASTER_KEY) - session = session_keys_from_master_key(_MASTER_KEY) - derived = { - envelope.signing_key.get_secret_value(), - envelope.encryption_key.get_secret_value(), - session.signing_key.get_secret_value(), - } + envelope: Final = envelope_keys_from_master_key(_MASTER_KEY) + session: Final = session_keys_from_master_key(_MASTER_KEY) + derived: Final = frozenset( + ( + envelope.signing_key.get_secret_value(), + envelope.encryption_key.get_secret_value(), + session.signing_key.get_secret_value(), + ) + ) assert len(derived) == 3 def test_new_envelopes_open_under_the_hkdf_keys_with_no_legacy_fallback(): - keys = envelope_keys_from_master_key(_MASTER_KEY) - result = resolve_bridge_envelope(_sealed_token(keys), keys, _NOW, _SERVER_ID) + keys: Final = envelope_keys_from_master_key(_MASTER_KEY) + result: Final = resolve_bridge_envelope(_sealed_token(keys), keys, _NOW, _SERVER_ID) assert isinstance(result, BridgeEnvelopeAdmitted) def test_legacy_envelope_opens_during_the_grace_window(): - legacy = _legacy_keys(_MASTER_KEY) - token = _sealed_token(legacy) - grace = legacy_envelope_keys_from_master_key(_MASTER_KEY, environ=_environ({LEGACY_KDF_GRACE_ENV_VAR: "true"})) + legacy: Final = _legacy_keys(_MASTER_KEY) + token: Final = _sealed_token(legacy) + grace: Final = legacy_envelope_keys_from_master_key( + _MASTER_KEY, environ=_environ(MappingProxyType({LEGACY_KDF_GRACE_ENV_VAR: "true"})) + ) assert grace == legacy - result = resolve_bridge_envelope( + result: Final = resolve_bridge_envelope( token, envelope_keys_from_master_key(_MASTER_KEY), _NOW, _SERVER_ID, legacy_keys=grace ) assert isinstance(result, BridgeEnvelopeAdmitted) @@ -398,32 +405,36 @@ def test_legacy_envelope_opens_during_the_grace_window(): def test_legacy_envelope_is_invalid_without_grace(): - token = _sealed_token(_legacy_keys(_MASTER_KEY)) - assert legacy_envelope_keys_from_master_key(_MASTER_KEY, environ=_environ({})) is None - result = resolve_bridge_envelope(token, envelope_keys_from_master_key(_MASTER_KEY), _NOW, _SERVER_ID) + token: Final = _sealed_token(_legacy_keys(_MASTER_KEY)) + assert legacy_envelope_keys_from_master_key(_MASTER_KEY, environ=_environ(MappingProxyType({}))) is None + result: Final = resolve_bridge_envelope(token, envelope_keys_from_master_key(_MASTER_KEY), _NOW, _SERVER_ID) assert isinstance(result, BridgeEnvelopeInvalid) def test_grace_is_disabled_under_fips(): - grace = legacy_envelope_keys_from_master_key( + grace: Final = legacy_envelope_keys_from_master_key( _MASTER_KEY, - environ=_environ({LEGACY_KDF_GRACE_ENV_VAR: "true", FIPS_MODE_ENV_VAR: "true"}), + environ=_environ(MappingProxyType({LEGACY_KDF_GRACE_ENV_VAR: "true", FIPS_MODE_ENV_VAR: "true"})), ) assert grace is None -@pytest.mark.parametrize("value", ["yes", "1", "on", "TRUE ", "True"]) +@pytest.mark.parametrize("value", ("yes", "1", "on", "TRUE ", "True")) def test_grace_is_enabled_only_by_case_insensitive_true(value: str): - result = legacy_envelope_keys_from_master_key(_MASTER_KEY, environ=_environ({LEGACY_KDF_GRACE_ENV_VAR: value})) - expected = value.strip().lower() == "true" + result: Final = legacy_envelope_keys_from_master_key( + _MASTER_KEY, environ=_environ(MappingProxyType({LEGACY_KDF_GRACE_ENV_VAR: value})) + ) + expected: Final = value.strip().lower() == "true" assert (result is not None) is expected def test_legacy_refresh_envelope_opens_during_grace(): - legacy = _legacy_keys(_MASTER_KEY) - refresh = _sealed_refresh(legacy) - grace = legacy_envelope_keys_from_master_key(_MASTER_KEY, environ=_environ({LEGACY_KDF_GRACE_ENV_VAR: "true"})) - result = open_bridge_refresh_envelope( + legacy: Final = _legacy_keys(_MASTER_KEY) + refresh: Final = _sealed_refresh(legacy) + grace: Final = legacy_envelope_keys_from_master_key( + _MASTER_KEY, environ=_environ(MappingProxyType({LEGACY_KDF_GRACE_ENV_VAR: "true"})) + ) + result: Final = open_bridge_refresh_envelope( refresh, envelope_keys_from_master_key(_MASTER_KEY), _NOW, _SERVER_ID, legacy_keys=grace ) assert isinstance(result, BridgeRefreshOpened) @@ -431,6 +442,6 @@ def test_legacy_refresh_envelope_opens_during_grace(): def test_legacy_refresh_envelope_is_invalid_without_grace(): - refresh = _sealed_refresh(_legacy_keys(_MASTER_KEY)) - result = open_bridge_refresh_envelope(refresh, envelope_keys_from_master_key(_MASTER_KEY), _NOW, _SERVER_ID) + refresh: Final = _sealed_refresh(_legacy_keys(_MASTER_KEY)) + result: Final = open_bridge_refresh_envelope(refresh, envelope_keys_from_master_key(_MASTER_KEY), _NOW, _SERVER_ID) assert isinstance(result, BridgeRefreshInvalid) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_credentials.py index 396e7b965e1..13886982b37 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_credentials.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_credentials.py @@ -1,7 +1,10 @@ """Tests for the session-token KDF and the edge/token-endpoint resolvers.""" import hashlib +from collections.abc import Callable, Mapping from datetime import datetime, timedelta, timezone +from types import MappingProxyType +from typing import Final import pytest from cryptography.hazmat.primitives import serialization @@ -12,7 +15,9 @@ from pydantic import SecretStr from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( envelope_keys_from_master_key, ) -from litellm.proxy._experimental.mcp_server.outbound_credentials import session_credentials +from litellm.proxy._experimental.mcp_server.outbound_credentials.key_derivation import ( + LEGACY_KDF_GRACE_ENV_VAR, +) from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credentials import ( NotSessionBearer, SessionBearerAdmitted, @@ -27,9 +32,6 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credent resolve_session_signing_keys, session_keys_from_master_key, ) -from litellm.proxy._experimental.mcp_server.outbound_credentials.key_derivation import ( - LEGACY_KDF_GRACE_ENV_VAR, -) from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import ( SESSION_TTL_SECONDS, AsymmetricSessionKeys, @@ -239,12 +241,12 @@ def test_signing_config_error_detail_never_leaks_key_material(): assert pem.splitlines()[1] not in resolved.detail -def _environ(values: dict[str, str]): +def _environ(values: Mapping[str, str]) -> Callable[[str], str | None]: return values.get def _legacy_session_keys(master_key: str) -> SessionKeys: - signing = hashlib.scrypt( + signing: Final = hashlib.scrypt( master_key.encode(), salt=b"litellm-mcp-gateway:session-signing:", n=2**15, @@ -257,53 +259,57 @@ def _legacy_session_keys(master_key: str) -> SessionKeys: def test_new_session_tokens_open_under_the_hkdf_key_with_no_legacy_fallback(): - result = resolve_session_bearer(f"Bearer {_access_token()}", KEYS, NOW) + result: Final = resolve_session_bearer(f"Bearer {_access_token()}", KEYS, NOW) assert isinstance(result, SessionBearerAdmitted) def test_legacy_session_token_opens_during_the_grace_window(): - legacy = _legacy_session_keys(MASTER_KEY) - minted = mint_session_token(PRINCIPAL, legacy, NOW) + legacy: Final = _legacy_session_keys(MASTER_KEY) + minted: Final = mint_session_token(PRINCIPAL, legacy, NOW) assert isinstance(minted, MintedSessionToken) - grace = legacy_session_keys_from_master_key(MASTER_KEY, KEYS, environ=_environ({LEGACY_KDF_GRACE_ENV_VAR: "true"})) + grace: Final = legacy_session_keys_from_master_key( + MASTER_KEY, KEYS, environ=_environ(MappingProxyType({LEGACY_KDF_GRACE_ENV_VAR: "true"})) + ) assert grace == legacy - result = resolve_session_bearer(f"Bearer {minted.token.get_secret_value()}", KEYS, NOW, legacy_keys=grace) + result: Final = resolve_session_bearer(f"Bearer {minted.token.get_secret_value()}", KEYS, NOW, legacy_keys=grace) assert isinstance(result, SessionBearerAdmitted) assert result.principal == PRINCIPAL def test_legacy_session_token_is_invalid_without_grace(): - minted = mint_session_token(PRINCIPAL, _legacy_session_keys(MASTER_KEY), NOW) + minted: Final = mint_session_token(PRINCIPAL, _legacy_session_keys(MASTER_KEY), NOW) assert isinstance(minted, MintedSessionToken) - assert legacy_session_keys_from_master_key(MASTER_KEY, KEYS, environ=_environ({})) is None - result = resolve_session_bearer(f"Bearer {minted.token.get_secret_value()}", KEYS, NOW) + assert legacy_session_keys_from_master_key(MASTER_KEY, KEYS, environ=_environ(MappingProxyType({}))) is None + result: Final = resolve_session_bearer(f"Bearer {minted.token.get_secret_value()}", KEYS, NOW) assert isinstance(result, SessionBearerInvalid) assert result.expired is False def test_legacy_session_key_is_never_derived_for_configured_rs256_keys(): - active = AsymmetricSessionKeys(private_key_pem=SecretStr(_rsa_private_pem()), kid="k") - grace = legacy_session_keys_from_master_key( - MASTER_KEY, active, environ=_environ({LEGACY_KDF_GRACE_ENV_VAR: "true"}) + active: Final = AsymmetricSessionKeys(private_key_pem=SecretStr(_rsa_private_pem()), kid="k") + grace: Final = legacy_session_keys_from_master_key( + MASTER_KEY, active, environ=_environ(MappingProxyType({LEGACY_KDF_GRACE_ENV_VAR: "true"})) ) assert grace is None def test_grace_is_disabled_under_fips(): - grace = legacy_session_keys_from_master_key( + grace: Final = legacy_session_keys_from_master_key( MASTER_KEY, KEYS, - environ=_environ({LEGACY_KDF_GRACE_ENV_VAR: "true", FIPS_MODE_ENV_VAR: "true"}), + environ=_environ(MappingProxyType({LEGACY_KDF_GRACE_ENV_VAR: "true", FIPS_MODE_ENV_VAR: "true"})), ) assert grace is None def test_legacy_session_refresh_token_opens_during_grace(): - legacy = _legacy_session_keys(MASTER_KEY) - minted = mint_session_refresh_token(PRINCIPAL, legacy, NOW) + legacy: Final = _legacy_session_keys(MASTER_KEY) + minted: Final = mint_session_refresh_token(PRINCIPAL, legacy, NOW) assert isinstance(minted, MintedSessionToken) - grace = legacy_session_keys_from_master_key(MASTER_KEY, KEYS, environ=_environ({LEGACY_KDF_GRACE_ENV_VAR: "true"})) - result = open_session_refresh_bearer( + grace: Final = legacy_session_keys_from_master_key( + MASTER_KEY, KEYS, environ=_environ(MappingProxyType({LEGACY_KDF_GRACE_ENV_VAR: "true"})) + ) + result: Final = open_session_refresh_bearer( minted.token.get_secret_value(), KEYS, NOW, @@ -315,37 +321,22 @@ def test_legacy_session_refresh_token_opens_during_grace(): def test_legacy_session_refresh_token_is_invalid_without_grace(): - minted = mint_session_refresh_token(PRINCIPAL, _legacy_session_keys(MASTER_KEY), NOW) + minted: Final = mint_session_refresh_token(PRINCIPAL, _legacy_session_keys(MASTER_KEY), NOW) assert isinstance(minted, MintedSessionToken) - result = open_session_refresh_bearer( + result: Final = open_session_refresh_bearer( minted.token.get_secret_value(), KEYS, NOW, expected_client_id="llm_client_abc" ) assert isinstance(result, SessionRefreshInvalid) def test_expired_token_under_the_primary_key_still_reports_expired_during_grace(): - legacy = _legacy_session_keys(MASTER_KEY) - minted = mint_session_token(PRINCIPAL, legacy, NOW) + legacy: Final = _legacy_session_keys(MASTER_KEY) + minted: Final = mint_session_token(PRINCIPAL, legacy, NOW) assert isinstance(minted, MintedSessionToken) - grace = legacy_session_keys_from_master_key(MASTER_KEY, KEYS, environ=_environ({LEGACY_KDF_GRACE_ENV_VAR: "true"})) - later = NOW + timedelta(seconds=SESSION_TTL_SECONDS + 1) - result = resolve_session_bearer(f"Bearer {minted.token.get_secret_value()}", KEYS, later, legacy_keys=grace) + grace: Final = legacy_session_keys_from_master_key( + MASTER_KEY, KEYS, environ=_environ(MappingProxyType({LEGACY_KDF_GRACE_ENV_VAR: "true"})) + ) + later: Final = NOW + timedelta(seconds=SESSION_TTL_SECONDS + 1) + result: Final = resolve_session_bearer(f"Bearer {minted.token.get_secret_value()}", KEYS, later, legacy_keys=grace) assert isinstance(result, SessionBearerInvalid) assert result.expired is True - - -def test_fallback_tries_the_active_key_before_legacy(monkeypatch): - calls: list[object] = [] - real = session_credentials.open_session_token - - def spy(candidate, keys, now): - calls.append(keys) - return real(candidate, keys, now) - - monkeypatch.setattr(session_credentials, "open_session_token", spy) - grace = legacy_session_keys_from_master_key(MASTER_KEY, KEYS, environ=_environ({LEGACY_KDF_GRACE_ENV_VAR: "true"})) - minted = mint_session_token(PRINCIPAL, _legacy_session_keys(MASTER_KEY), NOW) - assert isinstance(minted, MintedSessionToken) - result = resolve_session_bearer(f"Bearer {minted.token.get_secret_value()}", KEYS, NOW, legacy_keys=grace) - assert isinstance(result, SessionBearerAdmitted) - assert calls[0] is KEYS