mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-10 22:41:41 +00:00
Fix: Separate OAuth M2M authentication from UI SSO + Handle Introspection endpoint for Oauth2 (#15667)
* add patch for Oauth2 token * fix Oauth2 checks
This commit is contained in:
Ishaan Jaff
GitHub
parent
eee28a7af8
commit
ef28b593e4
2 changed files with 216 additions and 71 deletions
|
|
@ -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}")
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue