From ffcc8c792f7ba075019aa78aa8e19e0f8edd969e Mon Sep 17 00:00:00 2001 From: Atharva-Kanherkar Date: Wed, 28 Oct 2026 19:42:36 +0900 Subject: [PATCH] fix(proxy): authorize the resolved model before alias rewrites, not after Check the model the handler resolved at the point it is resolved, in common_processing_pre_call_logic and the image generation handler, before the global, key and team alias rewrites. The earlier placement in route_request saw the post-rewrite name, so a key allowed a global alias was denied once the alias was rewritten to its target. Checking the requested name keeps every alias grant that auth already honored when it read the body model. The new can_key_call_requested_models splits a comma batch and a list the way auth splits the request model and authorizes each name, so a key allowed every model of a batch still passes. Add allowlist tests covering single, comma and list models in and out of the allowlist, and an alias-only grant. Tracking issue: https://github.com/BerriAI/litellm/issues/48760 Signed-off-by: Atharva-Kanherkar --- litellm/proxy/auth/auth_checks.py | 36 +++++++++++ litellm/proxy/common_request_processing.py | 10 ++- litellm/proxy/image_endpoints/endpoints.py | 11 +++- litellm/proxy/route_llm_request.py | 34 +--------- tests/unit/proxy/auth/test_auth_checks.py | 53 ++++++++++++++++ tests/unit/proxy/test_route_llm_request.py | 73 ---------------------- 6 files changed, 108 insertions(+), 109 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 448643ef59b..e7f3bc1c37e 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -5561,6 +5561,42 @@ async def _personal_key_owner( return None +def _requested_model_names(model: object) -> Iterator[str]: + """Each model name a request can route to, split the way auth reads the request model. + + A comma-separated value is a batch the router fans out and a list is a set of + names, so every non-empty part is a name of its own. + """ + items: Final[tuple[object, ...]] = tuple(model) if isinstance(model, list) else (model,) # pyright: ignore[reportUnknownArgumentType, reportUnknownVariableType] # a list body model holds unknown element types + for item in items: + if not isinstance(item, str): + continue + for part in item.split(","): + name = part.strip() + if name: + yield name + + +async def can_key_call_requested_models( + model: object, + valid_token: UserAPIKeyAuth, + llm_router: litellm.Router, +) -> None: + """Authorize every model a request resolves to, before any alias rewrite. + + Call this where an endpoint model from the URL path or query is resolved into the + request. The names checked are the ones the caller asked for, so global, key and + team alias grants apply exactly as they do in auth, which reads the body model. + """ + for name in _requested_model_names(model): + await can_key_call_resolved_model( + model=name, + llm_model_list=llm_router.model_list, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # Router.model_list is an untyped list + valid_token=valid_token, + llm_router=llm_router, + ) + + async def can_key_call_resolved_model( model: str, llm_model_list: Sequence[object] | None, diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 24a36a7f2dc..9e58c662726 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -90,6 +90,7 @@ from litellm.litellm_core_utils.streaming_handler import ( from litellm.litellm_core_utils.tool_fields import DropPolicy, MalformedField, drop_policy, resolve_tool_fields from litellm.proxy._types import LiteLLMRoutes, ProxyErrorTypes, ProxyException, UserAPIKeyAuth from litellm.proxy.auth.auth_checks import ( + can_key_call_requested_models, can_key_call_resolved_model, is_model_cost_zero, request_skips_budget_checks, @@ -2364,13 +2365,20 @@ class ProxyBaseLLMRequestProcessing: if isinstance(model, str): reject_url_valued_destination("model", model) - self.data["model"] = resolve_inference_model( + resolved_model: Final = resolve_inference_model( self.data.get("model"), general_settings, user_model, model, kind="image_edit" if route_type == "aimage_edit" else "completion", ) + self.data["model"] = resolved_model + # Auth checked the body model; the endpoint model from the URL wins here, so + # authorize the resolved name before any alias rewrite. + if llm_router is not None: + await can_key_call_requested_models( + model=resolved_model, valid_token=user_api_key_dict, llm_router=llm_router + ) # override with user settings, these are params passed via cli if user_temperature: diff --git a/litellm/proxy/image_endpoints/endpoints.py b/litellm/proxy/image_endpoints/endpoints.py index e93b22cb716..3112e41c277 100644 --- a/litellm/proxy/image_endpoints/endpoints.py +++ b/litellm/proxy/image_endpoints/endpoints.py @@ -14,6 +14,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_str_from_messages, ) from litellm.proxy._types import * +from litellm.proxy.auth.auth_checks import can_key_call_requested_models from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth from litellm.proxy.common_request_processing import ( ProxyBaseLLMRequestProcessing, @@ -122,9 +123,16 @@ async def image_generation( if isinstance(model, str): reject_url_valued_destination("model", model) - data["model"] = resolve_inference_model( + resolved_model: Final = resolve_inference_model( data.get("model"), general_settings, user_model, model, kind="image_generation" ) + data["model"] = resolved_model + # Auth checked the body model; the deployment path model wins here, so + # authorize the resolved name before any alias rewrite. + if llm_router is not None: + await can_key_call_requested_models( + model=resolved_model, valid_token=user_api_key_dict, llm_router=llm_router + ) ### MODEL ALIAS MAPPING ### # check if model name in model alias map @@ -156,7 +164,6 @@ async def image_generation( route_type="aimage_generation", llm_router=llm_router, user_model=user_model, - user_api_key_dict=user_api_key_dict, ) response = await llm_call diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index 608985910c5..cb226929167 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -2,7 +2,7 @@ import asyncio import hashlib import json from collections import OrderedDict -from collections.abc import Callable, Iterator, Mapping +from collections.abc import Callable, Mapping from dataclasses import dataclass from datetime import datetime from types import MappingProxyType @@ -710,22 +710,6 @@ def _merge_router_settings_override( return _ROUTER_ONLY_SETTINGS.intersection(router_settings_from_override) -def _models_to_authorize(model: object) -> Iterator[str]: - """Each model the dispatch can reach, split the way auth splits the request model. - - A comma-separated value is a batch the router fans out, and a list is a set of - names, so every non-empty part is authorized on its own. - """ - items: Final[tuple[object, ...]] = tuple(model) if isinstance(model, list) else (model,) # pyright: ignore[reportUnknownArgumentType, reportUnknownVariableType] # a list body model holds unknown element types - for item in items: - if not isinstance(item, str): - continue - for part in item.split(","): - name = part.strip() - if name: - yield name - - async def route_request( data: dict, llm_router: LitellmRouter | None, @@ -736,22 +720,6 @@ async def route_request( """ Common helper to route the request """ - # Enforce the key and team allowlist on the model that will actually be routed, - # not only the body model auth read. Routes that bind the model from the URL - # resolve it into data["model"] before dispatch, so this is the single place - # that covers every route, including ones added later. - dispatched_model: Final = data.get("model") # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # route_request takes an untyped dict body - if user_api_key_dict is not None and llm_router is not None: - from litellm.proxy.auth.auth_checks import can_key_call_resolved_model - - for model_name in _models_to_authorize(dispatched_model): # pyright: ignore[reportUnknownArgumentType] # dispatched_model comes from the untyped dict body - await can_key_call_resolved_model( - model=model_name, - llm_model_list=llm_router.model_list, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # Router.model_list is an untyped list - valid_token=user_api_key_dict, - llm_router=llm_router, - ) - router_only_settings_from_override: Final = _merge_router_settings_override(data) # pyright: ignore[reportUnknownArgumentType] # route_request takes an untyped dict body try: return await _route_request_single_attempt( diff --git a/tests/unit/proxy/auth/test_auth_checks.py b/tests/unit/proxy/auth/test_auth_checks.py index e1afa81d372..bef9c892d5e 100644 --- a/tests/unit/proxy/auth/test_auth_checks.py +++ b/tests/unit/proxy/auth/test_auth_checks.py @@ -2380,3 +2380,56 @@ async def test_bulk_eviction_clears_the_redis_copies_before_waiting_on_a_wedged_ assert refilled == [] assert redis.entries == {} + + +def _requested_models_router() -> litellm.Router: + return litellm.Router( + model_list=[ + { + "model_name": name, + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "fake", "mock_response": f"from {name}"}, + } + for name in ("model-a", "model-b", "model-c") + ] + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "requested", + ["model-c", "model-a,model-c", ["model-a", "model-c"]], + ids=["single", "comma_batch", "list"], +) +async def test_can_key_call_requested_models_denies_any_model_outside_allowlist(requested): + from litellm.proxy._types import ProxyErrorTypes, ProxyException + from litellm.proxy.auth.auth_checks import can_key_call_requested_models + + key = UserAPIKeyAuth(api_key="sk-test", models=["model-a", "model-b"]) + with pytest.raises(ProxyException) as exc: + await can_key_call_requested_models(requested, key, _requested_models_router()) + assert exc.value.type == ProxyErrorTypes.key_model_access_denied.value + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "requested", + ["model-a", "model-a,model-b", ["model-a", "model-b"]], + ids=["single", "comma_batch", "list"], +) +async def test_can_key_call_requested_models_allows_models_inside_allowlist(requested): + from litellm.proxy.auth.auth_checks import can_key_call_requested_models + + key = UserAPIKeyAuth(api_key="sk-test", models=["model-a", "model-b"]) + # Returns None without raising for any form that stays inside the allowlist. + assert await can_key_call_requested_models(requested, key, _requested_models_router()) is None + + +@pytest.mark.asyncio +async def test_can_key_call_requested_models_preserves_a_global_alias_grant(monkeypatch): + from litellm.proxy.auth.auth_checks import can_key_call_requested_models + + monkeypatch.setattr(litellm, "model_alias_map", {"friendly": "model-a"}) + key = UserAPIKeyAuth(api_key="sk-test", models=["friendly"]) + # A key allowed the alias may call the alias name; resolution to model-a happens + # after this check, so the grant must not be denied here. + assert await can_key_call_requested_models("friendly", key, _requested_models_router()) is None diff --git a/tests/unit/proxy/test_route_llm_request.py b/tests/unit/proxy/test_route_llm_request.py index ad411995600..6044475652c 100644 --- a/tests/unit/proxy/test_route_llm_request.py +++ b/tests/unit/proxy/test_route_llm_request.py @@ -3025,76 +3025,3 @@ async def test_user_config_routers_discard_the_least_recently_used_router_past_c assert first_handler() is not None assert second_handler() is None routers.clear() - - -def _allowlist_router() -> litellm.Router: - return litellm.Router( - model_list=[ - { - "model_name": name, - "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "fake", "mock_response": f"from {name}"}, - } - for name in ("model-a", "model-b", "model-c") - ] - ) - - -def _hi() -> list[dict[str, str]]: - return [{"role": "user", "content": "hi"}] - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "requested", - ["model-c", "model-a,model-c", ["model-a", "model-c"]], - ids=["single", "comma_batch", "list"], -) -async def test_route_request_denies_any_routed_model_outside_key_allowlist(requested): - from litellm.proxy._types import ProxyErrorTypes, ProxyException - - key = UserAPIKeyAuth(api_key="sk-test", models=["model-a", "model-b"]) - - with pytest.raises(ProxyException) as exc: - await route_request( - {"model": requested, "messages": _hi()}, _allowlist_router(), None, "acompletion", user_api_key_dict=key - ) - assert exc.value.type == ProxyErrorTypes.key_model_access_denied.value - - -@pytest.mark.asyncio -async def test_route_request_serves_a_model_inside_key_allowlist(): - key = UserAPIKeyAuth(api_key="sk-test", models=["model-a", "model-b"]) - - llm_call = await route_request( - {"model": "model-a", "messages": _hi()}, _allowlist_router(), None, "acompletion", user_api_key_dict=key - ) - response = await llm_call - - assert response.choices[0].message.content == "from model-a" - - -@pytest.mark.asyncio -async def test_route_request_serves_a_comma_batch_when_key_allows_every_model(): - key = UserAPIKeyAuth(api_key="sk-test", models=["model-a", "model-b"]) - - llm_call = await route_request( - {"model": "model-a, model-b", "messages": _hi()}, - _allowlist_router(), - None, - "acompletion", - user_api_key_dict=key, - ) - responses = await llm_call - - assert sorted(response.choices[0].message.content for response in responses) == [ - "from model-a", - "from model-b", - ] - - -@pytest.mark.asyncio -async def test_route_request_without_a_key_leaves_model_access_to_auth(): - llm_call = await route_request({"model": "model-c", "messages": _hi()}, _allowlist_router(), None, "acompletion") - response = await llm_call - - assert response.choices[0].message.content == "from model-c"