mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
fix(auth): check the key allowlist on team alias targets
A team model alias used to satisfy the model access check on its own, so a key scoped to a subset of the team's models could call any model the team had aliased. The alias is now resolved to its target and the target goes through the same allowlist, wildcard and access group checks as a plain model name, for both key and team checks and for the header forwarding lookup that shares the helper Co-authored-by: songkuan-zheng <252822057+songkuan-zheng@users.noreply.github.com> Co-authored-by: songkuan-zheng <songkuan-zheng@users.noreply.github.com>
This commit is contained in:
parent
9a715df212
commit
a0756b71cf
2 changed files with 123 additions and 19 deletions
|
|
@ -3959,7 +3959,13 @@ def _check_model_access_helper(
|
|||
# Filter out models that are access_groups
|
||||
filtered_models: Final = [m for m in models if m not in access_groups]
|
||||
|
||||
if _model_in_team_aliases(model=model, team_model_aliases=team_model_aliases):
|
||||
team_alias_target: Final = team_model_aliases.get(model) if team_model_aliases else None
|
||||
if team_alias_target is not None and _check_model_access_helper(
|
||||
model=team_alias_target,
|
||||
llm_router=llm_router,
|
||||
models=models,
|
||||
team_id=team_id,
|
||||
):
|
||||
return True
|
||||
|
||||
if _model_matches_any_wildcard_pattern_in_list(model=model, allowed_model_list=filtered_models):
|
||||
|
|
@ -4045,24 +4051,6 @@ def _can_object_call_model(
|
|||
)
|
||||
|
||||
|
||||
def _model_in_team_aliases(model: str, team_model_aliases: dict[str, str] | None = None) -> bool:
|
||||
"""
|
||||
Returns True if `model` being accessed is an alias of a team model
|
||||
|
||||
- `model=gpt-4o`
|
||||
- `team_model_aliases={"gpt-4o": "gpt-4o-team-1"}`
|
||||
- returns True
|
||||
|
||||
- `model=gp-4o`
|
||||
- `team_model_aliases={"o-3": "o3-preview"}`
|
||||
- returns False
|
||||
"""
|
||||
if team_model_aliases:
|
||||
if model in team_model_aliases:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _resolve_key_models_for_auth_check(valid_token: UserAPIKeyAuth) -> list[str]:
|
||||
"""
|
||||
Expand key model sentinels before auth checks.
|
||||
|
|
|
|||
|
|
@ -403,6 +403,122 @@ async def test_can_key_call_model_all_team_models_no_team_id_is_unrestricted():
|
|||
)
|
||||
|
||||
|
||||
_TEAM_ALIASES: Final = {
|
||||
"fast": "openai/gpt-4.1-mini",
|
||||
"smart": "openai/gpt-4.1",
|
||||
"vision": "openai/gpt-4o",
|
||||
}
|
||||
|
||||
|
||||
def _aliased_key(models: list[str]) -> UserAPIKeyAuth:
|
||||
return UserAPIKeyAuth(
|
||||
api_key="sk-key",
|
||||
team_id="team-1",
|
||||
models=models,
|
||||
team_models=["openai/gpt-4.1-mini", "openai/gpt-4.1"],
|
||||
team_model_aliases=_TEAM_ALIASES,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_can_key_call_model_team_alias_target_outside_key_models_denied():
|
||||
"""A team alias whose target is not in key.models is denied."""
|
||||
from litellm.proxy.auth.auth_checks import can_key_call_model
|
||||
|
||||
valid_token: Final = _aliased_key(models=["openai/gpt-4.1-mini"])
|
||||
|
||||
for alias in ("smart", "vision"):
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await can_key_call_model(model=alias, llm_model_list=None, valid_token=valid_token, llm_router=None)
|
||||
assert exc_info.value.type == ProxyErrorTypes.key_model_access_denied
|
||||
assert exc_info.value.code == str(status.HTTP_403_FORBIDDEN)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_can_key_call_model_team_alias_target_in_key_models_allowed():
|
||||
"""A team alias whose target is in key.models is allowed, and a plain model name still resolves as before."""
|
||||
from litellm.proxy.auth.auth_checks import can_key_call_model
|
||||
|
||||
valid_token: Final = _aliased_key(models=["openai/gpt-4.1-mini"])
|
||||
|
||||
for model in ("fast", "openai/gpt-4.1-mini"):
|
||||
assert await can_key_call_model(model=model, llm_model_list=None, valid_token=valid_token, llm_router=None)
|
||||
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await can_key_call_model(model="openai/gpt-4.1", llm_model_list=None, valid_token=valid_token, llm_router=None)
|
||||
assert exc_info.value.type == ProxyErrorTypes.key_model_access_denied
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_can_key_call_model_team_alias_unrestricted_key_allowed():
|
||||
"""A key with no model restriction, or the all-proxy-models sentinel, can call any team alias."""
|
||||
from litellm.proxy._types import SpecialModelNames
|
||||
from litellm.proxy.auth.auth_checks import can_key_call_model
|
||||
|
||||
for models in ([], [SpecialModelNames.all_proxy_models.value]):
|
||||
valid_token: Final = _aliased_key(models=models)
|
||||
for alias in _TEAM_ALIASES:
|
||||
assert await can_key_call_model(model=alias, llm_model_list=None, valid_token=valid_token, llm_router=None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_can_key_call_model_team_alias_all_team_models_sentinel():
|
||||
"""With the all-team-models sentinel the alias target must be in team_models."""
|
||||
from litellm.proxy._types import SpecialModelNames
|
||||
from litellm.proxy.auth.auth_checks import can_key_call_model
|
||||
|
||||
valid_token: Final = _aliased_key(models=[SpecialModelNames.all_team_models.value])
|
||||
|
||||
for alias in ("fast", "smart"):
|
||||
assert await can_key_call_model(model=alias, llm_model_list=None, valid_token=valid_token, llm_router=None)
|
||||
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await can_key_call_model(model="vision", llm_model_list=None, valid_token=valid_token, llm_router=None)
|
||||
assert exc_info.value.type == ProxyErrorTypes.key_model_access_denied
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_can_team_access_model_team_alias_checks_target():
|
||||
"""A team alias is allowed only when its target is in team.models."""
|
||||
from litellm.proxy.auth.auth_checks import can_team_access_model
|
||||
|
||||
team_object: Final = LiteLLM_TeamTable(team_id="team-1", models=["openai/gpt-4.1-mini"])
|
||||
|
||||
assert await can_team_access_model(
|
||||
model="fast", team_object=team_object, llm_router=None, team_model_aliases=_TEAM_ALIASES
|
||||
)
|
||||
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await can_team_access_model(
|
||||
model="smart", team_object=team_object, llm_router=None, team_model_aliases=_TEAM_ALIASES
|
||||
)
|
||||
assert exc_info.value.type == ProxyErrorTypes.team_model_access_denied
|
||||
|
||||
|
||||
def test_check_model_access_helper_team_alias_target_matches_wildcard():
|
||||
"""A team alias target is matched against wildcard entries in the allowlist."""
|
||||
from litellm.proxy.auth.auth_checks import _check_model_access_helper
|
||||
|
||||
assert _check_model_access_helper(
|
||||
model="fast", llm_router=None, models=["openai/*"], team_model_aliases=_TEAM_ALIASES, team_id="team-1"
|
||||
)
|
||||
assert not _check_model_access_helper(
|
||||
model="fast", llm_router=None, models=["anthropic/*"], team_model_aliases=_TEAM_ALIASES, team_id="team-1"
|
||||
)
|
||||
|
||||
|
||||
def test_check_model_access_helper_self_referential_alias_terminates():
|
||||
"""An alias that points at its own name is decided by the allowlist alone."""
|
||||
from litellm.proxy.auth.auth_checks import _check_model_access_helper
|
||||
|
||||
self_alias: Final = {"same": "same"}
|
||||
|
||||
assert _check_model_access_helper(model="same", llm_router=None, models=["same"], team_model_aliases=self_alias)
|
||||
assert not _check_model_access_helper(
|
||||
model="same", llm_router=None, models=["other"], team_model_aliases=self_alias
|
||||
)
|
||||
|
||||
|
||||
def test_resolve_key_models_teamless_all_team_models_returns_empty():
|
||||
"""_resolve_key_models_for_auth_check must return [] for a teamless key
|
||||
with all-team-models, making it equivalent to an unscoped key (unrestricted
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue