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"