mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
Merge pull request #32748 from BerriAI/litellm_lit4338_envelope_module
feat(mcp): add sealed envelope module for dcr_bridge client-held credentials
This commit is contained in:
tin-berri
GitHub
commit
ab53d06a2e
2 changed files with 845 additions and 0 deletions
|
|
@ -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()
|
||||
|
|
@ -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"
|
||||
Loading…
Add table
Reference in a new issue