diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/envelope.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/envelope.py index 666ef87e6b2..446b36a9f84 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/envelope.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/envelope.py @@ -21,8 +21,15 @@ 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. Raising is reserved for programmer errors, which the pydantic -models reject at construction (e.g. a non-positive ``expires_in``). +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 @@ -85,10 +92,15 @@ class UpstreamTokenGrant(BaseModel): class EnvelopeKeys(BaseModel): - """Injected key material: the HS256 signing key and the symmetric encryption key.""" + """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=1) + signing_key: SecretStr = Field(min_length=32) encryption_key: SecretStr = Field(min_length=1) @@ -160,12 +172,22 @@ EnvelopeOpenError: TypeAlias = NotAnEnvelope | BadSignature | Expired | Malforme class _EnvelopeClaims(BaseModel): - """Decoded-claims boundary. ``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``.""" + """Decoded-claims boundary that pins the exact shape :func:`mint_envelope` emits. - model_config = ConfigDict(frozen=True) + ``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 @@ -228,10 +250,15 @@ def open_envelope( """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. + 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() + if len(candidate) > MAX_ENVELOPE_BYTES: + return MalformedPayload() claims = _decode_claims(candidate.removeprefix(ENVELOPE_PREFIX), keys.signing_key) if not isinstance(claims, _EnvelopeClaims): return claims @@ -267,17 +294,34 @@ 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, "require": ["iss", "iat", "exp"]}, + options={ + "verify_exp": False, + "verify_iat": False, + "verify_nbf": False, + "require": ["iss", "iat", "exp"], + }, ) except jwt.InvalidSignatureError: return BadSignature() - except jwt.InvalidTokenError: + except (jwt.InvalidTokenError, ValueError, TypeError): return MalformedPayload() try: return _EnvelopeClaims.model_validate(payload) 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 index b9d91c011c6..99f1bd6a178 100644 --- 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 @@ -9,11 +9,14 @@ error value, model repr, or raised exception ever contains the inner access toke """ 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 ( @@ -79,6 +82,22 @@ 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] @@ -221,6 +240,80 @@ def test_signed_empty_identity_claim_is_malformed_payload_not_a_raise(identity_c 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"})