tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-05 10:24:03 +00:00
Code Issues Packages Projects Releases Wiki Activity

chore(auth): require trusted proxy for header identity auth

Browse Source
This commit is contained in:
user
2026-04-29 21:20:21 -07:00
parent 722bc63e37
commit 2f4641752b
7 changed files with 300 additions and 85 deletions
Download Patch File Download Diff File Expand all files Collapse all files
enterprise/litellm_enterprise/proxy/auth/custom_sso_handler.py
+32 -22
View File
@@ -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,
)
)
litellm/integrations/custom_sso_handler.py
+11
View File
@@ -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"),
litellm/proxy/_types.py
+4
View File
@@ -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.",
litellm/proxy/auth/oauth2_proxy_hook.py
+7
View File
@@ -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 {}
)
litellm/proxy/auth/trusted_proxy_utils.py
+118
View File
@@ -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."
)
tests/test_litellm/proxy/auth/test_oauth2_proxy_hook.py
+38 -12
View File
@@ -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",
[
tests/test_litellm/proxy/management_endpoints/test_ui_sso.py
+90 -51
View File
@@ -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