Merge 8f0eae1c89 into 8d441ebc15

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,

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:

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

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