diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/envelope.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/envelope.py new file mode 100644 index 00000000000..298bc8d98cc --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/envelope.py @@ -0,0 +1,358 @@ +"""Client-held sealed envelope for the oauth_delegate DCR bridge. + +A DCR-bridge client holds ONE bearer that must carry BOTH a litellm identity and the +upstream OAuth grant, with zero server-side storage. The gateway token endpoint mints a +litellm-signed envelope (:func:`mint_envelope`); the MCP edge validates it, recovers the +identity claims and the inner upstream grant (:func:`open_envelope`), and forwards the +inner access token upstream. This module is pure and unwired: it imports nothing from +endpoint or edge code, reads no proxy globals, and takes all key material and the clock +as explicit parameters. + +Wire shape: ``llm_env_`` + an HS256 JWT (same signing approach as the BYOK session +bearer in ``byok_oauth_endpoints.py``). Registered claims are ``iss``/``iat``/``exp``; +custom claims are ``user_id``, ``server_id``, and ``grant``, where ``grant`` is the +upstream token grant serialized to JSON, encrypted with the repo's symmetric +encryption helpers (``encrypt_value``/``decrypt_value`` from +``encrypt_decrypt_utils`` — the same family ``encrypt_value_helper`` applies to +persisted DCR credentials), and base64url-encoded, so the inner token never appears +in plaintext anywhere in the envelope. + +Failures are values: :func:`open_envelope` returns one of the frozen +``EnvelopeOpenError`` variants (discriminated on ``tag``) for invalid, expired, +tampered, or undecryptable input, and :func:`mint_envelope` returns +``EnvelopeTooLarge`` for oversized grants. Error values carry tags and sizes only, +never token material. + +The pydantic input models reject programmer errors at construction (e.g. a +non-positive ``expires_in`` or an empty required field). :func:`open_envelope` is +additionally total over hostile, attacker-controlled input: it never raises, only +returns an ``EnvelopeOpenError``. :func:`mint_envelope` operates on a +gateway-supplied grant (an upstream IdP's UTF-8 JSON token response), so it does not +defend against non-UTF-8 field content that cannot survive JSON parsing; its only +value-typed failure is ``EnvelopeTooLarge``. +""" + +from __future__ import annotations + +import base64 +from datetime import datetime, timedelta +from typing import Literal, TypeAlias + +import jwt +from pydantic import BaseModel, ConfigDict, Field, SecretStr, ValidationError + +from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value, encrypt_value + +ENVELOPE_PREFIX = "llm_env_" +"""Marker prefix on every serialized envelope so the edge can cheaply tell an envelope +from a raw upstream token before doing any cryptography.""" + +ENVELOPE_ISSUER = "litellm-mcp-bridge" +"""``iss`` claim stamped into every envelope and required back on open.""" + +MAX_ENVELOPE_TTL_SECONDS = 3600 +"""Hard ceiling on envelope lifetime. ``exp`` is ``min(upstream expires_in, this cap)`` +(the cap alone when the upstream omits ``expires_in``), matching the 1h lifetime of the +BYOK session bearer this module's signing approach is borrowed from: a client-held +credential should never outlive a bounded window even when the upstream token does.""" + +MAX_ENVELOPE_BYTES = 12288 +"""Size cap on the final serialized envelope (prefix + JWT, in bytes). Upstream JWTs +commonly run 2-4KB; base64 plus encryption overhead roughly doubles that inside the +envelope, and common proxy/server header limits sit around 16KB total. 12288 leaves +comfortable headroom for a large upstream token while keeping the envelope safely +transmittable as a single Authorization header. Oversized grants are rejected with a +typed error, never truncated.""" + +_ENVELOPE_JWT_ALGORITHM = "HS256" + + +class EnvelopeIdentity(BaseModel): + """The litellm identity the envelope binds the inner grant to.""" + + model_config = ConfigDict(frozen=True) + user_id: str = Field(min_length=1) + server_id: str = Field(min_length=1) + + +class UpstreamTokenGrant(BaseModel): + """The upstream OAuth token response fields sealed inside the envelope. + + ``expires_in`` must be positive when present; a non-positive value is a programmer + error rejected at construction. Token fields are ``SecretStr`` so reprs never leak + them. + """ + + model_config = ConfigDict(frozen=True) + access_token: SecretStr = Field(min_length=1) + token_type: str = Field(min_length=1) + refresh_token: SecretStr | None = None + scope: str | None = None + expires_in: int | None = Field(default=None, gt=0) + + +class EnvelopeKeys(BaseModel): + """Injected key material: the HS256 signing key and the symmetric encryption key. + + ``signing_key`` must be at least 32 bytes: HS256's HMAC-SHA256 has a 256-bit + security level, RFC 7518 requires a key of at least that size, and a shorter key + makes PyJWT emit ``InsecureKeyLengthWarning``. + """ + + model_config = ConfigDict(frozen=True) + signing_key: SecretStr = Field(min_length=32) + encryption_key: SecretStr = Field(min_length=1) + + +class SealedEnvelope(BaseModel): + """A minted envelope: the client-held bearer value and when it expires.""" + + model_config = ConfigDict(frozen=True) + token: SecretStr + expires_at: datetime + + +class OpenedEnvelope(BaseModel): + """A validated envelope: the identity it was minted for and the recovered grant.""" + + model_config = ConfigDict(frozen=True) + identity: EnvelopeIdentity + grant: UpstreamTokenGrant + + +class EnvelopeTooLarge(BaseModel): + """The serialized envelope exceeded ``MAX_ENVELOPE_BYTES``; carries sizes only.""" + + model_config = ConfigDict(frozen=True) + tag: Literal["envelope_too_large"] = "envelope_too_large" + size_bytes: int + max_bytes: int + + +EnvelopeMintError: TypeAlias = EnvelopeTooLarge + + +class NotAnEnvelope(BaseModel): + """The candidate does not carry the envelope prefix.""" + + model_config = ConfigDict(frozen=True) + tag: Literal["not_an_envelope"] = "not_an_envelope" + + +class BadSignature(BaseModel): + """The JWT signature does not verify under the provided signing key.""" + + model_config = ConfigDict(frozen=True) + tag: Literal["bad_signature"] = "bad_signature" + + +class Expired(BaseModel): + """The envelope's ``exp`` is not in the future relative to the provided ``now``.""" + + model_config = ConfigDict(frozen=True) + tag: Literal["expired"] = "expired" + + +class MalformedPayload(BaseModel): + """The token is not a well-formed envelope: undecodable JWT, wrong issuer, missing + or mistyped claims, or a decrypted grant that fails validation.""" + + model_config = ConfigDict(frozen=True) + tag: Literal["malformed_payload"] = "malformed_payload" + + +class DecryptFailed(BaseModel): + """The signed ``grant`` blob could not be decrypted under the provided key.""" + + model_config = ConfigDict(frozen=True) + tag: Literal["decrypt_failed"] = "decrypt_failed" + + +EnvelopeOpenError: TypeAlias = NotAnEnvelope | BadSignature | Expired | MalformedPayload | DecryptFailed + + +class _EnvelopeClaims(BaseModel): + """Decoded-claims boundary that pins the exact shape :func:`mint_envelope` emits. + + ``user_id``/``server_id`` mirror the ``min_length`` constraints of + :class:`EnvelopeIdentity` so any claim set that validates here also constructs an + identity, keeping :func:`open_envelope` raise-free: a correctly signed JWT with an + empty identity claim fails here and maps to ``MalformedPayload``. + + ``strict`` rejects coerced types (``exp: "123"``, ``exp: 123.0``) rather than opening + on them, and ``extra="forbid"`` rejects any claim the gateway never mints (a hostile + ``nbf``/``aud``/... rides along on a re-signed token). Since PyJWT's own ``iat``/ + ``nbf``/``exp`` validators are disabled at decode (they raise on hostile claim types + and, for ``iat``/``nbf``, compare against the wall clock rather than the injected + ``now``), this model is the sole, total type gate for every registered claim. + """ + + model_config = ConfigDict(frozen=True, strict=True, extra="forbid") + iss: str + iat: int + exp: int + user_id: str = Field(min_length=1) + server_id: str = Field(min_length=1) + grant: str = Field(min_length=1) + + +class _GrantWire(BaseModel): + model_config = ConfigDict(frozen=True) + access_token: str + token_type: str + refresh_token: str | None = None + scope: str | None = None + expires_in: int | None = None + + +def is_envelope(candidate: str) -> bool: + """Cheap prefix check so the edge can route envelopes vs raw tokens without crypto.""" + return candidate.startswith(ENVELOPE_PREFIX) + + +def mint_envelope( + identity: EnvelopeIdentity, + grant: UpstreamTokenGrant, + keys: EnvelopeKeys, + now: datetime, +) -> SealedEnvelope | EnvelopeMintError: + """Seal ``grant`` for ``identity`` into a client-held envelope. + + ``exp`` is ``min(grant.expires_in, MAX_ENVELOPE_TTL_SECONDS)`` seconds from ``now`` + (the cap alone when ``expires_in`` is absent). Returns ``EnvelopeTooLarge`` when the + serialized envelope exceeds ``MAX_ENVELOPE_BYTES``. + """ + expires_at = now + timedelta(seconds=_envelope_ttl_seconds(grant.expires_in)) + claims = _EnvelopeClaims( + iss=ENVELOPE_ISSUER, + iat=int(now.timestamp()), + exp=int(expires_at.timestamp()), + user_id=identity.user_id, + server_id=identity.server_id, + grant=_encrypt_grant_blob(_grant_plaintext(grant), keys.encryption_key), + ) + token = ENVELOPE_PREFIX + jwt.encode( + claims.model_dump(), + keys.signing_key.get_secret_value(), + algorithm=_ENVELOPE_JWT_ALGORITHM, + ) + size_bytes = len(token.encode("utf-8")) + if size_bytes > MAX_ENVELOPE_BYTES: + return EnvelopeTooLarge(size_bytes=size_bytes, max_bytes=MAX_ENVELOPE_BYTES) + return SealedEnvelope(token=SecretStr(token), expires_at=expires_at) + + +def open_envelope( + candidate: str, + keys: EnvelopeKeys, + now: datetime, +) -> OpenedEnvelope | EnvelopeOpenError: + """Validate ``candidate`` and recover the identity and inner grant. + + Never raises for bad input: every invalid, expired, tampered, or undecryptable + candidate maps to a distinct ``EnvelopeOpenError`` variant. The recovered + ``grant.expires_in`` is the value the upstream reported at mint time and is not + re-derived, so it is stale by up to the envelope's lifetime; callers that need a + live remaining lifetime should use ``now`` against the upstream, not this field. + """ + if not is_envelope(candidate): + return NotAnEnvelope() + # UTF-8 byte length is never below character length, so a character count already over the + # cap rejects an oversize candidate in O(1) without encoding it; the exact byte check then + # runs only on candidates already bounded to <= MAX_ENVELOPE_BYTES characters. + if len(candidate) > MAX_ENVELOPE_BYTES: + return MalformedPayload() + if len(candidate.encode("utf-8", "surrogatepass")) > MAX_ENVELOPE_BYTES: + return MalformedPayload() + claims = _decode_claims(candidate.removeprefix(ENVELOPE_PREFIX), keys.signing_key) + if not isinstance(claims, _EnvelopeClaims): + return claims + if now.timestamp() >= claims.exp: + return Expired() + grant = _decrypt_grant(claims.grant, keys.encryption_key) + if not isinstance(grant, UpstreamTokenGrant): + return grant + return OpenedEnvelope( + identity=EnvelopeIdentity(user_id=claims.user_id, server_id=claims.server_id), + grant=grant, + ) + + +def _envelope_ttl_seconds(upstream_expires_in: int | None) -> int: + if upstream_expires_in is None: + return MAX_ENVELOPE_TTL_SECONDS + return min(upstream_expires_in, MAX_ENVELOPE_TTL_SECONDS) + + +def _grant_plaintext(grant: UpstreamTokenGrant) -> str: + wire = _GrantWire( + access_token=grant.access_token.get_secret_value(), + token_type=grant.token_type, + refresh_token=None if grant.refresh_token is None else grant.refresh_token.get_secret_value(), + scope=grant.scope, + expires_in=grant.expires_in, + ) + return wire.model_dump_json(exclude_none=True) + + +def _decode_claims( + compact: str, + signing_key: SecretStr, +) -> _EnvelopeClaims | BadSignature | MalformedPayload: + """Verify the HS256 signature and shape of an attacker-controlled compact JWT. + + ``compact`` is fully hostile and bounded to ``MAX_ENVELOPE_BYTES`` by the caller. + PyJWT's ``iat``/``nbf``/``exp`` validators are disabled: they raise on hostile claim + types and, for ``iat``/``nbf``, compare against the wall clock rather than the + injected ``now`` (``exp`` is checked by the caller against ``now``). Apart from a + signature mismatch (``BadSignature``), every decode failure is ``MalformedPayload``: + a non-UTF-8 candidate surfaces as ``UnicodeEncodeError`` (a ``ValueError``), a + non-string registered claim such as ``iss`` as a ``TypeError`` from PyJWT's claim + validators, and a wrong issuer or structurally invalid token as an + ``InvalidTokenError``. ``_EnvelopeClaims`` is the total type gate for the payload. + """ + try: + payload = jwt.decode( + compact, + signing_key.get_secret_value(), + algorithms=[_ENVELOPE_JWT_ALGORITHM], + issuer=ENVELOPE_ISSUER, + options={ + "verify_exp": False, + "verify_iat": False, + "verify_nbf": False, + "require": ["iss", "iat", "exp"], + }, + ) + except jwt.InvalidSignatureError: + return BadSignature() + except (jwt.InvalidTokenError, ValueError, TypeError): + return MalformedPayload() + try: + return _EnvelopeClaims.model_validate(payload) + except ValidationError: + return MalformedPayload() + + +def _encrypt_grant_blob(plaintext: str, encryption_key: SecretStr) -> str: + ciphertext = bytes(encrypt_value(value=plaintext, signing_key=encryption_key.get_secret_value())) + return base64.urlsafe_b64encode(ciphertext).decode("ascii") + + +def _decrypt_grant( + blob: str, + encryption_key: SecretStr, +) -> UpstreamTokenGrant | DecryptFailed | MalformedPayload: + from nacl.exceptions import CryptoError + + try: + plaintext = decrypt_value( + value=base64.urlsafe_b64decode(blob), + signing_key=encryption_key.get_secret_value(), + ) + except (CryptoError, ValueError): + return DecryptFailed() + try: + return UpstreamTokenGrant.model_validate_json(plaintext) + except ValidationError: + return MalformedPayload() diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_envelope.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_envelope.py new file mode 100644 index 00000000000..71de206aa1e --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_envelope.py @@ -0,0 +1,487 @@ +"""Spec tests for the sealed-envelope module (oauth_delegate DCR bridge). + +The envelope is the single client-held bearer carrying both a litellm identity and the +encrypted upstream grant, with zero server-side storage. These tests pin the security +contract: an envelope opens only under the exact keys that minted it, tampering with any +signed byte is detected, expiry is enforced against the injected clock (capped by the +module TTL ceiling), oversized envelopes are rejected rather than truncated, and no +error value, model repr, or raised exception ever contains the inner access token. +""" + +import base64 +import hashlib +import hmac +import json +from datetime import datetime, timedelta, timezone + +import jwt +import pytest +from cryptography.hazmat.primitives.asymmetric import rsa +from pydantic import SecretStr, ValidationError + +from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( + ENVELOPE_ISSUER, + ENVELOPE_PREFIX, + MAX_ENVELOPE_BYTES, + MAX_ENVELOPE_TTL_SECONDS, + BadSignature, + DecryptFailed, + EnvelopeIdentity, + EnvelopeKeys, + EnvelopeTooLarge, + Expired, + MalformedPayload, + NotAnEnvelope, + OpenedEnvelope, + SealedEnvelope, + UpstreamTokenGrant, + is_envelope, + mint_envelope, + open_envelope, +) +from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value, encrypt_value + +_NOW = datetime(2026, 7, 9, 12, 0, 0, tzinfo=timezone.utc) +_SIGNING_KEY = "unit-test-signing-key-0123456789abcdef0123456789abcdef" +_ENCRYPTION_KEY = "unit-test-encryption-key-fedcba9876543210fedcba9876543210" +_OTHER_SIGNING_KEY = "other-signing-key-aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" +_OTHER_ENCRYPTION_KEY = "other-encryption-key-bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb" +_KEYS = EnvelopeKeys(signing_key=SecretStr(_SIGNING_KEY), encryption_key=SecretStr(_ENCRYPTION_KEY)) +_WRONG_SIGNING = EnvelopeKeys(signing_key=SecretStr(_OTHER_SIGNING_KEY), encryption_key=SecretStr(_ENCRYPTION_KEY)) +_WRONG_ENCRYPTION = EnvelopeKeys(signing_key=SecretStr(_SIGNING_KEY), encryption_key=SecretStr(_OTHER_ENCRYPTION_KEY)) +_ACCESS_TOKEN = "upstream-access-token-do-not-leak-8f14e45fceea" +_REFRESH_TOKEN = "upstream-refresh-token-do-not-leak-1d0aa4b7" +_IDENTITY = EnvelopeIdentity(user_id="user-123", server_id="srv-456") + + +def _full_grant() -> UpstreamTokenGrant: + return UpstreamTokenGrant( + access_token=SecretStr(_ACCESS_TOKEN), + token_type="Bearer", + refresh_token=SecretStr(_REFRESH_TOKEN), + scope="read:tools write:tools", + expires_in=600, + ) + + +def _minimal_grant() -> UpstreamTokenGrant: + return UpstreamTokenGrant(access_token=SecretStr(_ACCESS_TOKEN), token_type="Bearer") + + +def _sealed_token(grant: UpstreamTokenGrant, keys: EnvelopeKeys = _KEYS) -> str: + sealed = mint_envelope(_IDENTITY, grant, keys, _NOW) + assert isinstance(sealed, SealedEnvelope) + return sealed.token.get_secret_value() + + +def _unverified_claims(sealed_token: str) -> dict[str, object]: + return jwt.decode(sealed_token.removeprefix(ENVELOPE_PREFIX), options={"verify_signature": False}) + + +def _forge(claims: dict[str, object], signing_key: str = _SIGNING_KEY) -> str: + return ENVELOPE_PREFIX + jwt.encode(claims, signing_key, algorithm="HS256") + + +def _b64url(raw: bytes) -> str: + return base64.urlsafe_b64encode(raw).rstrip(b"=").decode("ascii") + + +def _hand_crafted_hs256(payload: dict[str, object], signing_key: str = _SIGNING_KEY) -> str: + """Assemble an HS256 envelope from raw bytes, bypassing PyJWT's encode-side claim + guards (it refuses to build a token with a non-string ``iss``). This is the real + attacker path: a client crafts the compact JWT directly, so any registered claim can + carry a hostile JSON type.""" + header = _b64url(json.dumps({"alg": "HS256", "typ": "JWT"}).encode("utf-8")) + body = _b64url(json.dumps(payload).encode("utf-8")) + signing_input = f"{header}.{body}".encode("ascii") + signature = _b64url(hmac.new(signing_key.encode("utf-8"), signing_input, hashlib.sha256).digest()) + return ENVELOPE_PREFIX + f"{header}.{body}.{signature}" + + +def _tampered(sealed_token: str, segment: int, index: int) -> str: + parts = sealed_token.removeprefix(ENVELOPE_PREFIX).split(".") + original = parts[segment][index] + replacement = "A" if original in "QRST" else "Q" + mutated = parts[segment][:index] + replacement + parts[segment][index + 1 :] + rebuilt = ".".join(parts[:segment] + [mutated] + parts[segment + 1 :]) + return ENVELOPE_PREFIX + rebuilt + + +def test_round_trip_recovers_identity_and_grant_exactly(): + grant = _full_grant() + token = _sealed_token(grant) + assert is_envelope(token) + opened = open_envelope(token, _KEYS, _NOW) + assert isinstance(opened, OpenedEnvelope) + assert opened.identity == _IDENTITY + assert opened.grant == grant + assert opened.grant.access_token.get_secret_value() == _ACCESS_TOKEN + assert opened.grant.refresh_token is not None + assert opened.grant.refresh_token.get_secret_value() == _REFRESH_TOKEN + + +def test_minimal_grant_round_trips_without_none_leakage_into_claims(): + token = _sealed_token(_minimal_grant()) + claims = _unverified_claims(token) + blob = claims["grant"] + assert isinstance(blob, str) + plaintext = decrypt_value(value=base64.urlsafe_b64decode(blob), signing_key=_ENCRYPTION_KEY) + assert set(json.loads(plaintext)) == {"access_token", "token_type"} + opened = open_envelope(token, _KEYS, _NOW) + assert isinstance(opened, OpenedEnvelope) + assert opened.grant.refresh_token is None + assert opened.grant.scope is None + assert opened.grant.expires_in is None + + +def test_claim_layout_and_no_plaintext_token_in_envelope(): + token = _sealed_token(_full_grant()) + claims = _unverified_claims(token) + assert set(claims) == {"iss", "iat", "exp", "user_id", "server_id", "grant"} + assert claims["iss"] == ENVELOPE_ISSUER + assert claims["iat"] == int(_NOW.timestamp()) + assert claims["exp"] == int(_NOW.timestamp()) + 600 + assert claims["user_id"] == "user-123" + assert claims["server_id"] == "srv-456" + assert _ACCESS_TOKEN not in token + assert _ACCESS_TOKEN not in json.dumps(claims) + assert _REFRESH_TOKEN not in json.dumps(claims) + + +@pytest.mark.parametrize( + "expires_in, expected_ttl", + [ + (600, 600), + (MAX_ENVELOPE_TTL_SECONDS + 82800, MAX_ENVELOPE_TTL_SECONDS), + (None, MAX_ENVELOPE_TTL_SECONDS), + ], +) +def test_exp_is_min_of_upstream_expires_in_and_cap(expires_in, expected_ttl): + grant = UpstreamTokenGrant(access_token=SecretStr(_ACCESS_TOKEN), token_type="Bearer", expires_in=expires_in) + sealed = mint_envelope(_IDENTITY, grant, _KEYS, _NOW) + assert isinstance(sealed, SealedEnvelope) + assert sealed.expires_at == _NOW + timedelta(seconds=expected_ttl) + + +def test_expiry_honored_against_injected_clock(): + token = _sealed_token(_full_grant()) + assert isinstance(open_envelope(token, _KEYS, _NOW + timedelta(seconds=599)), OpenedEnvelope) + assert isinstance(open_envelope(token, _KEYS, _NOW + timedelta(seconds=600)), Expired) + assert isinstance(open_envelope(token, _KEYS, _NOW + timedelta(seconds=601)), Expired) + + +def test_ttl_cap_enforced_on_open_even_when_upstream_token_lives_longer(): + grant = UpstreamTokenGrant(access_token=SecretStr(_ACCESS_TOKEN), token_type="Bearer", expires_in=86400) + token = _sealed_token(grant) + just_before_cap = _NOW + timedelta(seconds=MAX_ENVELOPE_TTL_SECONDS - 1) + at_cap = _NOW + timedelta(seconds=MAX_ENVELOPE_TTL_SECONDS) + assert isinstance(open_envelope(token, _KEYS, just_before_cap), OpenedEnvelope) + assert isinstance(open_envelope(token, _KEYS, at_cap), Expired) + + +def test_tampering_any_payload_or_signature_byte_is_bad_signature(): + token = _sealed_token(_full_grant()) + parts = token.removeprefix(ENVELOPE_PREFIX).split(".") + for segment in (1, 2): + for index in range(len(parts[segment])): + result = open_envelope(_tampered(token, segment, index), _KEYS, _NOW) + assert isinstance(result, BadSignature), f"segment {segment} index {index}: {result!r}" + + +def test_tampering_header_bytes_never_opens(): + token = _sealed_token(_full_grant()) + parts = token.removeprefix(ENVELOPE_PREFIX).split(".") + for index in range(len(parts[0])): + result = open_envelope(_tampered(token, 0, index), _KEYS, _NOW) + assert isinstance(result, (BadSignature, MalformedPayload)), f"header index {index}: {result!r}" + + +def test_alg_none_is_rejected(): + claims = _unverified_claims(_sealed_token(_full_grant())) + unsigned = ENVELOPE_PREFIX + jwt.encode(claims, None, algorithm="none") + assert isinstance(open_envelope(unsigned, _KEYS, _NOW), MalformedPayload) + + +def test_wrong_signing_key_is_bad_signature(): + token = _sealed_token(_full_grant()) + assert isinstance(open_envelope(token, _WRONG_SIGNING, _NOW), BadSignature) + + +def test_wrong_encryption_key_is_decrypt_failed(): + token = _sealed_token(_full_grant()) + assert isinstance(open_envelope(token, _WRONG_ENCRYPTION, _NOW), DecryptFailed) + + +def test_ciphertext_swapped_from_another_envelope_is_decrypt_failed(): + claims_a = _unverified_claims(_sealed_token(_full_grant(), keys=_KEYS)) + claims_b = _unverified_claims(_sealed_token(_minimal_grant(), keys=_WRONG_ENCRYPTION)) + swapped = _forge({**claims_a, "grant": claims_b["grant"]}) + assert isinstance(open_envelope(swapped, _KEYS, _NOW), DecryptFailed) + + +def test_wrong_issuer_is_malformed_payload(): + claims = _unverified_claims(_sealed_token(_full_grant())) + assert isinstance(open_envelope(_forge({**claims, "iss": "evil-issuer"}), _KEYS, _NOW), MalformedPayload) + + +def test_missing_identity_claim_is_malformed_payload(): + claims = _unverified_claims(_sealed_token(_full_grant())) + forged = _forge({key: value for key, value in claims.items() if key != "user_id"}) + assert isinstance(open_envelope(forged, _KEYS, _NOW), MalformedPayload) + + +@pytest.mark.parametrize("identity_claim", ["user_id", "server_id"]) +def test_signed_empty_identity_claim_is_malformed_payload_not_a_raise(identity_claim): + claims = _unverified_claims(_sealed_token(_full_grant())) + forged = _forge({**claims, identity_claim: ""}) + assert isinstance(open_envelope(forged, _KEYS, _NOW), MalformedPayload) + intact = open_envelope(_forge(claims), _KEYS, _NOW) + assert isinstance(intact, OpenedEnvelope) + assert intact.identity == _IDENTITY + + +def test_lone_surrogate_candidate_is_malformed_payload_not_a_raise(): + surrogate_candidate = ENVELOPE_PREFIX + "\ud800abc.def.ghi" + result = open_envelope(surrogate_candidate, _KEYS, _NOW) + assert isinstance(result, MalformedPayload) + + +@pytest.mark.parametrize( + "override", + [ + {"iat": [1]}, + {"iat": {}}, + {"iat": float("inf")}, + {"nbf": None}, + {"nbf": [1]}, + ], +) +def test_hostile_iat_nbf_types_are_malformed_payload_not_a_raise(override): + claims = _unverified_claims(_sealed_token(_full_grant())) + forged = _forge({**claims, **override}) + result = open_envelope(forged, _KEYS, _NOW) + assert isinstance(result, MalformedPayload) + + +@pytest.mark.parametrize("hostile_iss", [["litellm-mcp-bridge"], 5, {"iss": "x"}]) +def test_non_string_issuer_claim_is_malformed_payload_not_a_raise(hostile_iss): + claims = _unverified_claims(_sealed_token(_full_grant())) + forged = _hand_crafted_hs256({**claims, "iss": hostile_iss}) + result = open_envelope(forged, _KEYS, _NOW) + assert isinstance(result, MalformedPayload) + + +@pytest.mark.parametrize("hostile_exp", ["600", 600.5, [600]]) +def test_non_int_exp_claim_is_malformed_payload_not_a_raise(hostile_exp): + claims = _unverified_claims(_sealed_token(_full_grant())) + forged = _hand_crafted_hs256({**claims, "exp": hostile_exp}) + result = open_envelope(forged, _KEYS, _NOW) + assert isinstance(result, MalformedPayload) + + +def test_unexpected_extra_claim_is_malformed_payload(): + claims = _unverified_claims(_sealed_token(_full_grant())) + forged = _forge({**claims, "role": "admin"}) + assert isinstance(open_envelope(forged, _KEYS, _NOW), MalformedPayload) + + +def test_future_iat_opens_against_injected_now_not_wall_clock(): + future = _NOW + timedelta(seconds=100_000) + sealed = mint_envelope(_IDENTITY, _full_grant(), _KEYS, future) + assert isinstance(sealed, SealedEnvelope) + opened = open_envelope(sealed.token.get_secret_value(), _KEYS, future) + assert isinstance(opened, OpenedEnvelope) + assert opened.identity == _IDENTITY + assert opened.grant == _full_grant() + + +def test_rs256_signed_token_is_rejected_against_the_hs256_pin(): + claims = _unverified_claims(_sealed_token(_full_grant())) + private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048) + rs256_token = ENVELOPE_PREFIX + jwt.encode(claims, private_key, algorithm="RS256") + result = open_envelope(rs256_token, _KEYS, _NOW) + assert isinstance(result, MalformedPayload) + + +@pytest.mark.parametrize("short_key", ["", "too-short", "x" * 31]) +def test_signing_key_below_hs256_minimum_is_rejected_at_construction(short_key): + with pytest.raises(ValidationError): + EnvelopeKeys(signing_key=SecretStr(short_key), encryption_key=SecretStr(_ENCRYPTION_KEY)) + + +def test_signing_key_at_hs256_minimum_is_accepted(): + keys = EnvelopeKeys(signing_key=SecretStr("y" * 32), encryption_key=SecretStr(_ENCRYPTION_KEY)) + assert keys.signing_key.get_secret_value() == "y" * 32 + + +def test_correctly_signed_garbage_grant_blob_is_decrypt_failed(): + claims = _unverified_claims(_sealed_token(_full_grant())) + forged = _forge({**claims, "grant": "not-a-ciphertext"}) + assert isinstance(open_envelope(forged, _KEYS, _NOW), DecryptFailed) + + +def test_decryptable_blob_that_is_not_a_grant_is_malformed_payload(): + claims = _unverified_claims(_sealed_token(_full_grant())) + wrong_shape = base64.urlsafe_b64encode( + bytes(encrypt_value(value=json.dumps({"nope": 1}), signing_key=_ENCRYPTION_KEY)) + ).decode("ascii") + forged = _forge({**claims, "grant": wrong_shape}) + assert isinstance(open_envelope(forged, _KEYS, _NOW), MalformedPayload) + + +def _mint_with_token_len(n: int) -> SealedEnvelope | EnvelopeTooLarge: + grant = UpstreamTokenGrant(access_token=SecretStr("a" * n), token_type="Bearer") + return mint_envelope(_IDENTITY, grant, _KEYS, _NOW) + + +def _largest_token_len_that_mints(lo: int, hi: int) -> int: + if hi - lo <= 1: + return lo + mid = (lo + hi) // 2 + if isinstance(_mint_with_token_len(mid), SealedEnvelope): + return _largest_token_len_that_mints(mid, hi) + return _largest_token_len_that_mints(lo, mid) + + +def test_oversized_grant_is_a_typed_mint_error_never_truncated(): + result = _mint_with_token_len(30000) + assert isinstance(result, EnvelopeTooLarge) + assert result.tag == "envelope_too_large" + assert result.size_bytes > MAX_ENVELOPE_BYTES + assert result.max_bytes == MAX_ENVELOPE_BYTES + + +def test_size_cap_boundary_just_under_succeeds_and_just_over_fails(): + assert isinstance(_mint_with_token_len(1), SealedEnvelope) + assert isinstance(_mint_with_token_len(30000), EnvelopeTooLarge) + largest = _largest_token_len_that_mints(1, 30000) + assert largest > 6000 + sealed = _mint_with_token_len(largest) + assert isinstance(sealed, SealedEnvelope) + assert len(sealed.token.get_secret_value().encode("utf-8")) <= MAX_ENVELOPE_BYTES + overflowing = _mint_with_token_len(largest + 1) + assert isinstance(overflowing, EnvelopeTooLarge) + assert overflowing.size_bytes > MAX_ENVELOPE_BYTES + opened = open_envelope(sealed.token.get_secret_value(), _KEYS, _NOW) + assert isinstance(opened, OpenedEnvelope) + + +def test_open_size_guard_measures_bytes_not_characters(): + """The open-side size guard must reject on UTF-8 byte length, matching mint's cap, so a + hostile multi-byte candidate whose character count is under the cap but whose byte count is + over it is rejected up front rather than reaching the expensive HMAC/decrypt path. Patching + _decode_claims to fail loudly proves the guard short-circuits before decode.""" + from unittest.mock import patch + + from litellm.proxy._experimental.mcp_server.outbound_credentials import envelope + + multibyte_body = "é" * 7000 # 7000 chars, 14000 UTF-8 bytes + candidate = ENVELOPE_PREFIX + multibyte_body + assert len(candidate) <= MAX_ENVELOPE_BYTES + assert len(candidate.encode("utf-8")) > MAX_ENVELOPE_BYTES + + with patch.object(envelope, "_decode_claims", side_effect=AssertionError("decode reached")) as decode: + result = open_envelope(candidate, _KEYS, _NOW) + + assert isinstance(result, MalformedPayload) + decode.assert_not_called() + + +def test_open_size_guard_rejects_oversize_character_count_before_decode(): + """A candidate whose character count already exceeds the cap is rejected up front, before the + decode path, so an arbitrarily long hostile string is not run through HMAC/decrypt. The cheap + character precheck makes this O(1) since UTF-8 byte length is never below character length.""" + from unittest.mock import patch + + from litellm.proxy._experimental.mcp_server.outbound_credentials import envelope + + candidate = ENVELOPE_PREFIX + ("a" * (MAX_ENVELOPE_BYTES + 1)) + assert len(candidate) > MAX_ENVELOPE_BYTES + + with patch.object(envelope, "_decode_claims", side_effect=AssertionError("decode reached")) as decode: + result = open_envelope(candidate, _KEYS, _NOW) + + assert isinstance(result, MalformedPayload) + decode.assert_not_called() + + +def test_is_envelope_detects_only_prefixed_values(): + assert is_envelope(_sealed_token(_full_grant())) + raw_jwt = jwt.encode({"sub": "user-123"}, _SIGNING_KEY, algorithm="HS256") + assert not is_envelope(raw_jwt) + assert not is_envelope("some-random-opaque-token") + assert not is_envelope("") + + +def test_open_on_non_envelope_input_is_not_an_envelope(): + raw_jwt = jwt.encode({"sub": "user-123"}, _SIGNING_KEY, algorithm="HS256") + assert isinstance(open_envelope(raw_jwt, _KEYS, _NOW), NotAnEnvelope) + assert isinstance(open_envelope("", _KEYS, _NOW), NotAnEnvelope) + assert isinstance(open_envelope(_ACCESS_TOKEN, _KEYS, _NOW), NotAnEnvelope) + + +def test_open_on_prefixed_garbage_is_malformed_payload(): + assert isinstance(open_envelope(ENVELOPE_PREFIX + "garbage", _KEYS, _NOW), MalformedPayload) + assert isinstance(open_envelope(ENVELOPE_PREFIX + _ACCESS_TOKEN, _KEYS, _NOW), MalformedPayload) + + +def test_no_result_value_ever_reveals_the_access_token(): + grant = _full_grant() + sealed = mint_envelope(_IDENTITY, grant, _KEYS, _NOW) + assert isinstance(sealed, SealedEnvelope) + token = sealed.token.get_secret_value() + oversized_grant = UpstreamTokenGrant(access_token=SecretStr(_ACCESS_TOKEN + "x" * 30000), token_type="Bearer") + values = ( + sealed, + open_envelope(token, _KEYS, _NOW), + mint_envelope(_IDENTITY, oversized_grant, _KEYS, _NOW), + open_envelope(_ACCESS_TOKEN, _KEYS, _NOW), + open_envelope(ENVELOPE_PREFIX + _ACCESS_TOKEN, _KEYS, _NOW), + open_envelope(token, _WRONG_SIGNING, _NOW), + open_envelope(token, _WRONG_ENCRYPTION, _NOW), + open_envelope(token, _KEYS, _NOW + timedelta(seconds=601)), + grant, + ) + for value in values: + assert _ACCESS_TOKEN not in repr(value) + assert _ACCESS_TOKEN not in str(value) + assert _REFRESH_TOKEN not in repr(value) + assert _REFRESH_TOKEN not in str(value) + + +def test_non_positive_expires_in_is_rejected_at_construction_without_leaking(): + for bad_expires_in in (0, -5): + with pytest.raises(ValidationError) as excinfo: + UpstreamTokenGrant( + access_token=SecretStr(_ACCESS_TOKEN), + token_type="Bearer", + expires_in=bad_expires_in, + ) + assert _ACCESS_TOKEN not in str(excinfo.value) + assert _ACCESS_TOKEN not in repr(excinfo.value) + + +def test_empty_identity_and_key_fields_are_rejected_at_construction(): + with pytest.raises(ValidationError): + EnvelopeIdentity(user_id="", server_id="srv-456") + with pytest.raises(ValidationError): + EnvelopeIdentity(user_id="user-123", server_id="") + with pytest.raises(ValidationError): + EnvelopeKeys(signing_key=SecretStr(""), encryption_key=SecretStr(_ENCRYPTION_KEY)) + with pytest.raises(ValidationError): + EnvelopeKeys(signing_key=SecretStr(_SIGNING_KEY), encryption_key=SecretStr("")) + with pytest.raises(ValidationError): + UpstreamTokenGrant(access_token=SecretStr(""), token_type="Bearer") + + +def test_public_models_are_frozen(): + sealed = mint_envelope(_IDENTITY, _full_grant(), _KEYS, _NOW) + assert isinstance(sealed, SealedEnvelope) + opened = open_envelope(sealed.token.get_secret_value(), _KEYS, _NOW) + assert isinstance(opened, OpenedEnvelope) + with pytest.raises(ValidationError): + sealed.token = SecretStr("overwritten") + with pytest.raises(ValidationError): + opened.grant = _minimal_grant() + with pytest.raises(ValidationError): + _IDENTITY.user_id = "someone-else"