mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
feat(mcp): identity-only session tokens for the gateway DCR front door
This commit is contained in:
Tin Chi Lo
parent
90195afa06
commit
5b2877f742
4 changed files with 887 additions and 0 deletions
|
|
@ -0,0 +1,190 @@
|
|||
"""Producer and consumer helpers for the gateway-level DCR session token.
|
||||
|
||||
The aggregate ``/mcp`` front door (``mcp_gateway_dcr``) issues the identity-only session
|
||||
tokens defined in :mod:`.session_token`. The gateway token endpoint mints them (producer)
|
||||
after SSO sign-in, and at the MCP admission edge the gateway derives the session signing
|
||||
key from the proxy ``master_key``, opens the bearer, and admits the request under the
|
||||
recovered litellm user (consumer), reloading the live user record and policy before
|
||||
anything runs. This module is the pure surface for both sides; the token-endpoint and
|
||||
admission wiring live in their respective call sites.
|
||||
|
||||
The signing key is derived with the same memory-hard scrypt construction as
|
||||
:func:`~.bridge_credentials.envelope_keys_from_master_key` but under a distinct domain
|
||||
label, so session tokens and bridge envelopes never share key material: a token of one
|
||||
family is unverifiable in the other by key separation, on top of the distinct issuers,
|
||||
prefixes, and claim shapes.
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
from datetime import datetime
|
||||
from functools import lru_cache
|
||||
from typing import Literal, TypeAlias
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, SecretStr
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import (
|
||||
OpenedSessionToken,
|
||||
SessionExpired,
|
||||
SessionKeys,
|
||||
SessionPrincipal,
|
||||
is_session_refresh_token,
|
||||
is_session_token,
|
||||
open_session_refresh_token,
|
||||
open_session_token,
|
||||
)
|
||||
|
||||
_SESSION_SIGNING_KEY_DOMAIN = b"litellm-mcp-gateway:session-signing:"
|
||||
|
||||
# scrypt work factors (RFC 7914), identical to the envelope KDF: memory-hard so a captured
|
||||
# session token is not a cheap offline oracle for the master key.
|
||||
_SCRYPT_N = 2**15
|
||||
_SCRYPT_R = 8
|
||||
_SCRYPT_P = 1
|
||||
_SCRYPT_MAXMEM = 128 * _SCRYPT_N * _SCRYPT_R * _SCRYPT_P * 2
|
||||
_DERIVED_KEY_BYTES = 32
|
||||
|
||||
|
||||
@lru_cache(maxsize=8)
|
||||
def session_keys_from_master_key(master_key: str) -> SessionKeys:
|
||||
"""Derive the session signing key from the proxy master key.
|
||||
|
||||
A memory-hard scrypt KDF (RFC 7914) over a session-specific domain-label salt yields a
|
||||
256-bit subkey from the one secret, so the producer (mint) and consumer (open) agree on
|
||||
the key without persisting any. The domain label differs from both envelope labels in
|
||||
:mod:`.bridge_credentials`, so compromise or misuse of one token family never crosses
|
||||
into the other. The result is cached (the master key is fixed for a process); rotating
|
||||
``master_key`` invalidates every outstanding session, which is the intended behavior
|
||||
for a signing-key change.
|
||||
"""
|
||||
signing = hashlib.scrypt(
|
||||
master_key.encode(),
|
||||
salt=_SESSION_SIGNING_KEY_DOMAIN,
|
||||
n=_SCRYPT_N,
|
||||
r=_SCRYPT_R,
|
||||
p=_SCRYPT_P,
|
||||
maxmem=_SCRYPT_MAXMEM,
|
||||
dklen=_DERIVED_KEY_BYTES,
|
||||
).hex()
|
||||
return SessionKeys(signing_key=SecretStr(signing))
|
||||
|
||||
|
||||
class NotSessionBearer(BaseModel):
|
||||
"""The bearer is not session-shaped; admission continues on its normal path."""
|
||||
|
||||
model_config = ConfigDict(frozen=True)
|
||||
tag: Literal["not_session_bearer"] = "not_session_bearer"
|
||||
|
||||
|
||||
class SessionBearerAdmitted(BaseModel):
|
||||
"""A valid session access token: the principal to admit under after a live reload."""
|
||||
|
||||
model_config = ConfigDict(frozen=True)
|
||||
tag: Literal["admitted"] = "admitted"
|
||||
principal: SessionPrincipal
|
||||
|
||||
|
||||
class SessionBearerInvalid(BaseModel):
|
||||
"""The bearer is session-shaped but must not admit (expired, tampered, wrong key, or a
|
||||
refresh token presented at the tool-call edge); admission fails closed with the
|
||||
``invalid_token`` challenge rather than falling through to another arm. ``expired``
|
||||
distinguishes a routine expiry (debug-log worthy) from a tampered or foreign token."""
|
||||
|
||||
model_config = ConfigDict(frozen=True)
|
||||
tag: Literal["invalid"] = "invalid"
|
||||
expired: bool = False
|
||||
|
||||
|
||||
SessionBearerResult: TypeAlias = NotSessionBearer | SessionBearerAdmitted | SessionBearerInvalid
|
||||
|
||||
|
||||
def _strip_bearer(value: str) -> str:
|
||||
parts = value.split(None, 1)
|
||||
if len(parts) == 2 and parts[0].lower() == "bearer":
|
||||
return parts[1]
|
||||
return value
|
||||
|
||||
|
||||
def is_session_bearer_shaped(authorization_value: str) -> bool:
|
||||
"""Cheap, keyless test that an ``Authorization`` value carries a session token of either
|
||||
kind (optional ``Bearer`` scheme stripped). The admission edge engages the session arm
|
||||
for an access token (to admit) and for a refresh token (to reject it explicitly, since
|
||||
a refresh credential is never usable at the tool-call edge); anything else falls
|
||||
through to normal admission."""
|
||||
candidate = _strip_bearer(authorization_value)
|
||||
return is_session_token(candidate) or is_session_refresh_token(candidate)
|
||||
|
||||
|
||||
def resolve_session_bearer(
|
||||
authorization_value: str,
|
||||
keys: SessionKeys,
|
||||
now: datetime,
|
||||
) -> SessionBearerResult:
|
||||
"""Classify an ``Authorization`` value presented at the aggregate MCP edge.
|
||||
|
||||
Strips an optional ``Bearer`` scheme, then returns ``NotSessionBearer`` for a
|
||||
non-session bearer (normal admission continues), ``SessionBearerAdmitted`` with the
|
||||
recovered principal for a valid access token, and ``SessionBearerInvalid`` for a
|
||||
session-shaped bearer that must not admit. Never raises: total over hostile input via
|
||||
:func:`~.session_token.open_session_token`.
|
||||
|
||||
A refresh token is ``SessionBearerInvalid`` here: it is a valid gateway credential but
|
||||
only ever presented back to the token endpoint, so admission must fail it closed rather
|
||||
than let it fall through to another arm.
|
||||
"""
|
||||
candidate = _strip_bearer(authorization_value)
|
||||
if is_session_refresh_token(candidate):
|
||||
return SessionBearerInvalid()
|
||||
if not is_session_token(candidate):
|
||||
return NotSessionBearer()
|
||||
opened = open_session_token(candidate, keys, now)
|
||||
if isinstance(opened, OpenedSessionToken):
|
||||
return SessionBearerAdmitted(principal=opened.principal)
|
||||
return SessionBearerInvalid(expired=isinstance(opened, SessionExpired))
|
||||
|
||||
|
||||
class SessionRefreshOpened(BaseModel):
|
||||
"""A valid session refresh token presented to the token endpoint: the principal to
|
||||
re-validate and renew under."""
|
||||
|
||||
model_config = ConfigDict(frozen=True)
|
||||
tag: Literal["opened"] = "opened"
|
||||
principal: SessionPrincipal
|
||||
|
||||
|
||||
class SessionRefreshInvalid(BaseModel):
|
||||
"""The presented refresh grant is not a valid session refresh token for this client
|
||||
(not refresh-shaped, will not open, or bound to a different ``client_id``); the token
|
||||
endpoint fails the refresh closed."""
|
||||
|
||||
model_config = ConfigDict(frozen=True)
|
||||
tag: Literal["invalid"] = "invalid"
|
||||
|
||||
|
||||
SessionRefreshResult: TypeAlias = SessionRefreshOpened | SessionRefreshInvalid
|
||||
|
||||
|
||||
def open_session_refresh_bearer(
|
||||
refresh_value: str,
|
||||
keys: SessionKeys,
|
||||
now: datetime,
|
||||
expected_client_id: str,
|
||||
) -> SessionRefreshResult:
|
||||
"""Open a session refresh token presented on a ``refresh_token`` grant.
|
||||
|
||||
The token-endpoint mirror of :func:`resolve_session_bearer`: strips an optional
|
||||
``Bearer`` scheme, then returns ``SessionRefreshOpened`` with the recovered principal,
|
||||
or ``SessionRefreshInvalid`` for anything that is not a valid session refresh token
|
||||
issued to ``expected_client_id``. Never raises. The client binding (RFC 6749 section 6)
|
||||
stops a refresh token stolen from one DCR client from being renewed through another;
|
||||
``client_id`` is not a secret (the caller presents it), so a plain equality check is
|
||||
sufficient and, unlike ``hmac.compare_digest`` on ``str``, does not raise on non-ASCII.
|
||||
"""
|
||||
candidate = _strip_bearer(refresh_value)
|
||||
if not is_session_refresh_token(candidate):
|
||||
return SessionRefreshInvalid()
|
||||
opened = open_session_refresh_token(candidate, keys, now)
|
||||
if not isinstance(opened, OpenedSessionToken):
|
||||
return SessionRefreshInvalid()
|
||||
if opened.principal.client_id != expected_client_id:
|
||||
return SessionRefreshInvalid()
|
||||
return SessionRefreshOpened(principal=opened.principal)
|
||||
|
|
@ -0,0 +1,356 @@
|
|||
"""Identity-only session tokens for the gateway-level (aggregate ``/mcp``) DCR front door.
|
||||
|
||||
A DCR client that signs in through LiteLLM SSO holds ONE bearer that carries ONLY a
|
||||
litellm identity; unlike the :mod:`.envelope` bridge bearer it seals no upstream
|
||||
credential, because the custody model vaults every upstream token server-side in
|
||||
``LiteLLM_MCPUserCredentials`` and egress resolves them by user at call time. The token
|
||||
is therefore a stable REFERENCE, not an authorization: admission reloads the live user
|
||||
record and policy on every request, so deactivating the user (or their team) kills
|
||||
outstanding sessions immediately without a revocation store.
|
||||
|
||||
Wire shape: ``llm_session_`` (access) / ``llm_srefresh_`` (refresh) + an HS256 JWT,
|
||||
the same signing approach as :mod:`.envelope`. Claims are ``iss``/``iat``/``exp``
|
||||
plus ``kind``, ``user_id``, and ``client_id``; ``client_id`` binds the refresh token
|
||||
to the DCR client it was issued to (RFC 6749 section 6) and is carried on the access
|
||||
token for parity and audit. There is no encrypted payload: nothing in a session token
|
||||
is secret beyond the signature, and reprs never print the signed value because minted
|
||||
tokens are ``SecretStr``.
|
||||
|
||||
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.
|
||||
Failures are values: :func:`open_session_token` and :func:`open_session_refresh_token`
|
||||
are total over hostile, attacker-controlled input and return a
|
||||
``SessionTokenOpenError`` variant rather than raising. PyJWT's ``iat``/``nbf``/``exp``
|
||||
validators are disabled for the same reasons documented in :mod:`.envelope` (they
|
||||
raise on hostile claim types and compare against the wall clock instead of the
|
||||
injected ``now``); the strict pydantic claims model is the sole, total type gate.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime, timedelta
|
||||
from typing import Literal, TypeAlias
|
||||
|
||||
import jwt
|
||||
from pydantic import BaseModel, ConfigDict, Field, SecretStr, ValidationError
|
||||
|
||||
SESSION_TOKEN_PREFIX = "llm_session_"
|
||||
"""Marker prefix on every serialized session ACCESS token so the admission edge can cheaply
|
||||
tell a gateway session from a litellm key, JWT, or bridge envelope before doing any
|
||||
cryptography. Distinct from the ``llm_env_``/``llm_refresh_`` envelope prefixes."""
|
||||
|
||||
SESSION_REFRESH_PREFIX = "llm_srefresh_"
|
||||
"""Marker prefix on every serialized session REFRESH token. A distinct prefix keeps the two
|
||||
credentials routable without crypto and, together with the signed ``kind`` claim, stops one
|
||||
from being presented where the other is expected: the refresh token is only ever presented
|
||||
back to the token endpoint, never at the MCP edge."""
|
||||
|
||||
SESSION_ISSUER = "litellm-mcp-gateway"
|
||||
"""``iss`` claim stamped into every session token and required back on open. Distinct from
|
||||
the envelope issuer so a token of one family can never validate in the other even under a
|
||||
hypothetical shared signing key."""
|
||||
|
||||
SESSION_TTL_SECONDS = 3600
|
||||
"""Session ACCESS token lifetime (1h), matching the access-envelope and BYOK session bearer
|
||||
windows: a client-held credential never outlives a bounded window, and each refresh
|
||||
re-validates the live user before re-minting."""
|
||||
|
||||
SESSION_REFRESH_TTL_SECONDS = 1209600
|
||||
"""Session REFRESH token lifetime (14 days), matching the refresh-envelope bound. Each
|
||||
renewal re-validates the sealed user against the live record (deactivation gates it) and
|
||||
rotates the refresh token, so the practical bound is idle time, not a fixed session."""
|
||||
|
||||
MAX_SESSION_TOKEN_BYTES = 4096
|
||||
"""Size cap on the serialized token (prefix + JWT, in bytes) and on any candidate accepted
|
||||
by the openers. Session claims are small; the only variable-length field is ``client_id``
|
||||
(a sealed DCR client record), and 4096 leaves ample headroom under common 8-16KB header
|
||||
limits while bounding hostile input before JWT parsing."""
|
||||
|
||||
_SESSION_JWT_ALGORITHM = "HS256"
|
||||
|
||||
SessionTokenKind = Literal["session", "session_refresh"]
|
||||
"""Which credential a session token is. Stamped into the signed claims and required to match
|
||||
on open, so a signature-valid token of one kind cannot be replayed as the other even if its
|
||||
wire prefix is swapped (the prefix is not part of the signed payload; this claim is)."""
|
||||
|
||||
|
||||
class SessionPrincipal(BaseModel):
|
||||
"""The litellm user a session token identifies and the DCR client it was issued to.
|
||||
|
||||
``user_id`` is the SSO-established litellm user subject, never a credential: admission
|
||||
reloads the live user record by it, so current role, team, and revocation state are
|
||||
enforced at use time rather than frozen at mint time. ``client_id`` is the (stateless,
|
||||
gateway-sealed) DCR client identifier the token was issued to; the token endpoint
|
||||
requires it to match on the refresh grant.
|
||||
"""
|
||||
|
||||
model_config = ConfigDict(frozen=True)
|
||||
user_id: str = Field(min_length=1)
|
||||
client_id: str = Field(min_length=1)
|
||||
|
||||
|
||||
class SessionKeys(BaseModel):
|
||||
"""Injected key material: the HS256 signing 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)
|
||||
|
||||
|
||||
class MintedSessionToken(BaseModel):
|
||||
"""A minted session token: the client-held bearer value and when it expires."""
|
||||
|
||||
model_config = ConfigDict(frozen=True)
|
||||
token: SecretStr
|
||||
expires_at: datetime
|
||||
|
||||
|
||||
class OpenedSessionToken(BaseModel):
|
||||
"""A validated session token of either kind: the principal it was minted for."""
|
||||
|
||||
model_config = ConfigDict(frozen=True)
|
||||
principal: SessionPrincipal
|
||||
|
||||
|
||||
class SessionTokenTooLarge(BaseModel):
|
||||
"""The serialized token exceeded ``MAX_SESSION_TOKEN_BYTES``; carries sizes only. Only
|
||||
reachable through an oversized ``client_id``, which registration should have bounded."""
|
||||
|
||||
model_config = ConfigDict(frozen=True)
|
||||
tag: Literal["session_token_too_large"] = "session_token_too_large"
|
||||
size_bytes: int
|
||||
max_bytes: int
|
||||
|
||||
|
||||
SessionTokenMintError: TypeAlias = SessionTokenTooLarge
|
||||
|
||||
|
||||
class NotASessionToken(BaseModel):
|
||||
"""The candidate does not carry the expected session prefix."""
|
||||
|
||||
model_config = ConfigDict(frozen=True)
|
||||
tag: Literal["not_a_session_token"] = "not_a_session_token"
|
||||
|
||||
|
||||
class SessionBadSignature(BaseModel):
|
||||
"""The JWT signature does not verify under the provided signing key."""
|
||||
|
||||
model_config = ConfigDict(frozen=True)
|
||||
tag: Literal["session_bad_signature"] = "session_bad_signature"
|
||||
|
||||
|
||||
class SessionExpired(BaseModel):
|
||||
"""The token's ``exp`` is not in the future relative to the provided ``now``."""
|
||||
|
||||
model_config = ConfigDict(frozen=True)
|
||||
tag: Literal["session_expired"] = "session_expired"
|
||||
|
||||
|
||||
class SessionMalformed(BaseModel):
|
||||
"""The token is not a well-formed session token: undecodable JWT, wrong issuer, wrong
|
||||
``kind``, or missing/mistyped/extra claims."""
|
||||
|
||||
model_config = ConfigDict(frozen=True)
|
||||
tag: Literal["session_malformed"] = "session_malformed"
|
||||
|
||||
|
||||
SessionTokenOpenError: TypeAlias = NotASessionToken | SessionBadSignature | SessionExpired | SessionMalformed
|
||||
|
||||
|
||||
class _SessionClaims(BaseModel):
|
||||
"""Decoded-claims boundary that pins the exact shape the mints emit.
|
||||
|
||||
``user_id``/``client_id`` mirror the ``min_length`` constraints of
|
||||
:class:`SessionPrincipal` so any claim set that validates here also constructs a
|
||||
principal, keeping the openers raise-free: a correctly signed JWT with an empty
|
||||
identity claim fails here and maps to ``SessionMalformed``. ``strict`` rejects coerced
|
||||
types (``exp: "123"``) and ``extra="forbid"`` rejects any claim the gateway never
|
||||
mints; PyJWT's own registered-claim validators are disabled at decode (see module
|
||||
docstring), so this model is the sole, total type gate for every claim.
|
||||
"""
|
||||
|
||||
model_config = ConfigDict(frozen=True, strict=True, extra="forbid")
|
||||
iss: str
|
||||
iat: int
|
||||
exp: int
|
||||
kind: SessionTokenKind
|
||||
user_id: str = Field(min_length=1)
|
||||
client_id: str = Field(min_length=1)
|
||||
|
||||
|
||||
def is_session_token(candidate: str) -> bool:
|
||||
"""Cheap prefix check for a session ACCESS token so the admission edge can route gateway
|
||||
sessions vs keys, JWTs, and envelopes without crypto."""
|
||||
return candidate.startswith(SESSION_TOKEN_PREFIX)
|
||||
|
||||
|
||||
def is_session_refresh_token(candidate: str) -> bool:
|
||||
"""Cheap prefix check for a session REFRESH token so the token endpoint can route a
|
||||
refresh grant without crypto."""
|
||||
return candidate.startswith(SESSION_REFRESH_PREFIX)
|
||||
|
||||
|
||||
def mint_session_token(
|
||||
principal: SessionPrincipal,
|
||||
keys: SessionKeys,
|
||||
now: datetime,
|
||||
) -> MintedSessionToken | SessionTokenMintError:
|
||||
"""Mint the short-lived session ACCESS token for ``principal``.
|
||||
|
||||
``exp`` is ``SESSION_TTL_SECONDS`` from ``now``. Returns ``SessionTokenTooLarge`` when
|
||||
the serialized token exceeds ``MAX_SESSION_TOKEN_BYTES``.
|
||||
"""
|
||||
return _mint(
|
||||
kind="session",
|
||||
prefix=SESSION_TOKEN_PREFIX,
|
||||
principal=principal,
|
||||
expires_at=now + timedelta(seconds=SESSION_TTL_SECONDS),
|
||||
keys=keys,
|
||||
now=now,
|
||||
)
|
||||
|
||||
|
||||
def mint_session_refresh_token(
|
||||
principal: SessionPrincipal,
|
||||
keys: SessionKeys,
|
||||
now: datetime,
|
||||
) -> MintedSessionToken | SessionTokenMintError:
|
||||
"""Mint the long-lived session REFRESH token for ``principal``.
|
||||
|
||||
``exp`` is ``SESSION_REFRESH_TTL_SECONDS`` from ``now``. Minting a distinct
|
||||
``kind="session_refresh"`` claim is what keeps a refresh token from ever opening as an
|
||||
access credential at the MCP edge.
|
||||
"""
|
||||
return _mint(
|
||||
kind="session_refresh",
|
||||
prefix=SESSION_REFRESH_PREFIX,
|
||||
principal=principal,
|
||||
expires_at=now + timedelta(seconds=SESSION_REFRESH_TTL_SECONDS),
|
||||
keys=keys,
|
||||
now=now,
|
||||
)
|
||||
|
||||
|
||||
def open_session_token(
|
||||
candidate: str,
|
||||
keys: SessionKeys,
|
||||
now: datetime,
|
||||
) -> OpenedSessionToken | SessionTokenOpenError:
|
||||
"""Validate a session ACCESS ``candidate`` and recover the principal.
|
||||
|
||||
Never raises for bad input: every invalid, expired, tampered, or wrong-kind candidate
|
||||
maps to a distinct ``SessionTokenOpenError`` variant.
|
||||
"""
|
||||
return _open(candidate, prefix=SESSION_TOKEN_PREFIX, expected_kind="session", keys=keys, now=now)
|
||||
|
||||
|
||||
def open_session_refresh_token(
|
||||
candidate: str,
|
||||
keys: SessionKeys,
|
||||
now: datetime,
|
||||
) -> OpenedSessionToken | SessionTokenOpenError:
|
||||
"""Validate a session REFRESH ``candidate`` and recover the principal.
|
||||
|
||||
Total over hostile input exactly like :func:`open_session_token`. The
|
||||
``kind="session_refresh"`` claim is required, so an access token re-prefixed as a
|
||||
refresh one is rejected as ``SessionMalformed``.
|
||||
"""
|
||||
return _open(candidate, prefix=SESSION_REFRESH_PREFIX, expected_kind="session_refresh", keys=keys, now=now)
|
||||
|
||||
|
||||
def _mint(
|
||||
kind: SessionTokenKind,
|
||||
prefix: str,
|
||||
principal: SessionPrincipal,
|
||||
expires_at: datetime,
|
||||
keys: SessionKeys,
|
||||
now: datetime,
|
||||
) -> MintedSessionToken | SessionTokenTooLarge:
|
||||
"""Sign the claims for either token kind and enforce the size cap. Shared by both mints
|
||||
so the JWT shape, issuer, and size guard cannot drift between access and refresh."""
|
||||
claims = _SessionClaims(
|
||||
iss=SESSION_ISSUER,
|
||||
iat=int(now.timestamp()),
|
||||
exp=int(expires_at.timestamp()),
|
||||
kind=kind,
|
||||
user_id=principal.user_id,
|
||||
client_id=principal.client_id,
|
||||
)
|
||||
token = prefix + jwt.encode(
|
||||
claims.model_dump(), keys.signing_key.get_secret_value(), algorithm=_SESSION_JWT_ALGORITHM
|
||||
)
|
||||
size_bytes = len(token.encode("utf-8"))
|
||||
if size_bytes > MAX_SESSION_TOKEN_BYTES:
|
||||
return SessionTokenTooLarge(size_bytes=size_bytes, max_bytes=MAX_SESSION_TOKEN_BYTES)
|
||||
return MintedSessionToken(token=SecretStr(token), expires_at=expires_at)
|
||||
|
||||
|
||||
def _open(
|
||||
candidate: str,
|
||||
prefix: str,
|
||||
expected_kind: SessionTokenKind,
|
||||
keys: SessionKeys,
|
||||
now: datetime,
|
||||
) -> OpenedSessionToken | SessionTokenOpenError:
|
||||
"""Prefix-route, size-bound, signature-verify, kind-check, and expiry-check an
|
||||
attacker-controlled candidate, shared by both openers so the security gate is identical
|
||||
for access and refresh. Returns the opened token or a distinct error; never raises."""
|
||||
if not candidate.startswith(prefix):
|
||||
return NotASessionToken()
|
||||
# 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 the cap in characters.
|
||||
if len(candidate) > MAX_SESSION_TOKEN_BYTES:
|
||||
return SessionMalformed()
|
||||
if len(candidate.encode("utf-8", "surrogatepass")) > MAX_SESSION_TOKEN_BYTES:
|
||||
return SessionMalformed()
|
||||
claims = _decode_claims(candidate.removeprefix(prefix), keys.signing_key)
|
||||
if not isinstance(claims, _SessionClaims):
|
||||
return claims
|
||||
if claims.kind != expected_kind:
|
||||
return SessionMalformed()
|
||||
if now.timestamp() >= claims.exp:
|
||||
return SessionExpired()
|
||||
return OpenedSessionToken(principal=SessionPrincipal(user_id=claims.user_id, client_id=claims.client_id))
|
||||
|
||||
|
||||
def _decode_claims(
|
||||
compact: str,
|
||||
signing_key: SecretStr,
|
||||
) -> _SessionClaims | SessionBadSignature | SessionMalformed:
|
||||
"""Verify the HS256 signature and shape of an attacker-controlled compact JWT.
|
||||
|
||||
``compact`` is fully hostile and bounded to ``MAX_SESSION_TOKEN_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, every decode failure is ``SessionMalformed``: a non-UTF-8 candidate surfaces
|
||||
as ``UnicodeEncodeError`` (a ``ValueError``), a non-string registered claim as a
|
||||
``TypeError`` from PyJWT's claim validators, and a wrong issuer or structurally invalid
|
||||
token as an ``InvalidTokenError``. ``_SessionClaims`` is the total type gate.
|
||||
"""
|
||||
try:
|
||||
payload = jwt.decode(
|
||||
compact,
|
||||
signing_key.get_secret_value(),
|
||||
algorithms=[_SESSION_JWT_ALGORITHM],
|
||||
issuer=SESSION_ISSUER,
|
||||
options={
|
||||
"verify_exp": False,
|
||||
"verify_iat": False,
|
||||
"verify_nbf": False,
|
||||
"require": ["iss", "iat", "exp"],
|
||||
},
|
||||
)
|
||||
except jwt.InvalidSignatureError:
|
||||
return SessionBadSignature()
|
||||
except (jwt.InvalidTokenError, ValueError, TypeError):
|
||||
return SessionMalformed()
|
||||
try:
|
||||
return _SessionClaims.model_validate(payload)
|
||||
except ValidationError:
|
||||
return SessionMalformed()
|
||||
|
|
@ -0,0 +1,135 @@
|
|||
"""Tests for the session-token KDF and the edge/token-endpoint resolvers."""
|
||||
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import (
|
||||
envelope_keys_from_master_key,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credentials import (
|
||||
NotSessionBearer,
|
||||
SessionBearerAdmitted,
|
||||
SessionBearerInvalid,
|
||||
SessionRefreshInvalid,
|
||||
SessionRefreshOpened,
|
||||
is_session_bearer_shaped,
|
||||
open_session_refresh_bearer,
|
||||
resolve_session_bearer,
|
||||
session_keys_from_master_key,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import (
|
||||
SESSION_TTL_SECONDS,
|
||||
MintedSessionToken,
|
||||
SessionPrincipal,
|
||||
mint_session_refresh_token,
|
||||
mint_session_token,
|
||||
)
|
||||
|
||||
NOW = datetime(2026, 1, 1, 12, 0, 0, tzinfo=timezone.utc)
|
||||
MASTER_KEY = "sk-master-key-for-tests"
|
||||
KEYS = session_keys_from_master_key(MASTER_KEY)
|
||||
PRINCIPAL = SessionPrincipal(user_id="user-123", client_id="llm_client_abc")
|
||||
|
||||
|
||||
def _access_token() -> str:
|
||||
minted = mint_session_token(PRINCIPAL, KEYS, NOW)
|
||||
assert isinstance(minted, MintedSessionToken)
|
||||
return minted.token.get_secret_value()
|
||||
|
||||
|
||||
def _refresh_token() -> str:
|
||||
minted = mint_session_refresh_token(PRINCIPAL, KEYS, NOW)
|
||||
assert isinstance(minted, MintedSessionToken)
|
||||
return minted.token.get_secret_value()
|
||||
|
||||
|
||||
def test_kdf_is_deterministic_and_key_length_is_256_bit():
|
||||
again = session_keys_from_master_key(MASTER_KEY)
|
||||
assert again.signing_key.get_secret_value() == KEYS.signing_key.get_secret_value()
|
||||
assert len(bytes.fromhex(KEYS.signing_key.get_secret_value())) == 32
|
||||
|
||||
|
||||
def test_kdf_domain_separated_from_envelope_keys():
|
||||
envelope_keys = envelope_keys_from_master_key(MASTER_KEY)
|
||||
session_signing = KEYS.signing_key.get_secret_value()
|
||||
assert session_signing != envelope_keys.signing_key.get_secret_value()
|
||||
assert session_signing != envelope_keys.encryption_key.get_secret_value()
|
||||
|
||||
|
||||
def test_kdf_differs_across_master_keys():
|
||||
other = session_keys_from_master_key("sk-a-different-master-key")
|
||||
assert other.signing_key.get_secret_value() != KEYS.signing_key.get_secret_value()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"value,expected",
|
||||
[
|
||||
("Bearer sk-1234", False),
|
||||
("sk-1234", False),
|
||||
("Bearer llm_env_abc", False),
|
||||
("Bearer llm_refresh_abc", False),
|
||||
("llm_session_abc", True),
|
||||
("Bearer llm_session_abc", True),
|
||||
("bearer llm_srefresh_abc", True),
|
||||
],
|
||||
)
|
||||
def test_is_session_bearer_shaped(value, expected):
|
||||
assert is_session_bearer_shaped(value) is expected
|
||||
|
||||
|
||||
def test_resolve_admits_valid_access_token_with_and_without_scheme():
|
||||
token = _access_token()
|
||||
for value in (token, f"Bearer {token}", f"bearer {token}"):
|
||||
result = resolve_session_bearer(value, KEYS, NOW)
|
||||
assert isinstance(result, SessionBearerAdmitted)
|
||||
assert result.principal == PRINCIPAL
|
||||
|
||||
|
||||
def test_resolve_passes_non_session_bearers_through():
|
||||
for value in ("Bearer sk-1234", "Bearer llm_env_whatever", "Bearer eyJhbGciOi"):
|
||||
assert isinstance(resolve_session_bearer(value, KEYS, NOW), NotSessionBearer)
|
||||
|
||||
|
||||
def test_resolve_fails_expired_token_closed_and_flags_expiry():
|
||||
token = _access_token()
|
||||
later = NOW + timedelta(seconds=SESSION_TTL_SECONDS + 1)
|
||||
result = resolve_session_bearer(f"Bearer {token}", KEYS, later)
|
||||
assert isinstance(result, SessionBearerInvalid)
|
||||
assert result.expired is True
|
||||
|
||||
|
||||
def test_resolve_fails_tampered_token_closed_without_expiry_flag():
|
||||
token = _access_token()
|
||||
tampered = token[:-2] + ("aa" if not token.endswith("aa") else "bb")
|
||||
result = resolve_session_bearer(f"Bearer {tampered}", KEYS, NOW)
|
||||
assert isinstance(result, SessionBearerInvalid)
|
||||
assert result.expired is False
|
||||
|
||||
|
||||
def test_resolve_rejects_refresh_token_at_the_edge():
|
||||
result = resolve_session_bearer(f"Bearer {_refresh_token()}", KEYS, NOW)
|
||||
assert isinstance(result, SessionBearerInvalid)
|
||||
assert result.expired is False
|
||||
|
||||
|
||||
def test_resolve_wrong_master_key_fails_closed():
|
||||
other_keys = session_keys_from_master_key("sk-rotated-master-key")
|
||||
result = resolve_session_bearer(f"Bearer {_access_token()}", other_keys, NOW)
|
||||
assert isinstance(result, SessionBearerInvalid)
|
||||
|
||||
|
||||
def test_refresh_grant_opens_for_the_issued_client():
|
||||
result = open_session_refresh_bearer(_refresh_token(), KEYS, NOW, expected_client_id="llm_client_abc")
|
||||
assert isinstance(result, SessionRefreshOpened)
|
||||
assert result.principal == PRINCIPAL
|
||||
|
||||
|
||||
def test_refresh_grant_rejects_a_different_client():
|
||||
result = open_session_refresh_bearer(_refresh_token(), KEYS, NOW, expected_client_id="llm_client_other")
|
||||
assert isinstance(result, SessionRefreshInvalid)
|
||||
|
||||
|
||||
def test_refresh_grant_rejects_access_token_presented_as_refresh():
|
||||
result = open_session_refresh_bearer(_access_token(), KEYS, NOW, expected_client_id="llm_client_abc")
|
||||
assert isinstance(result, SessionRefreshInvalid)
|
||||
|
|
@ -0,0 +1,206 @@
|
|||
"""Tests for the identity-only gateway session token (mint/open, hostile-input totality)."""
|
||||
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
||||
import jwt
|
||||
import pytest
|
||||
from pydantic import SecretStr, ValidationError
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import (
|
||||
MAX_SESSION_TOKEN_BYTES,
|
||||
SESSION_ISSUER,
|
||||
SESSION_REFRESH_PREFIX,
|
||||
SESSION_REFRESH_TTL_SECONDS,
|
||||
SESSION_TOKEN_PREFIX,
|
||||
SESSION_TTL_SECONDS,
|
||||
MintedSessionToken,
|
||||
NotASessionToken,
|
||||
OpenedSessionToken,
|
||||
SessionBadSignature,
|
||||
SessionExpired,
|
||||
SessionKeys,
|
||||
SessionMalformed,
|
||||
SessionPrincipal,
|
||||
SessionTokenTooLarge,
|
||||
is_session_refresh_token,
|
||||
is_session_token,
|
||||
mint_session_refresh_token,
|
||||
mint_session_token,
|
||||
open_session_refresh_token,
|
||||
open_session_token,
|
||||
)
|
||||
|
||||
NOW = datetime(2026, 1, 1, 12, 0, 0, tzinfo=timezone.utc)
|
||||
KEYS = SessionKeys(signing_key=SecretStr("k" * 32))
|
||||
OTHER_KEYS = SessionKeys(signing_key=SecretStr("x" * 32))
|
||||
PRINCIPAL = SessionPrincipal(user_id="user-123", client_id="llm_client_abc")
|
||||
|
||||
|
||||
def _mint_access() -> str:
|
||||
minted = mint_session_token(PRINCIPAL, KEYS, NOW)
|
||||
assert isinstance(minted, MintedSessionToken)
|
||||
return minted.token.get_secret_value()
|
||||
|
||||
|
||||
def _mint_refresh() -> str:
|
||||
minted = mint_session_refresh_token(PRINCIPAL, KEYS, NOW)
|
||||
assert isinstance(minted, MintedSessionToken)
|
||||
return minted.token.get_secret_value()
|
||||
|
||||
|
||||
def _sign_claims(payload: dict, prefix: str = SESSION_TOKEN_PREFIX, keys: SessionKeys = KEYS) -> str:
|
||||
return prefix + jwt.encode(payload, keys.signing_key.get_secret_value(), algorithm="HS256")
|
||||
|
||||
|
||||
def _valid_claims(**overrides) -> dict:
|
||||
base = {
|
||||
"iss": SESSION_ISSUER,
|
||||
"iat": int(NOW.timestamp()),
|
||||
"exp": int((NOW + timedelta(seconds=600)).timestamp()),
|
||||
"kind": "session",
|
||||
"user_id": "user-123",
|
||||
"client_id": "llm_client_abc",
|
||||
}
|
||||
return {**base, **overrides}
|
||||
|
||||
|
||||
def test_access_round_trip_recovers_principal_and_caps_ttl():
|
||||
minted = mint_session_token(PRINCIPAL, KEYS, NOW)
|
||||
assert isinstance(minted, MintedSessionToken)
|
||||
assert minted.expires_at == NOW + timedelta(seconds=SESSION_TTL_SECONDS)
|
||||
token = minted.token.get_secret_value()
|
||||
assert is_session_token(token)
|
||||
assert not is_session_refresh_token(token)
|
||||
opened = open_session_token(token, KEYS, NOW)
|
||||
assert isinstance(opened, OpenedSessionToken)
|
||||
assert opened.principal == PRINCIPAL
|
||||
|
||||
|
||||
def test_refresh_round_trip_recovers_principal_and_caps_ttl():
|
||||
minted = mint_session_refresh_token(PRINCIPAL, KEYS, NOW)
|
||||
assert isinstance(minted, MintedSessionToken)
|
||||
assert minted.expires_at == NOW + timedelta(seconds=SESSION_REFRESH_TTL_SECONDS)
|
||||
token = minted.token.get_secret_value()
|
||||
assert is_session_refresh_token(token)
|
||||
opened = open_session_refresh_token(token, KEYS, NOW)
|
||||
assert isinstance(opened, OpenedSessionToken)
|
||||
assert opened.principal == PRINCIPAL
|
||||
|
||||
|
||||
def test_access_token_reprefixed_as_refresh_is_rejected_by_signed_kind():
|
||||
body = _mint_access().removeprefix(SESSION_TOKEN_PREFIX)
|
||||
swapped = SESSION_REFRESH_PREFIX + body
|
||||
assert isinstance(open_session_refresh_token(swapped, KEYS, NOW), SessionMalformed)
|
||||
|
||||
|
||||
def test_refresh_token_reprefixed_as_access_is_rejected_by_signed_kind():
|
||||
body = _mint_refresh().removeprefix(SESSION_REFRESH_PREFIX)
|
||||
swapped = SESSION_TOKEN_PREFIX + body
|
||||
assert isinstance(open_session_token(swapped, KEYS, NOW), SessionMalformed)
|
||||
|
||||
|
||||
def test_refresh_token_is_not_an_access_token_at_the_edge():
|
||||
assert isinstance(open_session_token(_mint_refresh(), KEYS, NOW), NotASessionToken)
|
||||
|
||||
|
||||
def test_expired_access_token_is_expired_not_malformed():
|
||||
token = _mint_access()
|
||||
at_expiry = NOW + timedelta(seconds=SESSION_TTL_SECONDS)
|
||||
assert isinstance(open_session_token(token, KEYS, at_expiry), SessionExpired)
|
||||
after = NOW + timedelta(seconds=SESSION_TTL_SECONDS + 1)
|
||||
assert isinstance(open_session_token(token, KEYS, after), SessionExpired)
|
||||
|
||||
|
||||
def test_still_valid_one_second_before_expiry():
|
||||
token = _mint_access()
|
||||
just_before = NOW + timedelta(seconds=SESSION_TTL_SECONDS - 1)
|
||||
assert isinstance(open_session_token(token, KEYS, just_before), OpenedSessionToken)
|
||||
|
||||
|
||||
def test_tampered_signature_is_bad_signature():
|
||||
token = _mint_access()
|
||||
tampered = token[:-2] + ("aa" if not token.endswith("aa") else "bb")
|
||||
assert isinstance(open_session_token(tampered, KEYS, NOW), SessionBadSignature)
|
||||
|
||||
|
||||
def test_key_rotation_invalidates_outstanding_tokens():
|
||||
token = _mint_access()
|
||||
assert isinstance(open_session_token(token, OTHER_KEYS, NOW), SessionBadSignature)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"candidate,expected",
|
||||
[
|
||||
("sk-1234", NotASessionToken),
|
||||
("llm_env_something", NotASessionToken),
|
||||
("", NotASessionToken),
|
||||
(SESSION_TOKEN_PREFIX, SessionMalformed),
|
||||
(SESSION_TOKEN_PREFIX + "not-a-jwt", SessionMalformed),
|
||||
(SESSION_TOKEN_PREFIX + "\ud800garbage", SessionMalformed),
|
||||
(SESSION_TOKEN_PREFIX + "a" * (MAX_SESSION_TOKEN_BYTES + 1), SessionMalformed),
|
||||
],
|
||||
)
|
||||
def test_hostile_candidates_never_raise(candidate, expected):
|
||||
assert isinstance(open_session_token(candidate, KEYS, NOW), expected)
|
||||
|
||||
|
||||
def test_multibyte_candidate_over_byte_cap_but_under_char_cap_is_rejected():
|
||||
filler = "€" * (MAX_SESSION_TOKEN_BYTES // 3)
|
||||
candidate = SESSION_TOKEN_PREFIX + filler
|
||||
assert len(candidate) <= MAX_SESSION_TOKEN_BYTES
|
||||
assert isinstance(open_session_token(candidate, KEYS, NOW), SessionMalformed)
|
||||
|
||||
|
||||
def test_alg_none_token_is_rejected():
|
||||
unsigned = jwt.api_jws.encode(b'{"iss":"litellm-mcp-gateway"}', key=None, algorithm="none")
|
||||
assert isinstance(open_session_token(SESSION_TOKEN_PREFIX + unsigned, KEYS, NOW), SessionMalformed)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"claims",
|
||||
[
|
||||
_valid_claims(iss="wrong-issuer"),
|
||||
_valid_claims(exp=str(int((NOW + timedelta(seconds=600)).timestamp()))),
|
||||
_valid_claims(iat="evil"),
|
||||
_valid_claims(kind="access"),
|
||||
_valid_claims(user_id=""),
|
||||
_valid_claims(nbf=0),
|
||||
{k: v for k, v in _valid_claims().items() if k != "client_id"},
|
||||
{k: v for k, v in _valid_claims().items() if k != "exp"},
|
||||
],
|
||||
)
|
||||
def test_signed_but_malformed_claims_are_rejected_without_raising(claims):
|
||||
token = _sign_claims(claims)
|
||||
assert isinstance(open_session_token(token, KEYS, NOW), SessionMalformed)
|
||||
|
||||
|
||||
def test_signed_claims_with_exact_shape_open():
|
||||
token = _sign_claims(_valid_claims())
|
||||
opened = open_session_token(token, KEYS, NOW)
|
||||
assert isinstance(opened, OpenedSessionToken)
|
||||
assert opened.principal.user_id == "user-123"
|
||||
|
||||
|
||||
def test_oversized_client_id_fails_mint_with_typed_error_not_truncation():
|
||||
principal = SessionPrincipal(user_id="user-123", client_id="c" * (MAX_SESSION_TOKEN_BYTES + 100))
|
||||
minted = mint_session_token(principal, KEYS, NOW)
|
||||
assert isinstance(minted, SessionTokenTooLarge)
|
||||
assert minted.max_bytes == MAX_SESSION_TOKEN_BYTES
|
||||
|
||||
|
||||
def test_empty_principal_fields_rejected_at_construction():
|
||||
with pytest.raises(ValidationError):
|
||||
SessionPrincipal(user_id="", client_id="c")
|
||||
with pytest.raises(ValidationError):
|
||||
SessionPrincipal(user_id="u", client_id="")
|
||||
|
||||
|
||||
def test_short_signing_key_rejected_at_construction():
|
||||
with pytest.raises(ValidationError):
|
||||
SessionKeys(signing_key=SecretStr("short"))
|
||||
|
||||
|
||||
def test_minted_token_repr_never_leaks_value():
|
||||
minted = mint_session_token(PRINCIPAL, KEYS, NOW)
|
||||
assert isinstance(minted, MintedSessionToken)
|
||||
assert minted.token.get_secret_value() not in repr(minted)
|
||||
Loading…
Add table
Reference in a new issue