mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-07 02:59:05 +00:00
feat(mcp): add sealed envelope module for dcr_bridge client-held credentials
Pure, unwired module: mints and opens the single client-held bearer that carries both a litellm identity and the encrypted upstream OAuth grant with zero server-side storage. HS256 JWT signing (same approach as the BYOK session bearer) plus the existing encrypt_value/decrypt_value symmetric helpers, with all key material and the clock injected as parameters. Opening returns typed frozen error values (not_an_envelope, bad_signature, expired, malformed_payload, decrypt_failed); minting rejects envelopes over MAX_ENVELOPE_BYTES with a typed error instead of truncating. Error values and reprs never carry token material.
This commit is contained in:
Tin Chi Lo
parent
34602ff627
commit
65c80919aa
2 changed files with 649 additions and 0 deletions
|
|
@ -0,0 +1,304 @@
|
|||
"""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. Raising is reserved for programmer errors, which the pydantic
|
||||
models reject at construction (e.g. a non-positive ``expires_in``).
|
||||
"""
|
||||
|
||||
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."""
|
||||
|
||||
model_config = ConfigDict(frozen=True)
|
||||
signing_key: SecretStr = Field(min_length=1)
|
||||
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):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
iss: str
|
||||
iat: int
|
||||
exp: int
|
||||
user_id: str
|
||||
server_id: str
|
||||
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.
|
||||
"""
|
||||
if not is_envelope(candidate):
|
||||
return NotAnEnvelope()
|
||||
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:
|
||||
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"]},
|
||||
)
|
||||
except jwt.InvalidSignatureError:
|
||||
return BadSignature()
|
||||
except jwt.InvalidTokenError:
|
||||
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,345 @@
|
|||
"""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 json
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
||||
import jwt
|
||||
import pytest
|
||||
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 _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)
|
||||
|
||||
|
||||
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_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