From 3737d6a1f39d09019337b10190a2cd6d8e2d3d08 Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Wed, 22 Apr 2026 23:48:32 +0000 Subject: [PATCH] fix(auth): centralize common_checks to close authorization bypass MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Multiple paths through _user_api_key_auth_builder returned a UserAPIKeyAuth without running common_checks(): OAuth2 token validation, OAuth2 proxy header hook, JWT admin shortcut, master_key path, pass-through custom headers, the /user/auth route, and the allow_requests_on_db_unavailable fallback. An operator-configured key model-access list, max_budget, team_blocked flag, or team model scope was therefore silently skipped on those paths. The HA-fallback token was worse: it was a full proxy-admin synthetic, so a DB outage granted full admin to every caller. Fix three root causes (VERIA-18): 1. Centralize common_checks in the user_api_key_auth wrapper. The builder paths no longer call it; the wrapper runs it once after the builder returns, for every path. Introduces _run_centralized_common_checks which gathers team/user/project/end_user/global_spend context in parallel via asyncio.gather. Preserves the existing custom_auth_run_common_checks opt-out for custom-auth deployments. 2. Narrow is_database_connection_error — drop the blanket PrismaError catch that routed data-layer errors (UniqueViolationError, etc.) into the HA fallback. Only real connectivity failures plus the no_db_connection marker now qualify. 3. DB-unavailable fallback issues an INTERNAL_USER token with user_id DB_UNAVAILABLE_FALLBACK_USER_ID instead of proxy-admin. An outage can no longer escalate an anonymous caller. JWT admin / master_key tokens still grant admin via a synthesized admin user_object (so non_proxy_admin_allowed_routes_check in common_checks recognizes them); other common_checks branches (team_blocked, team_model_access) now apply uniformly. --- litellm/proxy/auth/auth_exception_handler.py | 23 +- litellm/proxy/auth/user_api_key_auth.py | 376 +++++++++++++----- litellm/proxy/db/exception_handler.py | 14 +- .../test_key_generate_prisma.py | 9 +- .../proxy/auth/test_auth_checks.py | 107 +++-- .../proxy/auth/test_auth_exception_handler.py | 65 ++- .../auth/test_custom_auth_end_user_budget.py | 12 +- .../proxy/auth/test_user_api_key_auth.py | 265 ++++++++++++ 8 files changed, 680 insertions(+), 191 deletions(-) diff --git a/litellm/proxy/auth/auth_exception_handler.py b/litellm/proxy/auth/auth_exception_handler.py index 9c306acd2c..5ded8136ef 100644 --- a/litellm/proxy/auth/auth_exception_handler.py +++ b/litellm/proxy/auth/auth_exception_handler.py @@ -8,11 +8,22 @@ from fastapi import HTTPException, Request, status import litellm from litellm._logging import verbose_proxy_logger -from litellm.proxy._types import ProxyErrorTypes, ProxyException, UserAPIKeyAuth +from litellm.proxy._types import ( + LitellmUserRoles, + ProxyErrorTypes, + ProxyException, + UserAPIKeyAuth, +) from litellm.proxy.auth.auth_utils import _get_request_ip_address from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.types.services import ServiceTypes +# Sentinel user_id for the synthetic UserAPIKeyAuth issued during a DB +# outage when allow_requests_on_db_unavailable is True. Downstream +# enforcement can key off this value; it must never collide with a real +# user_id. +DB_UNAVAILABLE_FALLBACK_USER_ID = "__db_unavailable_fallback__" + if TYPE_CHECKING: from opentelemetry.trace import Span as _Span @@ -47,7 +58,6 @@ class UserAPIKeyAuthExceptionHandler: """ from litellm.proxy.proxy_server import ( general_settings, - litellm_proxy_admin_name, proxy_logging_obj, ) @@ -63,10 +73,17 @@ class UserAPIKeyAuthExceptionHandler: duration=0.0, ) + # Non-admin restricted token so a DB outage cannot escalate + # an anonymous caller to proxy-admin privileges. + verbose_proxy_logger.warning( + "Auth: DB unavailable — issuing restricted INTERNAL_USER " + "fallback token (allow_requests_on_db_unavailable=True)" + ) return UserAPIKeyAuth( key_name="failed-to-connect-to-db", token="failed-to-connect-to-db", - user_id=litellm_proxy_admin_name, + user_id=DB_UNAVAILABLE_FALLBACK_USER_ID, + user_role=LitellmUserRoles.INTERNAL_USER, request_route=route, ) else: diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index ff4f63593c..8ef3d59043 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -29,6 +29,7 @@ from litellm.proxy.auth.auth_checks import ( _cache_key_object, _delete_cache_key_object, _get_user_role, + _is_model_cost_zero, _is_user_proxy_admin, _virtual_key_max_budget_alert_check, _virtual_key_max_budget_check, @@ -783,22 +784,12 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 user_id = result["user_id"] user_object = result["user_object"] end_user_id = result["end_user_id"] - end_user_object = result["end_user_object"] org_id = result["org_id"] - token = result["token"] team_membership: Optional[LiteLLM_TeamMembership] = result.get( "team_membership", None ) jwt_claims = result.get("jwt_claims", None) - global_proxy_spend = await get_global_proxy_spend( - litellm_proxy_admin_name=litellm_proxy_admin_name, - user_api_key_cache=user_api_key_cache, - prisma_client=prisma_client, - token=token, - proxy_logging_obj=proxy_logging_obj, - ) - if is_proxy_admin: return UserAPIKeyAuth( api_key=None, @@ -895,24 +886,6 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 valid_token.project_metadata = _jwt_project_obj.metadata valid_token.project_alias = _jwt_project_obj.project_alias - # run through common checks - _ = await common_checks( - request=request, - request_body=request_data, - team_object=team_object, - user_object=user_object, - end_user_object=end_user_object, - general_settings=general_settings, - global_proxy_spend=global_proxy_spend, - route=route, - llm_router=llm_router, - proxy_logging_obj=proxy_logging_obj, - valid_token=valid_token, - skip_budget_checks=skip_budget_checks, - project_object=_jwt_project_obj, - ) - - # return UserAPIKeyAuth object return cast(UserAPIKeyAuth, valid_token) #### ELSE #### @@ -1498,22 +1471,6 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 user_info=call_info, ) ) - with tracer.trace("litellm.proxy.auth.common_checks"): - _ = await common_checks( - request=request, - request_body=request_data, - team_object=_team_obj, - user_object=user_obj, - end_user_object=_end_user_object, - general_settings=general_settings, - global_proxy_spend=global_proxy_spend, - route=route, - llm_router=llm_router, - proxy_logging_obj=proxy_logging_obj, - valid_token=valid_token, - skip_budget_checks=skip_budget_checks, - project_object=_project_obj, - ) # Token passed all checks if valid_token is None: raise HTTPException(401, detail="Invalid API key") @@ -1567,6 +1524,245 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 ) +async def _safe_fetch(label: str, awaitable): + """Run an awaitable and return its result, logging and swallowing + HTTPException / ProxyException / PrismaError. Used in the centralized + authz gate so a DB outage fetching a team/user/project yields a + ``None`` object (and ``common_checks`` no-ops the corresponding + branch) rather than failing the request before authorization can + run against whatever limits are recorded on the token itself. + """ + try: + return await awaitable + except (HTTPException, ProxyException) as e: + verbose_proxy_logger.debug( + "centralized auth: %s fetch failed (%s: %s)", + label, + type(e).__name__, + e, + ) + raise + except Exception as e: + verbose_proxy_logger.debug( + "centralized auth: %s fetch swallowed (%s: %s)", + label, + type(e).__name__, + e, + ) + return None + + +def _team_obj_from_token(valid_token: UserAPIKeyAuth) -> LiteLLM_TeamTableCachedObj: + """Reconstruct a cached team object from the fields already on the + UserAPIKeyAuth. Used as a last resort when the team DB row can't be + fetched (e.g. outage) — enforcement still runs against the token's + own team_* fields.""" + return LiteLLM_TeamTableCachedObj( + team_id=valid_token.team_id, + max_budget=valid_token.team_max_budget, + soft_budget=valid_token.team_soft_budget, + spend=valid_token.team_spend, + tpm_limit=valid_token.team_tpm_limit, + rpm_limit=valid_token.team_rpm_limit, + blocked=valid_token.team_blocked, + models=valid_token.team_models, + metadata=valid_token.team_metadata, + object_permission_id=valid_token.team_object_permission_id, + ) + + +@tracer.wrap() +async def _run_centralized_common_checks( + user_api_key_auth_obj: UserAPIKeyAuth, + request: Request, + request_data: dict, + route: str, +) -> None: + """Run ``common_checks`` once at the ``user_api_key_auth`` wrapper + boundary, regardless of which ``_user_api_key_auth_builder`` path + returned. This is the single invariant enforcement point for key + model-access, budgets, guardrails, org, and vector-store checks. + + Invariants: + - ``PROXY_ADMIN`` tokens short-circuit (admins bypass these checks + today; preserving that preserves behavior and avoids five DB + fetches per admin request). + - ``user_custom_auth`` with ``custom_auth_run_common_checks`` unset + skips the gate — matches the existing custom-auth RPS guarantee. + Custom-auth deployments don't use OAuth2 / DB-fallback paths, so + the skip does not re-open any bypass. + """ + from litellm.proxy.proxy_server import ( + general_settings, + litellm_proxy_admin_name, + llm_router, + master_key, + prisma_client, + proxy_logging_obj, + user_api_key_cache, + user_custom_auth, + ) + + # master_key unset = no-auth dev mode. The builder returns an + # INTERNAL_USER token for any api_key and the proxy is + # unauthenticated by configuration; running common_checks would + # block every admin route on these deployments where that was + # previously not the contract. + if master_key is None: + return + + if user_custom_auth is not None and not general_settings.get( + "custom_auth_run_common_checks", False + ): + return + + parent_otel_span = user_api_key_auth_obj.parent_otel_span + end_user_id = get_end_user_id_from_request_body( + request_data, _safe_get_request_headers(request) + ) + + fetch_coros = [] + if user_api_key_auth_obj.team_id is not None: + fetch_coros.append( + _safe_fetch( + "team", + get_team_object( + team_id=user_api_key_auth_obj.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ), + ) + ) + else: + fetch_coros.append(_safe_fetch("team", _noop_none())) + + if user_api_key_auth_obj.user_id is not None: + fetch_coros.append( + _safe_fetch( + "user", + get_user_object( + user_id=user_api_key_auth_obj.user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + user_id_upsert=False, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ), + ) + ) + else: + fetch_coros.append(_safe_fetch("user", _noop_none())) + + if user_api_key_auth_obj.project_id is not None: + fetch_coros.append( + _safe_fetch( + "project", + get_project_object( + project_id=user_api_key_auth_obj.project_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ), + ) + ) + else: + fetch_coros.append(_safe_fetch("project", _noop_none())) + + if end_user_id: + fetch_coros.append( + _safe_fetch( + "end_user", + get_end_user_object( + end_user_id=end_user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + route=route, + ), + ) + ) + else: + fetch_coros.append(_safe_fetch("end_user", _noop_none())) + + fetch_coros.append( + _safe_fetch( + "global_spend", + get_global_proxy_spend( + litellm_proxy_admin_name=litellm_proxy_admin_name, + user_api_key_cache=user_api_key_cache, + prisma_client=prisma_client, + token=user_api_key_auth_obj.token or "", + proxy_logging_obj=proxy_logging_obj, + ), + ) + ) + + try: + ( + team_object, + user_object, + project_object, + end_user_object, + global_proxy_spend, + ) = await asyncio.gather(*fetch_coros, return_exceptions=False) + except HTTPException: + # get_team_object raises HTTPException when the team isn't found. + # Reconstruct from the token so enforcement can still run. + team_object = _team_obj_from_token(user_api_key_auth_obj) + user_object = None + project_object = None + end_user_object = None + global_proxy_spend = None + + # common_checks identifies admin via user_object, not the token + # (non_proxy_admin_allowed_routes_check). JWT admin shortcut and + # master_key tokens have no DB user row — synthesize one so admin + # route access is granted after centralization. + if ( + user_object is None + and user_api_key_auth_obj.user_role == LitellmUserRoles.PROXY_ADMIN + ): + user_object = LiteLLM_UserTable( + user_id=user_api_key_auth_obj.user_id or litellm_proxy_admin_name, + user_role=LitellmUserRoles.PROXY_ADMIN, + spend=0.0, + ) + + if project_object is not None: + user_api_key_auth_obj.project_metadata = project_object.metadata + user_api_key_auth_obj.project_alias = project_object.project_alias + + skip_budget_checks = False + model = get_model_from_request(request_data, route) + if model is not None and llm_router is not None: + skip_budget_checks = _is_model_cost_zero(model=model, llm_router=llm_router) + + _ = await common_checks( + request=request, + request_body=request_data, + team_object=team_object, + user_object=user_object, + end_user_object=end_user_object, + general_settings=general_settings, + global_proxy_spend=global_proxy_spend, + route=route, + llm_router=llm_router, + proxy_logging_obj=proxy_logging_obj, + valid_token=user_api_key_auth_obj, + skip_budget_checks=skip_budget_checks, + project_object=project_object, + ) + + +async def _noop_none() -> None: + """Sentinel coroutine for asyncio.gather when a fetch is unnecessary + (e.g. token has no team_id). Keeps the result tuple positional.""" + return None + + @tracer.wrap() async def user_api_key_auth( request: Request, @@ -1608,6 +1804,28 @@ async def user_api_key_auth( ## ENSURE DISABLE ROUTE WORKS ACROSS ALL USER AUTH FLOWS ## RouteChecks.should_call_route(route=route, valid_token=user_api_key_auth_obj) + # Single authorization point. Builder paths MUST NOT call common_checks. + # Route through the same exception handler the builder uses so + # authorization failures (ProxyException, or plain Exception from + # admin-only-route / model-access / budget checks) surface as + # ProxyException consistently with pre-refactor behavior. + try: + await _run_centralized_common_checks( + user_api_key_auth_obj=user_api_key_auth_obj, + request=request, + request_data=request_data, + route=route, + ) + except Exception as e: + return await UserAPIKeyAuthExceptionHandler._handle_authentication_error( + e=e, + request=request, + request_data=request_data, + route=route, + parent_otel_span=user_api_key_auth_obj.parent_otel_span, + api_key=api_key, + ) + end_user_id = get_end_user_id_from_request_body( request_data, _safe_get_request_headers(request) ) @@ -1929,55 +2147,10 @@ async def _run_post_custom_auth_checks( model=current_model, ) - # 5. Look up user object if user_id is set - user_object = None - if valid_token.user_id is not None: - try: - user_object = await get_user_object( - user_id=valid_token.user_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - user_id_upsert=False, - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - ) - except Exception: - # If user_role is PROXY_ADMIN on the token, create a synthetic user object - # so that admin route checks pass for custom auth - if valid_token.user_role == LitellmUserRoles.PROXY_ADMIN: - user_object = LiteLLM_UserTable( - user_id=valid_token.user_id, - user_role=LitellmUserRoles.PROXY_ADMIN, - spend=0.0, - ) - - # 6. Run common checks - if valid_token.team_id is not None: - try: - _team_obj = await get_team_object( - team_id=valid_token.team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - ) - except HTTPException: - _team_obj = LiteLLM_TeamTableCachedObj( - team_id=valid_token.team_id, - max_budget=valid_token.team_max_budget, - soft_budget=valid_token.team_soft_budget, - spend=valid_token.team_spend, - tpm_limit=valid_token.team_tpm_limit, - rpm_limit=valid_token.team_rpm_limit, - blocked=valid_token.team_blocked, - models=valid_token.team_models, - metadata=valid_token.team_metadata, - object_permission_id=valid_token.team_object_permission_id, - ) - else: - _team_obj = None - - _project_obj = None + # team / user / end_user / project context objects are fetched by + # the centralized common_checks gate in user_api_key_auth after + # this helper returns. Keep only the project fetch here because it + # mutates the token (project_metadata / project_alias). if valid_token.project_id is not None: _project_obj = await get_project_object( project_id=valid_token.project_id, @@ -1989,21 +2162,4 @@ async def _run_post_custom_auth_checks( valid_token.project_metadata = _project_obj.metadata valid_token.project_alias = _project_obj.project_alias - if general_settings.get("custom_auth_run_common_checks", False): - _ = await common_checks( - request=request, - request_body=request_data, - team_object=_team_obj, - user_object=user_object, - end_user_object=end_user_object, - general_settings=general_settings, - global_proxy_spend=None, - route=route, - llm_router=llm_router, - proxy_logging_obj=proxy_logging_obj, - valid_token=valid_token, - skip_budget_checks=False, - project_object=_project_obj, - ) - return valid_token diff --git a/litellm/proxy/db/exception_handler.py b/litellm/proxy/db/exception_handler.py index 213dd39adc..9e3f43c9b5 100644 --- a/litellm/proxy/db/exception_handler.py +++ b/litellm/proxy/db/exception_handler.py @@ -31,16 +31,12 @@ class PrismaDBExceptionHandler: @staticmethod def is_database_connection_error(e: Exception) -> bool: + """Connectivity failures + the ``no_db_connection`` marker only. + Data-layer errors (``UniqueViolationError``, etc.) must NOT match + — matching them would route non-outage exceptions into the HA + fallback and the authentication bypass it grants. """ - Returns True if the exception is from a database outage / connection error. - Any PrismaError qualifies — the DB failed to serve the request. - Used by allow_requests_on_db_unavailable logic and endpoint 503 responses. - """ - import prisma - - if isinstance(e, DB_CONNECTION_ERROR_TYPES): - return True - if isinstance(e, prisma.errors.PrismaError): + if PrismaDBExceptionHandler.is_database_transport_error(e): return True if isinstance(e, ProxyException) and e.type == ProxyErrorTypes.no_db_connection: return True diff --git a/tests/proxy_unit_tests/test_key_generate_prisma.py b/tests/proxy_unit_tests/test_key_generate_prisma.py index d4ad69437c..19ee9f75d8 100644 --- a/tests/proxy_unit_tests/test_key_generate_prisma.py +++ b/tests/proxy_unit_tests/test_key_generate_prisma.py @@ -3816,10 +3816,15 @@ async def test_user_api_key_auth_db_unavailable(): api_key="Bearer sk-123456789", ) - # Verify results + from litellm.proxy.auth.auth_exception_handler import ( + DB_UNAVAILABLE_FALLBACK_USER_ID, + ) + + # Verify results. user_id is the non-admin fallback sentinel so a DB + # outage cannot escalate an anonymous caller to proxy-admin. assert isinstance(result, UserAPIKeyAuth) assert result.key_name == "failed-to-connect-to-db" - assert result.user_id == litellm.proxy.proxy_server.litellm_proxy_admin_name + assert result.user_id == DB_UNAVAILABLE_FALLBACK_USER_ID @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 19fffffc65..4f239b9acf 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -1692,6 +1692,7 @@ async def test_virtual_key_max_budget_alert_check_global_fallback(): ) import litellm + original = litellm.default_key_max_budget_alert_emails try: litellm.default_key_max_budget_alert_emails = global_config @@ -1731,6 +1732,7 @@ async def test_virtual_key_max_budget_alert_check_per_key_merges_with_global(): ) import litellm + original = litellm.default_key_max_budget_alert_emails try: litellm.default_key_max_budget_alert_emails = global_config @@ -1792,58 +1794,77 @@ async def test_get_fuzzy_user_object_case_insensitive_email(): @pytest.mark.asyncio async def test_custom_auth_common_checks_opt_in(): """ - Test that _run_post_custom_auth_checks only runs common_checks when + Test that common_checks only runs for a custom-auth deployment when custom_auth_run_common_checks is explicitly set to True in general_settings. - By default (False), common_checks is skipped for backwards compatibility - with custom auth flows that existed before PR #22164. + After the centralization refactor, common_checks runs in the + ``user_api_key_auth`` wrapper via ``_run_centralized_common_checks`` + (not inside ``_run_post_custom_auth_checks``). The opt-in flag now + gates the centralized gate for custom-auth deployments, preserving + the pre-existing RPS guarantee for custom-auth hot paths. """ - from litellm.proxy.auth.user_api_key_auth import _run_post_custom_auth_checks + import litellm.proxy.proxy_server as _proxy_server_mod + from litellm.proxy.auth.user_api_key_auth import _run_centralized_common_checks - valid_token = UserAPIKeyAuth(token="test-token") + valid_token = UserAPIKeyAuth(token="test-token", user_id="u1") mock_request = MagicMock() - # Default (no flag) — common_checks should NOT be called - with ( - patch( - "litellm.proxy.auth.user_api_key_auth.common_checks", - new_callable=AsyncMock, - ) as mock_common, - patch( - "litellm.proxy.proxy_server.general_settings", - {}, - ), - ): - mock_common.return_value = True - result = await _run_post_custom_auth_checks( - valid_token=valid_token, - request=mock_request, - request_data={}, - route="/ldap/ngs/ready", - parent_otel_span=None, - ) - mock_common.assert_not_called() + def _attrs(flag, user_custom_auth): + return { + "prisma_client": None, + "user_api_key_cache": MagicMock(), + "proxy_logging_obj": MagicMock(), + "general_settings": ( + {"custom_auth_run_common_checks": True} if flag else {} + ), + "llm_router": None, + "user_custom_auth": user_custom_auth, + "litellm_proxy_admin_name": "admin", + "master_key": "sk-test-master", + } - # With flag=True — common_checks SHOULD be called - with ( - patch( + # Default (no flag) with custom auth configured — centralized gate + # SHOULD skip to preserve custom-auth RPS. + attrs = _attrs(flag=False, user_custom_auth=AsyncMock()) + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + with patch( "litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock, - ) as mock_common, - patch( - "litellm.proxy.proxy_server.general_settings", - {"custom_auth_run_common_checks": True}, - ), - ): - mock_common.return_value = True - result = await _run_post_custom_auth_checks( - valid_token=valid_token, - request=mock_request, - request_data={}, - route="/chat/completions", - parent_otel_span=None, - ) - mock_common.assert_called_once() + ) as mock_common: + await _run_centralized_common_checks( + user_api_key_auth_obj=valid_token, + request=mock_request, + request_data={}, + route="/chat/completions", + ) + mock_common.assert_not_called() + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + + # With flag=True and custom auth configured — common_checks SHOULD run. + attrs = _attrs(flag=True, user_custom_auth=AsyncMock()) + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + with patch( + "litellm.proxy.auth.user_api_key_auth.common_checks", + new_callable=AsyncMock, + ) as mock_common: + await _run_centralized_common_checks( + user_api_key_auth_obj=valid_token, + request=mock_request, + request_data={}, + route="/chat/completions", + ) + mock_common.assert_called_once() + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) # ===================================================================== diff --git a/tests/test_litellm/proxy/auth/test_auth_exception_handler.py b/tests/test_litellm/proxy/auth/test_auth_exception_handler.py index 3e780c6ee9..54b8c4b53d 100644 --- a/tests/test_litellm/proxy/auth/test_auth_exception_handler.py +++ b/tests/test_litellm/proxy/auth/test_auth_exception_handler.py @@ -29,6 +29,35 @@ from litellm.proxy._types import ProxyErrorTypes, ProxyException from litellm.proxy.auth.auth_exception_handler import UserAPIKeyAuthExceptionHandler +@pytest.mark.asyncio +@pytest.mark.parametrize( + "prisma_error", + [ + HTTPClientClosedError(), + ClientNotConnectedError(), + ], +) +async def test_handle_authentication_error_db_unavailable_connectivity(prisma_error): + """Transport-level / connectivity failures trigger the HA fallback.""" + handler = UserAPIKeyAuthExceptionHandler() + + mock_request = MagicMock() + with patch( + "litellm.proxy.proxy_server.general_settings", + {"allow_requests_on_db_unavailable": True}, + ): + result = await handler._handle_authentication_error( + prisma_error, + mock_request, + {}, + "/test", + None, + "test-key", + ) + assert result.key_name == "failed-to-connect-to-db" + assert result.token == "failed-to-connect-to-db" + + @pytest.mark.asyncio @pytest.mark.parametrize( "prisma_error", @@ -51,35 +80,31 @@ from litellm.proxy.auth.auth_exception_handler import UserAPIKeyAuthExceptionHan RecordNotFoundError( data={"user_facing_error": {"meta": {"table": "test_table"}}} ), - HTTPClientClosedError(), - ClientNotConnectedError(), ], ) -async def test_handle_authentication_error_db_unavailable(prisma_error): +async def test_handle_authentication_error_data_layer_errors_do_not_fall_back( + prisma_error, +): + """Data-layer PrismaError subclasses (UniqueViolation, RecordNotFound, + etc.) mean the DB IS reachable — they must propagate instead of + triggering the HA fallback, which would grant the restricted + INTERNAL_USER token to a request that should have returned 401.""" handler = UserAPIKeyAuthExceptionHandler() - # Mock request and other dependencies mock_request = MagicMock() - mock_request_data = {} - mock_route = "/test" - mock_span = None - mock_api_key = "test-key" - - # Test with DB connection error when requests are allowed with patch( "litellm.proxy.proxy_server.general_settings", {"allow_requests_on_db_unavailable": True}, ): - result = await handler._handle_authentication_error( - prisma_error, - mock_request, - mock_request_data, - mock_route, - mock_span, - mock_api_key, - ) - assert result.key_name == "failed-to-connect-to-db" - assert result.token == "failed-to-connect-to-db" + with pytest.raises(ProxyException): + await handler._handle_authentication_error( + prisma_error, + mock_request, + {}, + "/test", + None, + "test-key", + ) @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py b/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py index 3af0cac6fd..4084fa4f3a 100644 --- a/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py +++ b/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py @@ -10,10 +10,13 @@ from litellm.proxy._types import UserAPIKeyAuth @pytest.mark.asyncio async def test_custom_auth_run_post_custom_auth_checks_without_end_user_id(): - # Test backwards compatibility — common_checks only runs when opt-in flag is set + # common_checks now runs in the user_api_key_auth wrapper via + # _run_centralized_common_checks, gated by + # custom_auth_run_common_checks for custom-auth deployments. The + # helper itself no longer calls common_checks. valid_token = UserAPIKeyAuth(token="test_token") - # Default: common_checks should NOT be called + # Default: common_checks should NOT be called inside the helper with patch( "litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock ) as mock_common: @@ -29,7 +32,8 @@ async def test_custom_auth_run_post_custom_auth_checks_without_end_user_id(): assert getattr(result, "end_user_id", None) is None mock_common.assert_not_awaited() - # With opt-in flag: common_checks SHOULD be called + # With opt-in flag: still not from the helper — the centralized gate + # in the wrapper handles it. with ( patch( "litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock @@ -48,7 +52,7 @@ async def test_custom_auth_run_post_custom_auth_checks_without_end_user_id(): parent_otel_span=None, ) assert result.token == "test_token" - mock_common.assert_awaited_once() + mock_common.assert_not_awaited() @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 881aa0c7e6..f06836e403 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -23,6 +23,7 @@ from litellm.proxy._types import ( from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.auth.user_api_key_auth import ( + _run_centralized_common_checks, _run_post_custom_auth_checks, get_api_key, user_api_key_auth, @@ -1732,3 +1733,267 @@ async def test_user_api_key_auth_builder_no_blocking_calls(): finally: for k, v in _originals.items(): setattr(_proxy_server_mod, k, v) + + +# --------------------------------------------------------------------------- +# _run_centralized_common_checks — centralized authz gate +# --------------------------------------------------------------------------- + + +def _proxy_attrs_for_centralized_checks( + user_custom_auth=None, flag=False, master_key="sk-test-master" +): + """Build the minimal proxy_server module attributes that + _run_centralized_common_checks reads. + + ``master_key`` defaults to a non-None value because setting it to + None short-circuits the gate (no-auth dev mode); tests that want to + exercise that branch must pass ``master_key=None`` explicitly. + """ + return { + "prisma_client": None, + "user_api_key_cache": MagicMock(), + "proxy_logging_obj": MagicMock(), + "general_settings": ({"custom_auth_run_common_checks": True} if flag else {}), + "llm_router": None, + "user_custom_auth": user_custom_auth, + "litellm_proxy_admin_name": "admin", + "master_key": master_key, + } + + +@pytest.mark.asyncio +async def test_centralized_common_checks_runs_for_standard_auth(): + """Regardless of which _user_api_key_auth_builder path returned, the + wrapper must run common_checks. This is the structural fix: no + early-return path can skip authorization.""" + import litellm.proxy.proxy_server as _proxy_server_mod + from fastapi import Request + from starlette.datastructures import URL + + token = UserAPIKeyAuth(api_key="sk-test", user_id="u1") + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + + attrs = _proxy_attrs_for_centralized_checks(user_custom_auth=None) + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + with patch( + "litellm.proxy.auth.user_api_key_auth.common_checks", + new_callable=AsyncMock, + ) as mock_checks: + await _run_centralized_common_checks( + user_api_key_auth_obj=token, + request=request, + request_data={"model": "gpt-4o"}, + route="/chat/completions", + ) + mock_checks.assert_awaited_once() + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + + +@pytest.mark.asyncio +async def test_centralized_common_checks_skipped_for_custom_auth_without_flag(): + """Existing RPS guarantee: custom-auth deployments without + custom_auth_run_common_checks must not pay the centralized gate. + Custom-auth paths don't use OAuth2/DB-fallback so this skip does + not widen any bypass.""" + import litellm.proxy.proxy_server as _proxy_server_mod + from fastapi import Request + from starlette.datastructures import URL + + token = UserAPIKeyAuth(api_key="sk-test", user_id="u1") + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + + attrs = _proxy_attrs_for_centralized_checks( + user_custom_auth=AsyncMock(), flag=False + ) + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + with patch( + "litellm.proxy.auth.user_api_key_auth.common_checks", + new_callable=AsyncMock, + ) as mock_checks: + await _run_centralized_common_checks( + user_api_key_auth_obj=token, + request=request, + request_data={"model": "gpt-4o"}, + route="/chat/completions", + ) + mock_checks.assert_not_awaited() + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + + +@pytest.mark.asyncio +async def test_centralized_common_checks_runs_for_custom_auth_with_flag(): + """Custom-auth deployments that opt in via custom_auth_run_common_checks + get the centralized gate.""" + import litellm.proxy.proxy_server as _proxy_server_mod + from fastapi import Request + from starlette.datastructures import URL + + token = UserAPIKeyAuth(api_key="sk-test", user_id="u1") + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + + attrs = _proxy_attrs_for_centralized_checks(user_custom_auth=AsyncMock(), flag=True) + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + with patch( + "litellm.proxy.auth.user_api_key_auth.common_checks", + new_callable=AsyncMock, + ) as mock_checks: + await _run_centralized_common_checks( + user_api_key_auth_obj=token, + request=request, + request_data={"model": "gpt-4o"}, + route="/chat/completions", + ) + mock_checks.assert_awaited_once() + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + + +@pytest.mark.asyncio +async def test_centralized_common_checks_runs_for_oauth2_fallback_token(): + """VERIA-18 regression: an OAuth2 token that would previously early- + return without common_checks is now subject to it. If common_checks + raises, the gate propagates the failure.""" + import litellm.proxy.proxy_server as _proxy_server_mod + from fastapi import Request + from starlette.datastructures import URL + + token = UserAPIKeyAuth(api_key="oauth2-token", user_id="oauth-user") + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + + attrs = _proxy_attrs_for_centralized_checks(user_custom_auth=None) + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + with patch( + "litellm.proxy.auth.user_api_key_auth.common_checks", + new_callable=AsyncMock, + side_effect=ProxyException( + message="Key not allowed to access model", + type=ProxyErrorTypes.key_model_access_denied, + param="model", + code=401, + ), + ): + with pytest.raises(ProxyException) as exc: + await _run_centralized_common_checks( + user_api_key_auth_obj=token, + request=request, + request_data={"model": "gpt-4"}, + route="/chat/completions", + ) + assert exc.value.type == ProxyErrorTypes.key_model_access_denied + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + + +@pytest.mark.asyncio +async def test_centralized_common_checks_tolerates_db_errors_when_fetching_context(): + """DB-outage scenario: the fallback UserAPIKeyAuth is issued when the + DB is down, then the gate tries to fetch team/user/etc. Those fetches + fail — the gate must swallow and still call common_checks with None + objects so enforcement runs against whatever the token recorded.""" + import litellm.proxy.proxy_server as _proxy_server_mod + from fastapi import Request + from starlette.datastructures import URL + + from litellm.proxy.auth.auth_exception_handler import ( + DB_UNAVAILABLE_FALLBACK_USER_ID, + ) + + token = UserAPIKeyAuth( + api_key="fallback", + user_id=DB_UNAVAILABLE_FALLBACK_USER_ID, + team_id="team-x", + ) + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + + attrs = _proxy_attrs_for_centralized_checks(user_custom_auth=None) + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + with ( + patch( + "litellm.proxy.auth.user_api_key_auth.get_team_object", + new_callable=AsyncMock, + side_effect=Exception("DB down"), + ), + patch( + "litellm.proxy.auth.user_api_key_auth.common_checks", + new_callable=AsyncMock, + ) as mock_checks, + ): + await _run_centralized_common_checks( + user_api_key_auth_obj=token, + request=request, + request_data={"model": "gpt-4o"}, + route="/chat/completions", + ) + mock_checks.assert_awaited_once() + # team_object kwarg should have ended up as None after the + # DB fetch was swallowed. + assert mock_checks.call_args.kwargs["team_object"] is None + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + + +@pytest.mark.asyncio +async def test_centralized_common_checks_short_circuits_when_master_key_unset(): + """master_key=None is no-auth dev mode — admin-only routes and + common_checks must not run. Deployments in this mode have no proxy- + level authentication, so applying authz would block every admin + route for a test/dev setup that was previously wide-open.""" + import litellm.proxy.proxy_server as _proxy_server_mod + from fastapi import Request + from starlette.datastructures import URL + + from litellm.proxy._types import LitellmUserRoles + + token = UserAPIKeyAuth( + api_key="sk-test", user_id="u", user_role=LitellmUserRoles.INTERNAL_USER + ) + request = Request(scope={"type": "http"}) + request._url = URL(url="/get/config/callbacks") + + attrs = _proxy_attrs_for_centralized_checks(user_custom_auth=None, master_key=None) + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + with patch( + "litellm.proxy.auth.user_api_key_auth.common_checks", + new_callable=AsyncMock, + ) as mock_checks: + await _run_centralized_common_checks( + user_api_key_auth_obj=token, + request=request, + request_data={}, + route="/get/config/callbacks", + ) + mock_checks.assert_not_awaited() + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v)