diff --git a/litellm/proxy/auth/oauth2_check.py b/litellm/proxy/auth/oauth2_check.py index 4851c27012b..c9d6d4b7df8 100644 --- a/litellm/proxy/auth/oauth2_check.py +++ b/litellm/proxy/auth/oauth2_check.py @@ -1,80 +1,222 @@ -from litellm.proxy._types import UserAPIKeyAuth +import base64 +import os +from typing import Dict, Tuple, cast + +import httpx + +from litellm._logging import verbose_proxy_logger +from litellm.llms.custom_httpx.http_handler import ( + get_async_httpx_client, + httpxSpecialProvider, +) +from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth -async def check_oauth2_token(token: str) -> UserAPIKeyAuth: +class Oauth2Handler: """ - Makes a request to the token info endpoint to validate the OAuth2 token. - - Args: - token (str): The OAuth2 token to validate. - - Returns: - Literal[True]: If the token is valid. - - Raises: - ValueError: If the token is invalid, the request fails, or the token info endpoint is not set. + Handles OAuth2 token validation. """ - import os - import httpx + @staticmethod + def _is_introspection_endpoint( + token_info_endpoint: str, + oauth_client_id: str | None, + oauth_client_secret: str | None, + ) -> bool: + """ + Determine if this is an introspection endpoint (requires POST) or token info endpoint (uses GET). - from litellm._logging import verbose_proxy_logger - from litellm.llms.custom_httpx.http_handler import ( - get_async_httpx_client, - httpxSpecialProvider, - ) - from litellm.proxy._types import CommonProxyErrors - from litellm.proxy.proxy_server import premium_user + Args: + token_info_endpoint: The OAuth2 endpoint URL + oauth_client_id: OAuth2 client ID + oauth_client_secret: OAuth2 client secret - if premium_user is not True: - raise ValueError( - "Oauth2 token validation is only available for premium users" - + CommonProxyErrors.not_premium_user.value + Returns: + bool: True if this is an introspection endpoint + """ + return ( + "introspect" in token_info_endpoint.lower() + and oauth_client_id is not None + and oauth_client_secret is not None ) - verbose_proxy_logger.debug("Oauth2 token validation for token=%s", token) - # Get the token info endpoint from environment variable - token_info_endpoint = os.getenv("OAUTH_TOKEN_INFO_ENDPOINT") - user_id_field_name = os.environ.get("OAUTH_USER_ID_FIELD_NAME", "sub") - user_role_field_name = os.environ.get("OAUTH_USER_ROLE_FIELD_NAME", "role") - user_team_id_field_name = os.environ.get("OAUTH_USER_TEAM_ID_FIELD_NAME", "team_id") + @staticmethod + def _prepare_introspection_request( + token: str, + oauth_client_id: str | None, + oauth_client_secret: str | None, + ) -> Tuple[Dict[str, str], Dict[str, str]]: + """ + Prepare headers and data for OAuth2 introspection endpoint (RFC 7662). - if not token_info_endpoint: - raise ValueError("OAUTH_TOKEN_INFO_ENDPOINT environment variable is not set") + Args: + token: The OAuth2 token to validate + oauth_client_id: OAuth2 client ID + oauth_client_secret: OAuth2 client secret - client = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Check) - headers = {"Authorization": f"Bearer {token}", "Content-Type": "application/json"} + Returns: + Tuple of (headers, data) for the introspection request + """ + headers = {"Content-Type": "application/x-www-form-urlencoded"} + data = {"token": token} - try: - response = await client.get(token_info_endpoint, headers=headers) + # Add client authentication if credentials are provided + if oauth_client_id and oauth_client_secret: + # Use HTTP Basic authentication for client credentials + credentials = base64.b64encode( + f"{oauth_client_id}:{oauth_client_secret}".encode() + ).decode() + headers["Authorization"] = f"Basic {credentials}" + elif oauth_client_id: + # For public clients, include client_id in the request body + data["client_id"] = oauth_client_id - # if it's a bad token we expect it to raise an HTTPStatusError - response.raise_for_status() + return headers, data - # If we get here, the request was successful - data = response.json() + @staticmethod + def _prepare_token_info_request(token: str) -> Dict[str, str]: + """ + Prepare headers for generic token info endpoint. - verbose_proxy_logger.debug( - "Oauth2 token validation for token=%s, response from /token/info=%s", - token, - data, + Args: + token: The OAuth2 token to validate + + Returns: + Dict of headers for the token info request + """ + return {"Authorization": f"Bearer {token}", "Content-Type": "application/json"} + + @staticmethod + def _extract_user_info( + response_data: Dict, + user_id_field_name: str, + user_role_field_name: str, + user_team_id_field_name: str, + ) -> Tuple[str | None, str | None, str | None]: + """ + Extract user information from OAuth2 response. + + Args: + response_data: The response data from OAuth2 endpoint + user_id_field_name: Field name for user ID + user_role_field_name: Field name for user role + user_team_id_field_name: Field name for team ID + + Returns: + Tuple of (user_id, user_role, user_team_id) + """ + user_id = response_data.get(user_id_field_name) + user_team_id = response_data.get(user_team_id_field_name) + user_role = response_data.get(user_role_field_name) + + return user_id, user_role, user_team_id + + @staticmethod + async def check_oauth2_token(token: str) -> UserAPIKeyAuth: + """ + Makes a request to the token introspection endpoint to validate the OAuth2 token. + + This function implements OAuth2 token introspection according to RFC 7662. + It supports both generic token info endpoints (GET) and OAuth2 introspection endpoints (POST). + + Args: + token (str): The OAuth2 token to validate. + + Returns: + UserAPIKeyAuth: If the token is valid, containing user information. + + Raises: + ValueError: If the token is invalid, the request fails, or the token info endpoint is not set. + """ + from litellm.proxy.proxy_server import premium_user + + if premium_user is not True: + raise ValueError( + "Oauth2 token validation is only available for premium users" + + CommonProxyErrors.not_premium_user.value + ) + + verbose_proxy_logger.debug("Oauth2 token validation for token=%s", token) + + # Get the token info endpoint from environment variable + token_info_endpoint = os.getenv("OAUTH_TOKEN_INFO_ENDPOINT") + user_id_field_name = os.environ.get("OAUTH_USER_ID_FIELD_NAME", "sub") + user_role_field_name = os.environ.get("OAUTH_USER_ROLE_FIELD_NAME", "role") + user_team_id_field_name = os.environ.get( + "OAUTH_USER_TEAM_ID_FIELD_NAME", "team_id" ) - # You might want to add additional checks here based on the response - # For example, checking if the token is expired or has the correct scope - user_id = data.get(user_id_field_name) - user_team_id = data.get(user_team_id_field_name) - user_role = data.get(user_role_field_name) + # OAuth2 client credentials for introspection endpoint authentication + oauth_client_id = os.environ.get("OAUTH_CLIENT_ID") + oauth_client_secret = os.environ.get("OAUTH_CLIENT_SECRET") - return UserAPIKeyAuth( - api_key=token, - team_id=user_team_id, - user_id=user_id, - user_role=user_role, + if not token_info_endpoint: + raise ValueError( + "OAUTH_TOKEN_INFO_ENDPOINT environment variable is not set" + ) + + client = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Check) + + # Determine if this is an introspection endpoint (requires POST) or token info endpoint (uses GET) + is_introspection_endpoint = Oauth2Handler._is_introspection_endpoint( + token_info_endpoint=token_info_endpoint, + oauth_client_id=oauth_client_id, + oauth_client_secret=oauth_client_secret, ) - except httpx.HTTPStatusError as e: - # This will catch any 4xx or 5xx errors - raise ValueError(f"Oauth 2.0 Token validation failed: {e}") - except Exception as e: - # This will catch any other errors (like network issues) - raise ValueError(f"An error occurred during token validation: {e}") + + try: + if is_introspection_endpoint: + # OAuth2 Token Introspection (RFC 7662) - requires POST with form data + verbose_proxy_logger.debug("Using OAuth2 introspection endpoint (POST)") + + headers, data = Oauth2Handler._prepare_introspection_request( + token=token, + oauth_client_id=oauth_client_id, + oauth_client_secret=oauth_client_secret, + ) + + response = await client.post( + token_info_endpoint, headers=headers, data=data + ) + else: + # Generic token info endpoint - uses GET with Bearer token + verbose_proxy_logger.debug("Using generic token info endpoint (GET)") + headers = Oauth2Handler._prepare_token_info_request(token=token) + response = await client.get(token_info_endpoint, headers=headers) + + # if it's a bad token we expect it to raise an HTTPStatusError + response.raise_for_status() + + # If we get here, the request was successful + data = response.json() + + verbose_proxy_logger.debug( + "Oauth2 token validation for token=%s, response from endpoint=%s", + token, + data, + ) + + # For introspection endpoints, check if token is active + if is_introspection_endpoint and not data.get("active", True): + raise ValueError("Token is not active") + + # Extract user information from response + user_id, user_role, user_team_id = Oauth2Handler._extract_user_info( + response_data=data, + user_id_field_name=user_id_field_name, + user_role_field_name=user_role_field_name, + user_team_id_field_name=user_team_id_field_name, + ) + + return UserAPIKeyAuth( + api_key=token, + team_id=user_team_id, + user_id=user_id, + user_role=cast(LitellmUserRoles, user_role), + ) + except httpx.HTTPStatusError as e: + # This will catch any 4xx or 5xx errors + raise ValueError(f"Oauth 2.0 Token validation failed: {e}") + except Exception as e: + # This will catch any other errors (like network issues) + raise ValueError(f"An error occurred during token validation: {e}") diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 63047f37b1c..a92e8c8cbd4 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -47,7 +47,7 @@ from litellm.proxy.auth.auth_utils import ( route_in_additonal_public_routes, ) from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler -from litellm.proxy.auth.oauth2_check import check_oauth2_token +from litellm.proxy.auth.oauth2_check import Oauth2Handler from litellm.proxy.auth.oauth2_proxy_hook import handle_oauth2_proxy_request from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.common_utils.http_parsing_utils import ( @@ -455,17 +455,20 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 ########## End of Route Checks Before Reading DB / Cache for "token" ######## if general_settings.get("enable_oauth2_auth", False) is True: - # return UserAPIKeyAuth object - # helper to check if the api_key is a valid oauth2 token - from litellm.proxy.proxy_server import premium_user + # Only apply OAuth2 M2M authentication to LLM API routes, not UI/management routes + # This allows UI SSO to work separately from API M2M authentication + if RouteChecks.is_llm_api_route(route=route): + # return UserAPIKeyAuth object + # helper to check if the api_key is a valid oauth2 token + from litellm.proxy.proxy_server import premium_user - if premium_user is not True: - raise ValueError( - "Oauth2 token validation is only available for premium users" - + CommonProxyErrors.not_premium_user.value - ) + if premium_user is not True: + raise ValueError( + "Oauth2 token validation is only available for premium users" + + CommonProxyErrors.not_premium_user.value + ) - return await check_oauth2_token(token=api_key) + return await Oauth2Handler.check_oauth2_token(token=api_key) if general_settings.get("enable_oauth2_proxy_auth", False) is True: return await handle_oauth2_proxy_request(request=request)