mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-10 22:41:41 +00:00
fix(auth): centralize common_checks to close authorization bypass
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.
This commit is contained in:
parent
c95cac5d46
commit
3737d6a1f3
8 changed files with 680 additions and 191 deletions
|
|
@ -8,11 +8,22 @@ from fastapi import HTTPException, Request, status
|
||||||
|
|
||||||
import litellm
|
import litellm
|
||||||
from litellm._logging import verbose_proxy_logger
|
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.auth.auth_utils import _get_request_ip_address
|
||||||
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
|
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
|
||||||
from litellm.types.services import ServiceTypes
|
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:
|
if TYPE_CHECKING:
|
||||||
from opentelemetry.trace import Span as _Span
|
from opentelemetry.trace import Span as _Span
|
||||||
|
|
||||||
|
|
@ -47,7 +58,6 @@ class UserAPIKeyAuthExceptionHandler:
|
||||||
"""
|
"""
|
||||||
from litellm.proxy.proxy_server import (
|
from litellm.proxy.proxy_server import (
|
||||||
general_settings,
|
general_settings,
|
||||||
litellm_proxy_admin_name,
|
|
||||||
proxy_logging_obj,
|
proxy_logging_obj,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
@ -63,10 +73,17 @@ class UserAPIKeyAuthExceptionHandler:
|
||||||
duration=0.0,
|
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(
|
return UserAPIKeyAuth(
|
||||||
key_name="failed-to-connect-to-db",
|
key_name="failed-to-connect-to-db",
|
||||||
token="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,
|
request_route=route,
|
||||||
)
|
)
|
||||||
else:
|
else:
|
||||||
|
|
|
||||||
|
|
@ -29,6 +29,7 @@ from litellm.proxy.auth.auth_checks import (
|
||||||
_cache_key_object,
|
_cache_key_object,
|
||||||
_delete_cache_key_object,
|
_delete_cache_key_object,
|
||||||
_get_user_role,
|
_get_user_role,
|
||||||
|
_is_model_cost_zero,
|
||||||
_is_user_proxy_admin,
|
_is_user_proxy_admin,
|
||||||
_virtual_key_max_budget_alert_check,
|
_virtual_key_max_budget_alert_check,
|
||||||
_virtual_key_max_budget_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_id = result["user_id"]
|
||||||
user_object = result["user_object"]
|
user_object = result["user_object"]
|
||||||
end_user_id = result["end_user_id"]
|
end_user_id = result["end_user_id"]
|
||||||
end_user_object = result["end_user_object"]
|
|
||||||
org_id = result["org_id"]
|
org_id = result["org_id"]
|
||||||
token = result["token"]
|
|
||||||
team_membership: Optional[LiteLLM_TeamMembership] = result.get(
|
team_membership: Optional[LiteLLM_TeamMembership] = result.get(
|
||||||
"team_membership", None
|
"team_membership", None
|
||||||
)
|
)
|
||||||
jwt_claims = result.get("jwt_claims", 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:
|
if is_proxy_admin:
|
||||||
return UserAPIKeyAuth(
|
return UserAPIKeyAuth(
|
||||||
api_key=None,
|
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_metadata = _jwt_project_obj.metadata
|
||||||
valid_token.project_alias = _jwt_project_obj.project_alias
|
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)
|
return cast(UserAPIKeyAuth, valid_token)
|
||||||
|
|
||||||
#### ELSE ####
|
#### ELSE ####
|
||||||
|
|
@ -1498,22 +1471,6 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
|
||||||
user_info=call_info,
|
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
|
# Token passed all checks
|
||||||
if valid_token is None:
|
if valid_token is None:
|
||||||
raise HTTPException(401, detail="Invalid API key")
|
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()
|
@tracer.wrap()
|
||||||
async def user_api_key_auth(
|
async def user_api_key_auth(
|
||||||
request: Request,
|
request: Request,
|
||||||
|
|
@ -1608,6 +1804,28 @@ async def user_api_key_auth(
|
||||||
## ENSURE DISABLE ROUTE WORKS ACROSS ALL USER AUTH FLOWS ##
|
## ENSURE DISABLE ROUTE WORKS ACROSS ALL USER AUTH FLOWS ##
|
||||||
RouteChecks.should_call_route(route=route, valid_token=user_api_key_auth_obj)
|
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(
|
end_user_id = get_end_user_id_from_request_body(
|
||||||
request_data, _safe_get_request_headers(request)
|
request_data, _safe_get_request_headers(request)
|
||||||
)
|
)
|
||||||
|
|
@ -1929,55 +2147,10 @@ async def _run_post_custom_auth_checks(
|
||||||
model=current_model,
|
model=current_model,
|
||||||
)
|
)
|
||||||
|
|
||||||
# 5. Look up user object if user_id is set
|
# team / user / end_user / project context objects are fetched by
|
||||||
user_object = None
|
# the centralized common_checks gate in user_api_key_auth after
|
||||||
if valid_token.user_id is not None:
|
# this helper returns. Keep only the project fetch here because it
|
||||||
try:
|
# mutates the token (project_metadata / project_alias).
|
||||||
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
|
|
||||||
if valid_token.project_id is not None:
|
if valid_token.project_id is not None:
|
||||||
_project_obj = await get_project_object(
|
_project_obj = await get_project_object(
|
||||||
project_id=valid_token.project_id,
|
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_metadata = _project_obj.metadata
|
||||||
valid_token.project_alias = _project_obj.project_alias
|
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
|
return valid_token
|
||||||
|
|
|
||||||
|
|
@ -31,16 +31,12 @@ class PrismaDBExceptionHandler:
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def is_database_connection_error(e: Exception) -> bool:
|
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.
|
if PrismaDBExceptionHandler.is_database_transport_error(e):
|
||||||
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):
|
|
||||||
return True
|
return True
|
||||||
if isinstance(e, ProxyException) and e.type == ProxyErrorTypes.no_db_connection:
|
if isinstance(e, ProxyException) and e.type == ProxyErrorTypes.no_db_connection:
|
||||||
return True
|
return True
|
||||||
|
|
|
||||||
|
|
@ -3816,10 +3816,15 @@ async def test_user_api_key_auth_db_unavailable():
|
||||||
api_key="Bearer sk-123456789",
|
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 isinstance(result, UserAPIKeyAuth)
|
||||||
assert result.key_name == "failed-to-connect-to-db"
|
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
|
@pytest.mark.asyncio
|
||||||
|
|
|
||||||
|
|
@ -1692,6 +1692,7 @@ async def test_virtual_key_max_budget_alert_check_global_fallback():
|
||||||
)
|
)
|
||||||
|
|
||||||
import litellm
|
import litellm
|
||||||
|
|
||||||
original = litellm.default_key_max_budget_alert_emails
|
original = litellm.default_key_max_budget_alert_emails
|
||||||
try:
|
try:
|
||||||
litellm.default_key_max_budget_alert_emails = global_config
|
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
|
import litellm
|
||||||
|
|
||||||
original = litellm.default_key_max_budget_alert_emails
|
original = litellm.default_key_max_budget_alert_emails
|
||||||
try:
|
try:
|
||||||
litellm.default_key_max_budget_alert_emails = global_config
|
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
|
@pytest.mark.asyncio
|
||||||
async def test_custom_auth_common_checks_opt_in():
|
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.
|
custom_auth_run_common_checks is explicitly set to True in general_settings.
|
||||||
|
|
||||||
By default (False), common_checks is skipped for backwards compatibility
|
After the centralization refactor, common_checks runs in the
|
||||||
with custom auth flows that existed before PR #22164.
|
``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()
|
mock_request = MagicMock()
|
||||||
|
|
||||||
# Default (no flag) — common_checks should NOT be called
|
def _attrs(flag, user_custom_auth):
|
||||||
with (
|
return {
|
||||||
patch(
|
"prisma_client": None,
|
||||||
"litellm.proxy.auth.user_api_key_auth.common_checks",
|
"user_api_key_cache": MagicMock(),
|
||||||
new_callable=AsyncMock,
|
"proxy_logging_obj": MagicMock(),
|
||||||
) as mock_common,
|
"general_settings": (
|
||||||
patch(
|
{"custom_auth_run_common_checks": True} if flag else {}
|
||||||
"litellm.proxy.proxy_server.general_settings",
|
),
|
||||||
{},
|
"llm_router": None,
|
||||||
),
|
"user_custom_auth": user_custom_auth,
|
||||||
):
|
"litellm_proxy_admin_name": "admin",
|
||||||
mock_common.return_value = True
|
"master_key": "sk-test-master",
|
||||||
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()
|
|
||||||
|
|
||||||
# With flag=True — common_checks SHOULD be called
|
# Default (no flag) with custom auth configured — centralized gate
|
||||||
with (
|
# SHOULD skip to preserve custom-auth RPS.
|
||||||
patch(
|
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",
|
"litellm.proxy.auth.user_api_key_auth.common_checks",
|
||||||
new_callable=AsyncMock,
|
new_callable=AsyncMock,
|
||||||
) as mock_common,
|
) as mock_common:
|
||||||
patch(
|
await _run_centralized_common_checks(
|
||||||
"litellm.proxy.proxy_server.general_settings",
|
user_api_key_auth_obj=valid_token,
|
||||||
{"custom_auth_run_common_checks": True},
|
request=mock_request,
|
||||||
),
|
request_data={},
|
||||||
):
|
route="/chat/completions",
|
||||||
mock_common.return_value = True
|
)
|
||||||
result = await _run_post_custom_auth_checks(
|
mock_common.assert_not_called()
|
||||||
valid_token=valid_token,
|
finally:
|
||||||
request=mock_request,
|
for k, v in originals.items():
|
||||||
request_data={},
|
setattr(_proxy_server_mod, k, v)
|
||||||
route="/chat/completions",
|
|
||||||
parent_otel_span=None,
|
# With flag=True and custom auth configured — common_checks SHOULD run.
|
||||||
)
|
attrs = _attrs(flag=True, user_custom_auth=AsyncMock())
|
||||||
mock_common.assert_called_once()
|
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)
|
||||||
|
|
||||||
|
|
||||||
# =====================================================================
|
# =====================================================================
|
||||||
|
|
|
||||||
|
|
@ -29,6 +29,35 @@ from litellm.proxy._types import ProxyErrorTypes, ProxyException
|
||||||
from litellm.proxy.auth.auth_exception_handler import UserAPIKeyAuthExceptionHandler
|
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.asyncio
|
||||||
@pytest.mark.parametrize(
|
@pytest.mark.parametrize(
|
||||||
"prisma_error",
|
"prisma_error",
|
||||||
|
|
@ -51,35 +80,31 @@ from litellm.proxy.auth.auth_exception_handler import UserAPIKeyAuthExceptionHan
|
||||||
RecordNotFoundError(
|
RecordNotFoundError(
|
||||||
data={"user_facing_error": {"meta": {"table": "test_table"}}}
|
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()
|
handler = UserAPIKeyAuthExceptionHandler()
|
||||||
|
|
||||||
# Mock request and other dependencies
|
|
||||||
mock_request = MagicMock()
|
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(
|
with patch(
|
||||||
"litellm.proxy.proxy_server.general_settings",
|
"litellm.proxy.proxy_server.general_settings",
|
||||||
{"allow_requests_on_db_unavailable": True},
|
{"allow_requests_on_db_unavailable": True},
|
||||||
):
|
):
|
||||||
result = await handler._handle_authentication_error(
|
with pytest.raises(ProxyException):
|
||||||
prisma_error,
|
await handler._handle_authentication_error(
|
||||||
mock_request,
|
prisma_error,
|
||||||
mock_request_data,
|
mock_request,
|
||||||
mock_route,
|
{},
|
||||||
mock_span,
|
"/test",
|
||||||
mock_api_key,
|
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.asyncio
|
||||||
|
|
|
||||||
|
|
@ -10,10 +10,13 @@ from litellm.proxy._types import UserAPIKeyAuth
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_custom_auth_run_post_custom_auth_checks_without_end_user_id():
|
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")
|
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(
|
with patch(
|
||||||
"litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock
|
"litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock
|
||||||
) as mock_common:
|
) 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
|
assert getattr(result, "end_user_id", None) is None
|
||||||
mock_common.assert_not_awaited()
|
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 (
|
with (
|
||||||
patch(
|
patch(
|
||||||
"litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock
|
"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,
|
parent_otel_span=None,
|
||||||
)
|
)
|
||||||
assert result.token == "test_token"
|
assert result.token == "test_token"
|
||||||
mock_common.assert_awaited_once()
|
mock_common.assert_not_awaited()
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
|
|
|
||||||
|
|
@ -23,6 +23,7 @@ from litellm.proxy._types import (
|
||||||
from litellm.proxy.auth.handle_jwt import JWTHandler
|
from litellm.proxy.auth.handle_jwt import JWTHandler
|
||||||
from litellm.proxy.auth.route_checks import RouteChecks
|
from litellm.proxy.auth.route_checks import RouteChecks
|
||||||
from litellm.proxy.auth.user_api_key_auth import (
|
from litellm.proxy.auth.user_api_key_auth import (
|
||||||
|
_run_centralized_common_checks,
|
||||||
_run_post_custom_auth_checks,
|
_run_post_custom_auth_checks,
|
||||||
get_api_key,
|
get_api_key,
|
||||||
user_api_key_auth,
|
user_api_key_auth,
|
||||||
|
|
@ -1732,3 +1733,267 @@ async def test_user_api_key_auth_builder_no_blocking_calls():
|
||||||
finally:
|
finally:
|
||||||
for k, v in _originals.items():
|
for k, v in _originals.items():
|
||||||
setattr(_proxy_server_mod, k, v)
|
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)
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue