diff --git a/litellm/proxy/image_endpoints/endpoints.py b/litellm/proxy/image_endpoints/endpoints.py index 28fb9baf2aa..e93b22cb716 100644 --- a/litellm/proxy/image_endpoints/endpoints.py +++ b/litellm/proxy/image_endpoints/endpoints.py @@ -156,6 +156,7 @@ 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 cb226929167..ddb6f6bb184 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -720,6 +720,21 @@ 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 and isinstance(dispatched_model, str): + from litellm.proxy.auth.auth_checks import can_key_call_resolved_model + + await can_key_call_resolved_model( + model=dispatched_model, + 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/integration/authorization/test_url_route_model_access.py b/tests/integration/authorization/test_url_route_model_access.py deleted file mode 100644 index d6917fa61d5..00000000000 --- a/tests/integration/authorization/test_url_route_model_access.py +++ /dev/null @@ -1,83 +0,0 @@ -"""Virtual-key model allowlist enforcement on routes that bind the model from the URL. - -Tracking issue: https://github.com/BerriAI/litellm/issues/48760 - -A virtual key scoped with ``models: ["allowed-model"]`` is correctly denied a model -named in the request body, but the Google ``generateContent`` and -``streamGenerateContent`` routes and the ``/openai/deployments/{model}/images/...`` -routes bind the model from the URL and route on it, while the access check reads the -body model. A body model the key may use plus a denied model in the URL is served the -denied model. The open PR #41852 closes the chat, completions and embeddings routes, -not these. - -These skeleton cases name the routes the fix must cover. They are skipped until the -fix lands; the implementation re-checks the resolved ``self.data["model"]`` at the -single request-processing choke point so every route is covered at once. See the -tracking issue for the full reproduction, the N-of-N matrix and the fix-control run. -""" - -from __future__ import annotations - -from dataclasses import dataclass -from typing import Final - -import pytest - -pytestmark = pytest.mark.skip( - reason="WIP, enforcing the model allowlist on URL-bound model routes, see issue #48760" -) - - -@dataclass(frozen=True, slots=True) -class _UrlRouteCase: - name: str - method: str - path: str - body_model: str - served_denied: bool - - -# A key scoped to "allowed-model". Each case sends "denied-model" in the URL and -# "allowed-model" in the body. served_denied is the behaviour on current main that -# the fix must turn into a 403. -_CASES: Final = ( - _UrlRouteCase( - "generate_content", - "POST", - "/v1beta/models/denied-model:generateContent", - "allowed-model", - True, - ), - _UrlRouteCase( - "stream_generate_content", - "POST", - "/v1beta/models/denied-model:streamGenerateContent", - "allowed-model", - True, - ), - _UrlRouteCase( - "images_generations_deployment", - "POST", - "/openai/deployments/denied-model/images/generations", - "allowed-model", - True, - ), - _UrlRouteCase( - "images_edits_deployment", - "POST", - "/openai/deployments/denied-model/images/edits", - "allowed-model", - True, - ), -) - - -@pytest.mark.parametrize("case", _CASES, ids=lambda case: case.name) -def test_url_route_model_outside_allowlist_is_denied(case: _UrlRouteCase) -> None: - """A URL-bound model outside the key allowlist must be denied before dispatch. - - The served model, not just the body model, must be inside the key's allowed set. - The request below sends an allowed body model and a denied URL model, and must - return 403 key_model_access_denied with no upstream call. - """ - raise NotImplementedError("fix and assertions pending, see issue #48760") diff --git a/tests/unit/proxy/test_route_llm_request.py b/tests/unit/proxy/test_route_llm_request.py index 6044475652c..f560e6543b6 100644 --- a/tests/unit/proxy/test_route_llm_request.py +++ b/tests/unit/proxy/test_route_llm_request.py @@ -3025,3 +3025,66 @@ 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": "allowed-model", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "fake", "mock_response": "ok"}, + }, + { + "model_name": "denied-model", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "fake", "mock_response": "ok"}, + }, + ] + ) + + +@pytest.mark.asyncio +async def test_route_request_denies_resolved_model_outside_key_allowlist(): + from litellm.proxy._types import ProxyErrorTypes, ProxyException, UserAPIKeyAuth + + router = _allowlist_router() + key = UserAPIKeyAuth(api_key="sk-test", models=["allowed-model"]) + data = {"model": "denied-model", "messages": [{"role": "user", "content": "hi"}]} + + with pytest.raises(ProxyException) as exc: + await route_request(data, 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_allows_resolved_model_inside_key_allowlist(): + from litellm.proxy._types import UserAPIKeyAuth + + router = _allowlist_router() + key = UserAPIKeyAuth(api_key="sk-test", models=["allowed-model"]) + data = {"model": "allowed-model", "messages": [{"role": "user", "content": "hi"}]} + + response = await route_request(data, router, None, "acompletion", user_api_key_dict=key) + assert await response if asyncio.iscoroutine(response) else response + + +@pytest.mark.asyncio +async def test_route_request_without_a_key_skips_the_allowlist_check(): + router = _allowlist_router() + data = {"model": "denied-model", "messages": [{"role": "user", "content": "hi"}]} + + response = await route_request(data, router, None, "acompletion") + assert await response if asyncio.iscoroutine(response) else response + + +@pytest.mark.asyncio +async def test_route_request_list_model_is_not_blocked_by_the_allowlist_check(): + from litellm.proxy._types import UserAPIKeyAuth + + router = _allowlist_router() + key = UserAPIKeyAuth(api_key="sk-test", models=["allowed-model"]) + data = {"model": ["allowed-model", "denied-model"], "messages": [{"role": "user", "content": "hi"}]} + + # A list model value is handled by the per-model auth path, not this backstop, + # so the check must not raise on it. + with contextlib.suppress(Exception): + await route_request(data, router, None, "acompletion", user_api_key_dict=key)