diff --git a/litellm/proxy/auth_v2/errors.py b/litellm/proxy/auth_v2/errors.py index a7c5950790f..965589e3a58 100644 --- a/litellm/proxy/auth_v2/errors.py +++ b/litellm/proxy/auth_v2/errors.py @@ -3,6 +3,9 @@ from __future__ import annotations from typing import Optional from fastapi import HTTPException +from fastapi.responses import JSONResponse + +SCIM_ERROR_SCHEMA = "urn:ietf:params:scim:api:messages:2.0:Error" class AuthError(HTTPException): @@ -52,3 +55,16 @@ def forbidden_permission() -> AuthError: def account_disabled() -> AuthError: return AuthError(403, "Account disabled") + + +def scim_error_response(exc: Exception) -> JSONResponse: + status_code = exc.status_code if isinstance(exc, HTTPException) else 500 + detail = exc.detail if isinstance(exc, HTTPException) else "Internal server error" + return JSONResponse( + status_code=status_code, + content={ + "schemas": [SCIM_ERROR_SCHEMA], + "status": str(status_code), + "detail": str(detail), + }, + ) diff --git a/litellm/proxy/auth_v2/oidc/__init__.py b/litellm/proxy/auth_v2/oidc/__init__.py index 0e9c3e12a15..1112293efd3 100644 --- a/litellm/proxy/auth_v2/oidc/__init__.py +++ b/litellm/proxy/auth_v2/oidc/__init__.py @@ -1,4 +1,9 @@ from .config import OIDCProviderConfig -from .router import build_oidc_router +from .router import _provider_key, _user_from_userinfo, build_oidc_router -__all__ = ["OIDCProviderConfig", "build_oidc_router"] +__all__ = [ + "OIDCProviderConfig", + "build_oidc_router", + "_provider_key", + "_user_from_userinfo", +] diff --git a/litellm/proxy/auth_v2/saml/__init__.py b/litellm/proxy/auth_v2/saml/__init__.py index c6eb35a995e..c8e4e7548bb 100644 --- a/litellm/proxy/auth_v2/saml/__init__.py +++ b/litellm/proxy/auth_v2/saml/__init__.py @@ -1,4 +1,15 @@ from .config import SAMLConfig -from .router import build_saml_router +from .router import ( + _map_attributes, + _metadata_source, + _user_from_mapped, + build_saml_router, +) -__all__ = ["SAMLConfig", "build_saml_router"] +__all__ = [ + "SAMLConfig", + "build_saml_router", + "_map_attributes", + "_metadata_source", + "_user_from_mapped", +] diff --git a/litellm/proxy/auth_v2/saml/router.py b/litellm/proxy/auth_v2/saml/router.py index c5018ac4169..20569d7324b 100644 --- a/litellm/proxy/auth_v2/saml/router.py +++ b/litellm/proxy/auth_v2/saml/router.py @@ -1,7 +1,7 @@ from __future__ import annotations import time -from typing import TYPE_CHECKING, Any, Dict, List, Optional, cast +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, cast from fastapi import APIRouter, HTTPException, Request from fastapi.responses import RedirectResponse, Response @@ -129,13 +129,44 @@ def build_sp_client(config: SAMLConfig) -> Saml2Client: class SAMLProtocolStore: - def __init__(self, replay_ttl_seconds: int) -> None: - self.outstanding: Dict[str, str] = {} + def __init__( + self, + replay_ttl_seconds: int, + outstanding_ttl_seconds: int = 300, + max_outstanding: int = 10000, + ) -> None: + self._outstanding: Dict[str, Tuple[float, str]] = {} self._seen_assertions: Dict[str, float] = {} self._replay_ttl = replay_ttl_seconds + self._outstanding_ttl = outstanding_ttl_seconds + self._max_outstanding = max_outstanding def remember_request(self, request_id: str, relay_state: str) -> None: - self.outstanding[request_id] = relay_state + now = time.time() + self._evict_outstanding(now) + self._outstanding[request_id] = (now + self._outstanding_ttl, relay_state) + + def outstanding_relays(self) -> Dict[str, str]: + now = time.time() + return { + rid: relay for rid, (exp, relay) in self._outstanding.items() if exp >= now + } + + def consume_request(self, request_id: str) -> Optional[str]: + entry = self._outstanding.pop(request_id, None) + if entry is None: + return None + expires_at, relay = entry + return relay if expires_at >= time.time() else None + + def _evict_outstanding(self, now: float) -> None: + for rid in [r for r, (exp, _) in self._outstanding.items() if exp < now]: + self._outstanding.pop(rid, None) + overflow = len(self._outstanding) - self._max_outstanding + 1 + if overflow > 0: + oldest = sorted(self._outstanding, key=lambda r: self._outstanding[r][0]) + for rid in oldest[:overflow]: + self._outstanding.pop(rid, None) def consume_assertion(self, assertion_id: str) -> bool: now = time.time() @@ -185,7 +216,7 @@ def build_saml_router(auth: "AuthSecurity") -> APIRouter: authn_response = client.parse_authn_request_response( saml_response, BINDING_HTTP_POST, - outstanding=protocol.outstanding or None, + outstanding=protocol.outstanding_relays() or None, ) except Exception as exc: raise HTTPException( @@ -196,7 +227,7 @@ def build_saml_router(auth: "AuthSecurity") -> APIRouter: in_response_to = getattr(authn_response, "in_response_to", None) bound_relay = ( - protocol.outstanding.pop(in_response_to, None) if in_response_to else None + protocol.consume_request(in_response_to) if in_response_to else None ) assertion = getattr(authn_response, "assertion", None) diff --git a/litellm/proxy/auth_v2/scim/router.py b/litellm/proxy/auth_v2/scim/router.py index e31cd2ce743..a106446f367 100644 --- a/litellm/proxy/auth_v2/scim/router.py +++ b/litellm/proxy/auth_v2/scim/router.py @@ -52,7 +52,9 @@ def _error(status_code: int, detail: str) -> JSONResponse: class _ScimRoute(APIRoute): """Render authentication failures with the SCIM Error schema (RFC 7644).""" - def get_route_handler(self) -> Callable[[Request], Coroutine[Any, Any, Response]]: + def get_route_handler( # type: ignore[override] + self, + ) -> Callable[[Request], Coroutine[Any, Any, Response]]: handler = super().get_route_handler() async def scim_handler(request: Request) -> Response: diff --git a/litellm/proxy/auth_v2/security.py b/litellm/proxy/auth_v2/security.py index d311019bec2..ef0bc89e782 100644 --- a/litellm/proxy/auth_v2/security.py +++ b/litellm/proxy/auth_v2/security.py @@ -40,6 +40,9 @@ class AuthSecurity: uvicorn with ``--no-proxy-headers`` and let this module resolve the client IP, or leave ``trusted_proxy_cidrs`` empty and rely on uvicorn's ``--forwarded-allow-ips``. Do not enable both. + + To return SCIM-shaped error bodies, register ``errors.scim_error_response`` as + the host app's exception handler for the SCIM routes. """ def __init__(