diff --git a/litellm/proxy/auth/key_disallowed_models.py b/litellm/proxy/auth/key_disallowed_models.py index 9eac408e5f8..33cbbe39ae1 100644 --- a/litellm/proxy/auth/key_disallowed_models.py +++ b/litellm/proxy/auth/key_disallowed_models.py @@ -1,4 +1,5 @@ from collections.abc import Collection, Iterable +from fnmatch import fnmatchcase from typing import Final from fastapi import status @@ -29,13 +30,38 @@ def _names_model_resolves_through( if llm_router is not None and model in llm_router.registry().aliases() else None ) + deployment: Final = llm_router.registry().deployment(model) if llm_router is not None else None + deployment_names: Final = ( + () if deployment is None else (deployment.model_name, deployment.model_info.team_public_model_name) + ) return frozenset( name - for name in (model, after_team_alias, after_key_alias, after_global_alias, dispatched, router_alias_target) + for name in ( + model, + after_team_alias, + after_key_alias, + after_global_alias, + dispatched, + router_alias_target, + *deployment_names, + ) if name is not None ) +def _matches(name: str, entry: str) -> bool: + return name == entry or ("*" in entry and fnmatchcase(name, entry)) + + +def _is_disallowed( + model: str, valid_token: UserAPIKeyAuth, llm_router: RegistryReader | None, disallowed: frozenset[str] +) -> bool: + names: Final = _names_model_resolves_through(model, valid_token, llm_router) + return any( + _matches(name, entry) for name in names for entry in disallowed + ) # comprehension-ok: the name x entry product is the match + + def first_disallowed_model( model: str | Iterable[str], valid_token: UserAPIKeyAuth, llm_router: RegistryReader | None ) -> str | None: @@ -44,7 +70,7 @@ def first_disallowed_model( return None requested: Final = (model,) if isinstance(model, str) else tuple(model) return next( - (m for m in requested if _names_model_resolves_through(m, valid_token, llm_router) & disallowed), + (m for m in requested if _is_disallowed(m, valid_token, llm_router, disallowed)), None, ) @@ -68,11 +94,7 @@ def without_key_disallowed_models( models: Iterable[str], valid_token: UserAPIKeyAuth, llm_router: RegistryReader | None ) -> tuple[str, ...]: disallowed: Final = key_disallowed_models(valid_token) - return tuple( - m - for m in models - if not disallowed or not _names_model_resolves_through(m, valid_token, llm_router) & disallowed - ) + return tuple(m for m in models if not disallowed or not _is_disallowed(m, valid_token, llm_router, disallowed)) def disallowed_models_removed_by(data: BaseModel | None, stored: Collection[str] | None) -> frozenset[str]: diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 36ffe74dc17..5b19e8a750f 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -3035,6 +3035,8 @@ async def prepare_key_update_data( data_json.pop("key", None) data_json.pop("new_key", None) data_json.pop("grace_period", None) # Request-only param, not a DB column + if "disallowed_models" in data_json and data_json["disallowed_models"] is None: + data_json["disallowed_models"] = [] if _resends_stored_routes(data, existing_key_row): data_json.pop("allowed_routes", None) if ( diff --git a/tests/integration/authorization/test_key_disallowed_models.py b/tests/integration/authorization/test_key_disallowed_models.py index 2000fbfe8e1..06e36f400d9 100644 --- a/tests/integration/authorization/test_key_disallowed_models.py +++ b/tests/integration/authorization/test_key_disallowed_models.py @@ -138,6 +138,11 @@ def test_adding_and_clearing_disallowed_models_takes_effect_on_the_next_request( assert info["disallowed_models"] == [grant.disallowed], info gateway.post("/key/update", {"key": key, "disallowed_models": []}) _assert_served(gateway, key, grant.disallowed) + gateway.post("/key/update", {"key": key, "disallowed_models": [grant.disallowed]}) + _assert_denied(gateway, key, {"model": grant.disallowed}, grant.disallowed) + gateway.post("/key/update", {"key": key, "disallowed_models": None}) + assert _stored_disallowed_models(key) == [{"disallowed_models": []}] + _assert_served(gateway, key, grant.disallowed) def test_a_key_owner_cannot_remove_disallowed_models_but_can_add_them(gateway: Gateway) -> None: @@ -177,3 +182,12 @@ def test_an_admin_of_the_keys_team_can_remove_disallowed_models(gateway: Gateway assert cleared.status_code == 200, f"{cleared.status_code} {cleared.text}" assert _stored_disallowed_models(key) == [{"disallowed_models": []}] _assert_served(gateway, key, disallowed) + + +def test_a_deployment_id_cannot_reach_a_disallowed_model(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + deployment_id: Final = f"integration-deployment-{uuid.uuid4().hex}" + disallowed: Final = scenario.model(model_info={"id": deployment_id}) + key: Final = scenario.key(disallowed_models=[disallowed]) + _assert_denied(gateway, key, {"model": deployment_id}, deployment_id) + _assert_denied(gateway, key, {"model": scenario.model(), "fallbacks": [deployment_id]}, deployment_id) diff --git a/tests/unit/proxy/auth/test_key_disallowed_models.py b/tests/unit/proxy/auth/test_key_disallowed_models.py index b4bb90981e3..143d689c930 100644 --- a/tests/unit/proxy/auth/test_key_disallowed_models.py +++ b/tests/unit/proxy/auth/test_key_disallowed_models.py @@ -52,6 +52,44 @@ def test_a_key_alias_that_a_global_alias_redirects_to_a_disallowed_model_is_refu _denied_model(error) +def test_a_deployment_id_of_a_disallowed_model_group_is_refused_and_unlisted() -> None: + router: Final = litellm.Router( + model_list=[ + { + "model_name": "gpt-denied", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-fixture"}, + "model_info": {"id": "deployment-of-denied"}, + }, + { + "model_name": "gpt-allowed", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-fixture"}, + "model_info": {"id": "deployment-of-allowed"}, + }, + ] + ) + token: Final = _token(disallowed_models=("gpt-denied",)) + with pytest.raises(ProxyException) as error: + ensure_key_may_call_model("deployment-of-denied", token, router) + _denied_model(error) + assert first_disallowed_model(["deployment-of-allowed"], token, router) is None + assert without_key_disallowed_models(["deployment-of-denied", "deployment-of-allowed"], token, router) == ( + "deployment-of-allowed", + ) + + +@pytest.mark.parametrize( + "requested,refused", + [ + pytest.param("openai/gpt-4o", True, id="matches-provider-wildcard"), + pytest.param("anthropic/claude", False, id="other-provider"), + pytest.param("openai", False, id="bare-provider-name"), + ], +) +def test_a_wildcard_entry_refuses_every_model_it_matches(requested: str, refused: bool) -> None: + token: Final = _token(disallowed_models=("openai/*",)) + assert (first_disallowed_model([requested], token, None) == requested) is refused + + @pytest.mark.parametrize("disallowed", [None, (), ("gpt-other",)]) def test_a_model_outside_disallowed_models_is_not_refused(disallowed: tuple[str, ...] | None) -> None: assert first_disallowed_model(["gpt-allowed"], _token(disallowed_models=disallowed), None) is None