diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 1ef5b1072ed..7448ea7d322 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -5577,6 +5577,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 e810a5ab329..525e751e8ef 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -91,6 +91,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, @@ -2439,13 +2440,23 @@ class ProxyBaseLLMRequestProcessing: if isinstance(model, str): reject_url_valued_destination("model", model) - self.data["model"] = resolve_inference_model( - self.data.get("model"), + body_model: Final[object] = self.data.get("model") # pyright: ignore[reportUnknownVariableType] # self.data is an untyped dict body + resolved_model: Final = resolve_inference_model( + body_model, general_settings, user_model, model, kind="image_edit" if route_type == "aimage_edit" else "completion", ) + self.data["model"] = resolved_model + # Auth read the body model. When the endpoint model from the URL overrides it, + # the served model was never checked, so authorize the resolved name here, before + # any alias rewrite. A request whose model came from the body was already checked + # by auth, so it is left alone. + if llm_router is not None and isinstance(model, str) and model and model != body_model: + 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 28fb9baf2aa..0802cb329c5 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,18 @@ async def image_generation( if isinstance(model, str): reject_url_valued_destination("model", model) - data["model"] = resolve_inference_model( - data.get("model"), general_settings, user_model, model, kind="image_generation" + body_model: Final = data.get("model") # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # data is an untyped dict body + resolved_model: Final = resolve_inference_model( + body_model, general_settings, user_model, model, kind="image_generation" ) + data["model"] = resolved_model + # Auth read the body model. When the deployment or query model from the URL + # overrides it, the served model was never checked, so authorize the resolved + # name here, before any alias rewrite. A body model was already checked by auth. + if llm_router is not None and isinstance(model, str) and model and model != body_model: + 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 diff --git a/tests/unit/proxy/auth/test_auth_checks.py b/tests/unit/proxy/auth/test_auth_checks.py index e783f01769e..357bddc250e 100644 --- a/tests/unit/proxy/auth/test_auth_checks.py +++ b/tests/unit/proxy/auth/test_auth_checks.py @@ -2425,3 +2425,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