From 2a726aa659ebd4218f7b2b4c9d950ca9c935a02f Mon Sep 17 00:00:00 2001 From: yucheng Date: Thu, 24 Sep 2026 08:47:29 +0000 Subject: [PATCH] fix(auth): satisfy type gates and adapt unknown-alg test for approved JWT algorithms (LIT-8429) Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../mcp_server/mcp_server_manager.py | 2 +- litellm/proxy/auth/handle_jwt.py | 49 +++++++++---------- .../mcp_jwt_signer/mcp_jwt_signer.py | 4 +- .../test_token_endpoint.py | 24 +++++---- ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 +- 5 files changed, 43 insertions(+), 38 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 2ff52dd5edc..e5ff4f72867 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -1485,7 +1485,7 @@ def _warn_on_shared_identifier_prefixes(servers: Iterable[MCPServer]) -> None: def _stored_client_assertion_signing_alg(value: object, server_name: str) -> ApprovedJwtAlgorithm: if value in APPROVED_JWT_ALGORITHMS: - return cast(ApprovedJwtAlgorithm, value) + return cast(ApprovedJwtAlgorithm, value) # cast-ok: value was checked against APPROVED_JWT_ALGORITHMS if value is not None: verbose_logger.warning( "MCP server %s: client_assertion_signing_alg %r is not an approved algorithm, using RS256", diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 0e0d0376578..4fff3b99469 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -221,7 +221,10 @@ class JWTHandler: # Supported algos: https://pyjwt.readthedocs.io/en/stable/algorithms.html # "Warning: Make sure not to mix symmetric and asymmetric algorithms that interpret # the key in different ways (e.g. HS* and RS*)." - SUPPORTED_JWT_ALGORITHMS = [*APPROVED_JWT_ALGORITHMS, *LEGACY_JWT_ALGORITHMS] + SUPPORTED_JWT_ALGORITHMS = [ # mutable-ok: list kept for backward compatibility + *APPROVED_JWT_ALGORITHMS, + *LEGACY_JWT_ALGORITHMS, + ] LITELLM_JWT_ISSUER_CLAIM = "_litellm_jwt_issuer" LITELLM_USER_ID_CLAIM = "_litellm_user_id" LITELLM_USER_EMAIL_CLAIM = "_litellm_user_email" @@ -955,10 +958,11 @@ class JWTHandler: log_context=f"kid={kid}", ) + allowed: Final = self.allowed_algorithms() usable_keys: Final[JWKKeyValue] = ( - list(jwks_keys_for(keys, self.allowed_algorithms())) + list(jwks_keys_for(keys, allowed)) # mutable-ok: parse_keys consumes a JWKKeyValue list if isinstance(keys, list) - else next(iter(jwks_keys_for((keys,), self.allowed_algorithms())), {}) + else next(iter(jwks_keys_for((keys,), allowed)), {}) # mutable-ok: single-key dict is a JWKKeyValue ) public_key: Final = self.parse_keys(keys=usable_keys, kid=kid) if public_key is not None: @@ -1231,32 +1235,25 @@ class JWTHandler: ) ) - if isinstance(public_key, dict): - public_key_obj: Final = PyJWK.from_dict(self._get_jwk_from_public_key(public_key=public_key)).key - payload: Final = jwt.decode( - token, - public_key_obj, - algorithms=self.allowed_algorithms(), - options=decode_options, - audience=audience, - issuer=issuer, - leeway=self.leeway, - ) - else: - cert: Final = x509.load_pem_x509_certificate(public_key.encode(), default_backend()) - key: Final = cert.public_key().public_bytes( + key_obj: Final = ( + PyJWK.from_dict(self._get_jwk_from_public_key(public_key=public_key)).key + if isinstance(public_key, dict) + else x509.load_pem_x509_certificate(public_key.encode(), default_backend()) + .public_key() + .public_bytes( serialization.Encoding.PEM, serialization.PublicFormat.SubjectPublicKeyInfo, ) - payload = jwt.decode( - token, - key, - algorithms=self.allowed_algorithms(), - audience=audience, - issuer=issuer, - options=decode_options, - leeway=self.leeway, - ) + ) + payload: Final = jwt.decode( + token, + key_obj, + algorithms=self.allowed_algorithms(), + audience=audience, + issuer=issuer, + options=decode_options, + leeway=self.leeway, + ) self._warn_deprecated_signing_algorithm(token) return payload diff --git a/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py b/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py index dd0611176d7..1a3fab21b7c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py +++ b/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py @@ -460,7 +460,9 @@ class MCPJWTSigner(CustomGuardrail): from jwt import PyJWKSet try: - jwks_set: Final = PyJWKSet.from_dict({"keys": list(approved_keys)}) + jwks_set: Final = PyJWKSet.from_dict( + {"keys": list(approved_keys)} # mutable-ok: PyJWKSet.from_dict takes a dict + ) except Exception as exc: raise jwt.exceptions.PyJWKSetError(f"Failed to parse JWKS from {jwks_uri!r}: {exc}") from exc diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_token_endpoint.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_token_endpoint.py index db3a1a386a3..b306e53c0b2 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_token_endpoint.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_token_endpoint.py @@ -13,11 +13,12 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx import jwt -import litellm import pytest from cryptography.hazmat.primitives import serialization -from cryptography.hazmat.primitives.asymmetric import rsa +from cryptography.hazmat.primitives.asymmetric import ec, rsa +from pydantic import SecretStr +import litellm from litellm.proxy._experimental.mcp_server.outbound_credentials.result import ( Error, Ok, @@ -34,12 +35,8 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( CredError, PrivateKeyJwtAuth, ) -from pydantic import SecretStr -_PATCH_TARGET = ( - "litellm.proxy._experimental.mcp_server.outbound_credentials." - "token_endpoint.get_async_httpx_client" -) +_PATCH_TARGET = "litellm.proxy._experimental.mcp_server.outbound_credentials.token_endpoint.get_async_httpx_client" _ENDPOINT = "https://idp.example.com/oauth2/token" _CLIENT_ID = "litellm-client-id" @@ -50,6 +47,15 @@ _PRIVATE_PEM = _RSA_KEY.private_bytes( serialization.PrivateFormat.PKCS8, serialization.NoEncryption(), ).decode() +_EC_PRIVATE_PEM = ( + ec.generate_private_key(ec.SECP256R1()) + .private_bytes( + serialization.Encoding.PEM, + serialization.PrivateFormat.PKCS8, + serialization.NoEncryption(), + ) + .decode() +) _PUBLIC_PEM = ( _RSA_KEY.public_key() .public_bytes( @@ -185,9 +191,9 @@ async def test_fetch_network_error_maps_to_upstream_unavailable(raised): "auth", [ PrivateKeyJwtAuth(private_key=SecretStr("not-a-pem-key"), signing_alg="RS256"), - PrivateKeyJwtAuth(private_key=SecretStr(_PRIVATE_PEM), signing_alg="XX999"), + PrivateKeyJwtAuth(private_key=SecretStr(_EC_PRIVATE_PEM), signing_alg="RS256"), ], - ids=["garbage-key", "unknown-alg"], + ids=["garbage-key", "key-alg-mismatch"], ) async def test_fetch_unsignable_client_assertion_is_misconfigured_not_a_crash(auth): client = AsyncMock() diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 2938e65cfde..af40e742905 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -34841,7 +34841,7 @@ export interface components { /** Aws Session Token */ aws_session_token?: string | null; /** Client Assertion Signing Alg */ - client_assertion_signing_alg?: string | null; + client_assertion_signing_alg?: ("RS256" | "RS384" | "RS512" | "PS256" | "PS384" | "PS512" | "ES256" | "ES384" | "ES512") | null; /** Client Id */ client_id?: string | null; /** Client Private Key */