diff --git a/enterprise/litellm_enterprise/proxy/auth/custom_sso_handler.py b/enterprise/litellm_enterprise/proxy/auth/custom_sso_handler.py index a368232038..e8f104c262 100644 --- a/enterprise/litellm_enterprise/proxy/auth/custom_sso_handler.py +++ b/enterprise/litellm_enterprise/proxy/auth/custom_sso_handler.py @@ -10,28 +10,21 @@ has already authenticated the user) and you need to extract user information fro custom headers or other request attributes. """ -from typing import TYPE_CHECKING, Dict, Optional, Union, cast +from typing import cast from fastapi import Request from fastapi.responses import RedirectResponse -if TYPE_CHECKING: - from fastapi_sso.sso.base import OpenID -else: - from typing import Any as OpenID - -from litellm.proxy.management_endpoints.types import CustomOpenID - class EnterpriseCustomSSOHandler: """ Enterprise Custom SSO Handler for LiteLLM Proxy - + This class provides methods for handling custom SSO authentication flows where users can implement their own authentication logic by processing request headers and returning user information in OpenID format. """ - + @staticmethod async def handle_custom_ui_sso_sign_in( request: Request, @@ -40,16 +33,16 @@ class EnterpriseCustomSSOHandler: Allow a user to execute their custom code to parse incoming request headers and return a OpenID object Use this when you have an OAuth proxy in front of LiteLLM (where the OAuth proxy has already authenticated the user) - + Args: request: The FastAPI request object containing headers and other request data - + Returns: RedirectResponse: Redirect response that sends the user to the LiteLLM UI with authentication token - + Raises: ValueError: If custom_ui_sso_sign_in_handler is not configured - + Example: This method is typically called when a user has already been authenticated by an external OAuth proxy and the proxy has added custom headers containing user information. @@ -60,27 +53,44 @@ class EnterpriseCustomSSOHandler: from litellm.integrations.custom_sso_handler import CustomSSOLoginHandler from litellm.proxy.proxy_server import ( CommonProxyErrors, + general_settings, premium_user, user_custom_ui_sso_sign_in_handler, ) + from litellm.proxy.auth.trusted_proxy_utils import ( + require_trusted_proxy_request, + ) + if premium_user is not True: raise ValueError(CommonProxyErrors.not_premium_user.value) - + if user_custom_ui_sso_sign_in_handler is None: - raise ValueError("custom_ui_sso_sign_in_handler is not configured. Please set it in general_settings.") - - custom_sso_login_handler = cast(CustomSSOLoginHandler, user_custom_ui_sso_sign_in_handler) - openid_response: OpenID = await custom_sso_login_handler.handle_custom_ui_sso_sign_in( + raise ValueError( + "custom_ui_sso_sign_in_handler is not configured. Please set it in general_settings." + ) + + require_trusted_proxy_request( request=request, + general_settings=general_settings, + feature_name="Custom UI SSO", ) - + + custom_sso_login_handler = cast( + CustomSSOLoginHandler, user_custom_ui_sso_sign_in_handler + ) + openid_response: OpenID = ( + await custom_sso_login_handler.handle_custom_ui_sso_sign_in( + request=request, + ) + ) + # Import here to avoid circular imports from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler - + return await SSOAuthenticationHandler.get_redirect_response_from_openid( result=openid_response, request=request, received_response=None, generic_client_id=None, ui_access_mode=None, - ) \ No newline at end of file + ) diff --git a/litellm/integrations/custom_sso_handler.py b/litellm/integrations/custom_sso_handler.py index 7f60decabc..202e488e0e 100644 --- a/litellm/integrations/custom_sso_handler.py +++ b/litellm/integrations/custom_sso_handler.py @@ -18,6 +18,17 @@ class CustomSSOLoginHandler(CustomLogger): self, request: Request, ) -> OpenID: + from litellm.proxy.auth.trusted_proxy_utils import ( + require_trusted_proxy_request, + ) + from litellm.proxy.proxy_server import general_settings + + require_trusted_proxy_request( + request=request, + general_settings=general_settings, + feature_name="Custom UI SSO", + ) + request_headers_dict = dict(request.headers) return OpenID( id=request_headers_dict.get("x-litellm-user-id"), diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 92c920ca59..5165c7fd50 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2374,6 +2374,10 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): None, description="CIDR ranges of trusted reverse proxies. When set, X-Forwarded-For headers are only trusted from these IPs.", ) + trusted_proxy_ranges: Optional[List[str]] = Field( + None, + description="CIDR ranges of trusted reverse proxies allowed to provide identity headers for header-based auth paths such as enable_oauth2_proxy_auth and custom_ui_sso_sign_in_handler.", + ) store_model_in_db: Optional[bool] = Field( None, description="If True, models and config are stored in and loaded from the database. Default is False.", diff --git a/litellm/proxy/auth/oauth2_proxy_hook.py b/litellm/proxy/auth/oauth2_proxy_hook.py index 389a5b2b9e..9fc4c4fb53 100644 --- a/litellm/proxy/auth/oauth2_proxy_hook.py +++ b/litellm/proxy/auth/oauth2_proxy_hook.py @@ -4,6 +4,7 @@ from fastapi import Request from litellm._logging import verbose_proxy_logger from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.trusted_proxy_utils import require_trusted_proxy_request # OAuth2-proxy header trust is for **identity assertion** from a trusted # upstream auth proxy (oauth2-proxy, Authelia, etc.). The allowlist below @@ -57,6 +58,12 @@ async def handle_oauth2_proxy_request(request: Request) -> UserAPIKeyAuth: from litellm.proxy.proxy_server import general_settings verbose_proxy_logger.debug("Handling oauth2 proxy request") + require_trusted_proxy_request( + request=request, + general_settings=general_settings, + feature_name="OAuth2 proxy auth", + ) + oauth2_config_mappings: Dict[str, str] = ( general_settings.get("oauth2_config_mappings") or {} ) diff --git a/litellm/proxy/auth/trusted_proxy_utils.py b/litellm/proxy/auth/trusted_proxy_utils.py new file mode 100644 index 0000000000..df7b3080f2 --- /dev/null +++ b/litellm/proxy/auth/trusted_proxy_utils.py @@ -0,0 +1,118 @@ +import ipaddress +from typing import Any, Dict, List, Optional, Union + +from fastapi import Request + +from litellm._logging import verbose_proxy_logger + +TRUSTED_PROXY_RANGES_KEY = "trusted_proxy_ranges" +TrustedProxyNetwork = Union[ipaddress.IPv4Network, ipaddress.IPv6Network] + + +def _get_proxy_general_settings() -> Dict[str, Any]: + try: + from litellm.proxy.proxy_server import general_settings + + return general_settings or {} + except ImportError: + return {} + + +def _normalize_cidr_ranges(configured_ranges: Any, *, setting_name: str) -> List[str]: + if not configured_ranges: + return [] + if isinstance(configured_ranges, str): + return [ + raw_range.strip() + for raw_range in configured_ranges.split(",") + if raw_range.strip() + ] + if isinstance(configured_ranges, (list, tuple, set)): + return [ + str(raw_range).strip() + for raw_range in configured_ranges + if str(raw_range).strip() + ] + verbose_proxy_logger.warning( + "Invalid %s value: expected a list of CIDR ranges, got %s", + setting_name, + type(configured_ranges).__name__, + ) + return [] + + +def parse_trusted_proxy_ranges( + configured_ranges: Any, + *, + setting_name: str = TRUSTED_PROXY_RANGES_KEY, +) -> List[TrustedProxyNetwork]: + networks: List[TrustedProxyNetwork] = [] + for cidr in _normalize_cidr_ranges(configured_ranges, setting_name=setting_name): + try: + networks.append(ipaddress.ip_network(cidr, strict=False)) + except ValueError: + verbose_proxy_logger.warning( + "Invalid CIDR in %s: %s, skipping", setting_name, cidr + ) + return networks + + +def _get_direct_client_ip(request: Request) -> Optional[str]: + client = getattr(request, "client", None) + client_host = getattr(client, "host", None) + if isinstance(client_host, str): + return client_host + return None + + +def _is_ip_in_networks( + client_ip: Optional[str], networks: List[TrustedProxyNetwork] +) -> bool: + if not client_ip or not networks: + return False + try: + addr = ipaddress.ip_address(client_ip.strip()) + except ValueError: + return False + return any(addr in network for network in networks) + + +def require_trusted_proxy_request( + *, + request: Request, + general_settings: Optional[Dict[str, Any]] = None, + feature_name: str, + setting_name: str = TRUSTED_PROXY_RANGES_KEY, +) -> None: + """ + Fail closed unless the direct TCP peer is one of the configured + trusted reverse proxies. + + Header-based auth paths must validate the direct peer, not + X-Forwarded-For, because the direct peer is the actor supplying the + identity headers. + """ + if general_settings is None: + general_settings = _get_proxy_general_settings() + + trusted_networks = parse_trusted_proxy_ranges( + general_settings.get(setting_name), setting_name=setting_name + ) + if not trusted_networks: + raise ValueError( + f"{feature_name} requires general_settings.{setting_name} before " + "trusting identity headers from an upstream proxy." + ) + + direct_client_ip = _get_direct_client_ip(request) + if not _is_ip_in_networks(direct_client_ip, trusted_networks): + verbose_proxy_logger.warning( + "%s rejected identity headers from untrusted direct client IP %r", + feature_name, + direct_client_ip, + ) + raise ValueError( + f"{feature_name} only accepts identity headers from configured " + f"trusted proxy ranges. Direct client IP {direct_client_ip!r} " + "is not trusted." + ) diff --git a/tests/test_litellm/proxy/auth/test_oauth2_proxy_hook.py b/tests/test_litellm/proxy/auth/test_oauth2_proxy_hook.py index 9d0bdcf351..dcbfd281e0 100644 --- a/tests/test_litellm/proxy/auth/test_oauth2_proxy_hook.py +++ b/tests/test_litellm/proxy/auth/test_oauth2_proxy_hook.py @@ -3,12 +3,14 @@ Regression tests for the OAuth2-proxy header-forgery fix (GHSA-5c3m-qffq-4r9m). The hook reads HTTP request headers per ``oauth2_config_mappings`` and -constructs a ``UserAPIKeyAuth`` from them. Without the -identity-only allowlist any field could be mapped — including -``user_role``, which Pydantic coerces from the string -``"proxy_admin"`` into ``LitellmUserRoles.PROXY_ADMIN``. An attacker -who reaches the proxy directly (or via a misconfigured reverse -proxy) sets the mapped header and gains full admin privileges. +constructs a ``UserAPIKeyAuth`` from them. The fix has two parts: + +1. Only requests from configured trusted proxy CIDR ranges may provide + identity headers. +2. Only identity fields may be mapped from those headers. Without the + identity-only allowlist any field could be mapped — including + ``user_role``, which Pydantic coerces from the string + ``"proxy_admin"`` into ``LitellmUserRoles.PROXY_ADMIN``. """ import os @@ -27,9 +29,10 @@ from litellm.proxy.auth.oauth2_proxy_hook import ( ) -def _request_with_headers(headers: dict) -> Request: +def _request_with_headers(headers: dict, *, client_host: str = "127.0.0.1") -> Request: scope = { "type": "http", + "client": (client_host, 12345), "headers": [(k.lower().encode(), v.encode()) for k, v in headers.items()], } request = Request(scope=scope) @@ -40,19 +43,24 @@ def _request_with_headers(headers: dict) -> Request: @pytest.fixture def configure_proxy(monkeypatch): """ - Yields a callable that sets ``oauth2_config_mappings`` on the - proxy_server module for the duration of one test. Default mapping - is a single ``user_id -> x-user-id`` (identity-only). + Yields a callable that sets ``oauth2_config_mappings`` and + ``trusted_proxy_ranges`` on the proxy_server module for the duration + of one test. Defaults to a single identity mapping and localhost as + a trusted proxy. """ import litellm.proxy.proxy_server as proxy_server - def _configure(*, mappings=None): + def _configure(*, mappings=None, trusted_proxy_ranges=("127.0.0.1/32",)): if mappings is None: mappings = {"user_id": "x-user-id"} + settings = { + "oauth2_config_mappings": mappings, + "trusted_proxy_ranges": trusted_proxy_ranges, + } monkeypatch.setattr( proxy_server, "general_settings", - {"oauth2_config_mappings": mappings}, + settings, raising=False, ) @@ -70,6 +78,24 @@ async def test_returns_auth_for_simple_user_id_mapping(configure_proxy): assert auth.user_role is None +@pytest.mark.asyncio +async def test_rejects_identity_headers_without_trusted_proxy_ranges(configure_proxy): + configure_proxy(trusted_proxy_ranges=None) + request = _request_with_headers({"x-user-id": "alice"}) + + with pytest.raises(ValueError, match="trusted_proxy_ranges"): + await handle_oauth2_proxy_request(request) + + +@pytest.mark.asyncio +async def test_rejects_identity_headers_from_untrusted_direct_client(configure_proxy): + configure_proxy(trusted_proxy_ranges=["10.0.0.0/24"]) + request = _request_with_headers({"x-user-id": "alice"}, client_host="203.0.113.10") + + with pytest.raises(ValueError, match="not trusted"): + await handle_oauth2_proxy_request(request) + + @pytest.mark.parametrize( "privileged_field", [ diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index eecfcaa035..92759c56cb 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -4,7 +4,6 @@ import os import sys from unittest.mock import AsyncMock, MagicMock, patch -import httpx import pytest from fastapi import HTTPException, Request @@ -25,7 +24,6 @@ from litellm.proxy.management_endpoints.ui_sso import ( SSOAuthenticationHandler, _setup_team_mappings, _sync_user_role_from_jwt_role_map, - determine_role_from_groups, normalize_email, process_sso_jwt_access_token, ) @@ -1849,6 +1847,7 @@ class TestCustomUISSO: "x-forwarded-for": "192.168.1.1", } mock_request.base_url = "https://test.litellm.ai/" + mock_request.client.host = "10.0.0.10" # Mock the custom handler mock_custom_handler = MagicMock(spec=CustomSSOLoginHandler) @@ -1874,36 +1873,73 @@ class TestCustomUISSO: "litellm.proxy.proxy_server.user_custom_ui_sso_sign_in_handler", mock_custom_handler, ): - with patch.object( - SSOAuthenticationHandler, - "get_redirect_response_from_openid", - return_value=mock_redirect_response, - ) as mock_get_redirect: - # Act - result = ( + with patch( + "litellm.proxy.proxy_server.general_settings", + {"trusted_proxy_ranges": ["10.0.0.0/24"]}, + ): + with patch.object( + SSOAuthenticationHandler, + "get_redirect_response_from_openid", + return_value=mock_redirect_response, + ) as mock_get_redirect: + # Act + result = await EnterpriseCustomSSOHandler.handle_custom_ui_sso_sign_in( + request=mock_request + ) + + # Assert + # Verify the custom handler was called with the request + mock_custom_handler.handle_custom_ui_sso_sign_in.assert_called_once_with( + request=mock_request + ) + + # Verify the redirect response was generated with correct OpenID + mock_get_redirect.assert_called_once_with( + result=expected_openid, + request=mock_request, + received_response=None, + generic_client_id=None, + ui_access_mode=None, + ) + + # Verify the result is the redirect response + assert result == mock_redirect_response + assert result.status_code == 303 + + @pytest.mark.asyncio + async def test_handle_custom_ui_sso_sign_in_rejects_untrusted_proxy(self): + """Custom UI SSO rejects spoofed identity headers from direct clients.""" + from enterprise.litellm_enterprise.proxy.auth.custom_sso_handler import ( + EnterpriseCustomSSOHandler, + ) + from litellm.integrations.custom_sso_handler import CustomSSOLoginHandler + + mock_request = MagicMock(spec=Request) + mock_request.headers = { + "x-litellm-user-id": "admin", + "x-litellm-user-email": "admin@example.com", + } + mock_request.base_url = "https://test.litellm.ai/" + mock_request.client.host = "203.0.113.10" + + mock_custom_handler = MagicMock(spec=CustomSSOLoginHandler) + mock_custom_handler.handle_custom_ui_sso_sign_in = AsyncMock() + + with patch("litellm.proxy.proxy_server.premium_user", True): + with patch( + "litellm.proxy.proxy_server.user_custom_ui_sso_sign_in_handler", + mock_custom_handler, + ): + with patch( + "litellm.proxy.proxy_server.general_settings", + {"trusted_proxy_ranges": ["10.0.0.0/24"]}, + ): + with pytest.raises(ValueError, match="not trusted"): await EnterpriseCustomSSOHandler.handle_custom_ui_sso_sign_in( request=mock_request ) - ) - # Assert - # Verify the custom handler was called with the request - mock_custom_handler.handle_custom_ui_sso_sign_in.assert_called_once_with( - request=mock_request - ) - - # Verify the redirect response was generated with correct OpenID - mock_get_redirect.assert_called_once_with( - result=expected_openid, - request=mock_request, - received_response=None, - generic_client_id=None, - ui_access_mode=None, - ) - - # Verify the result is the redirect response - assert result == mock_redirect_response - assert result.status_code == 303 + mock_custom_handler.handle_custom_ui_sso_sign_in.assert_not_called() @pytest.mark.asyncio async def test_custom_ui_sso_handler_execution_with_real_class(self): @@ -1954,6 +1990,7 @@ class TestCustomUISSO: "x-forwarded-for": "10.0.0.1", } mock_request.base_url = "https://custom.litellm.ai/" + mock_request.client.host = "10.0.0.20" # Mock the redirect response method mock_redirect_response = MagicMock() @@ -1964,34 +2001,36 @@ class TestCustomUISSO: "litellm.proxy.proxy_server.user_custom_ui_sso_sign_in_handler", test_handler_instance, ): - with patch.object( - SSOAuthenticationHandler, - "get_redirect_response_from_openid", - return_value=mock_redirect_response, - ) as mock_get_redirect: - # Act - result = ( - await EnterpriseCustomSSOHandler.handle_custom_ui_sso_sign_in( + with patch( + "litellm.proxy.proxy_server.general_settings", + {"trusted_proxy_ranges": ["10.0.0.0/24"]}, + ): + with patch.object( + SSOAuthenticationHandler, + "get_redirect_response_from_openid", + return_value=mock_redirect_response, + ) as mock_get_redirect: + # Act + result = await EnterpriseCustomSSOHandler.handle_custom_ui_sso_sign_in( request=mock_request ) - ) - # Assert that our custom handler was executed - assert test_handler_instance.method_called is True - assert test_handler_instance.received_request == mock_request + # Assert that our custom handler was executed + assert test_handler_instance.method_called is True + assert test_handler_instance.received_request == mock_request - # Verify the redirect response was called with the OpenID from our custom handler - mock_get_redirect.assert_called_once() - call_args = mock_get_redirect.call_args.kwargs + # Verify the redirect response was called with the OpenID from our custom handler + mock_get_redirect.assert_called_once() + call_args = mock_get_redirect.call_args.kwargs - # Verify the OpenID object has the expected values from our custom handler - openid_result = call_args["result"] - assert openid_result.id == "custom_test_user_456" - assert openid_result.email == "custom@example.com" - assert openid_result.first_name == "Custom" - assert openid_result.last_name == "Handler" - assert openid_result.display_name == "Custom Handler Test" - assert openid_result.provider == "custom" + # Verify the OpenID object has the expected values from our custom handler + openid_result = call_args["result"] + assert openid_result.id == "custom_test_user_456" + assert openid_result.email == "custom@example.com" + assert openid_result.first_name == "Custom" + assert openid_result.last_name == "Handler" + assert openid_result.display_name == "Custom Handler Test" + assert openid_result.provider == "custom" # Verify the request and other parameters were passed correctly assert call_args["request"] == mock_request