mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
fix(proxy): authorize the resolved model before alias rewrites, not after
Check the model the handler resolved at the point it is resolved, in common_processing_pre_call_logic and the image generation handler, before the global, key and team alias rewrites. The earlier placement in route_request saw the post-rewrite name, so a key allowed a global alias was denied once the alias was rewritten to its target. Checking the requested name keeps every alias grant that auth already honored when it read the body model. The new can_key_call_requested_models splits a comma batch and a list the way auth splits the request model and authorizes each name, so a key allowed every model of a batch still passes. Add allowlist tests covering single, comma and list models in and out of the allowlist, and an alias-only grant. Tracking issue: https://github.com/BerriAI/litellm/issues/48760 Signed-off-by: Atharva-Kanherkar <atharvakanherkar25@gmail.com>
This commit is contained in:
parent
75f5cdf134
commit
ffcc8c792f
6 changed files with 108 additions and 109 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue