mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-03 02:22:24 +00:00
fix(proxy): scope team BYOK rows on /v1/model/info to caller's teams
Listing the full router model_list let any authenticated key without explicit model restrictions enumerate other teams' BYOK deployments (public name, team_id, api_base) via /v1/model/info. Reuse the existing _get_caller_byok_team_scope check so non-admin callers only see global deployments plus their own team's BYOK rows; admins keep the full view.
This commit is contained in:
parent
5f701a475c
commit
2e0a32ba9f
2 changed files with 133 additions and 9 deletions
|
|
@ -11115,6 +11115,22 @@ async def _get_caller_byok_team_scope(
|
|||
return set(user_row.teams or [])
|
||||
|
||||
|
||||
def _byok_row_outside_caller_teams(
|
||||
model_info_dict: Dict[str, Any], allowed_team_ids: Optional[Set[str]]
|
||||
) -> bool:
|
||||
"""Whether a team BYOK row belongs to a team the caller is not a member of.
|
||||
|
||||
`team_id` is only set on team BYOK rows; non-team rows fall through
|
||||
unaffected. `allowed_team_ids is None` means no scoping (e.g. admins).
|
||||
"""
|
||||
if allowed_team_ids is None:
|
||||
return False
|
||||
team_id = model_info_dict.get("team_id")
|
||||
if team_id is None:
|
||||
return False
|
||||
return team_id not in allowed_team_ids
|
||||
|
||||
|
||||
# Hard cap on rows the DB-side BYOK search may pull when results need to be
|
||||
# sorted across the full match set. Without this, an authenticated caller
|
||||
# can hit `/v2/model/info?search=<broad>&sortBy=<field>` and force the
|
||||
|
|
@ -11236,15 +11252,7 @@ async def _apply_search_filter_to_models(
|
|||
)
|
||||
|
||||
def _is_byok_outside_caller_teams(model_info_dict: Dict[str, Any]) -> bool:
|
||||
# `team_id` is only set on team BYOK rows. Non-team rows fall
|
||||
# through unaffected — they are gated by other paths (router
|
||||
# membership, direct_access, include_team_models).
|
||||
if allowed_team_ids is None:
|
||||
return False
|
||||
team_id = model_info_dict.get("team_id")
|
||||
if team_id is None:
|
||||
return False
|
||||
return team_id not in allowed_team_ids
|
||||
return _byok_row_outside_caller_teams(model_info_dict, allowed_team_ids)
|
||||
|
||||
def _model_matches_search(m: Dict[str, Any]) -> bool:
|
||||
# Team BYOK models persist an internal `model_name`
|
||||
|
|
@ -12634,6 +12642,21 @@ async def model_info_v1( # noqa: PLR0915
|
|||
allowed_model_names=allowed_model_names,
|
||||
)
|
||||
|
||||
# Team BYOK deployments carry an internal routing key and other teams'
|
||||
# public name/team_id/api_base; drop the ones the caller cannot access so
|
||||
# listing the full router model_list does not leak cross-team metadata.
|
||||
allowed_team_ids = await _get_caller_byok_team_scope(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
all_models = [
|
||||
model
|
||||
for model in all_models
|
||||
if not _byok_row_outside_caller_teams(
|
||||
model.get("model_info") or {}, allowed_team_ids
|
||||
)
|
||||
]
|
||||
|
||||
all_models = [
|
||||
_translate_model_name_for_response(
|
||||
_enrich_model_info_with_litellm_data(model=model, llm_router=llm_router)
|
||||
|
|
|
|||
|
|
@ -242,3 +242,104 @@ async def test_model_info_v1_restricted_key_filters_deployments(monkeypatch):
|
|||
resp = await ps.model_info_v1(user_api_key_dict=caller, litellm_model_id=None)
|
||||
|
||||
assert [m["model_name"] for m in resp["data"]] == ["gpt-4"]
|
||||
|
||||
|
||||
def _other_team_row() -> dict:
|
||||
return {
|
||||
"model_name": "model_name_team-other_9f2c1",
|
||||
"litellm_params": {
|
||||
"model": "azure/gpt-5.2-low-rpm-testing",
|
||||
"api_base": "https://team-other-private.example.com",
|
||||
},
|
||||
"model_info": {
|
||||
"id": "byok-id-other",
|
||||
"team_id": "team-other",
|
||||
"team_public_model_name": "team-claude-sonnet",
|
||||
"db_model": True,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_model_info_v1_unrestricted_key_hides_other_team_byok(monkeypatch):
|
||||
"""Unrestricted non-admin keys must not enumerate other teams' BYOK
|
||||
deployments, but must still see global models and their own team's."""
|
||||
team_row = _team_row()
|
||||
other_team_row = _other_team_row()
|
||||
global_row = {
|
||||
"model_name": "gpt-4",
|
||||
"litellm_params": {"model": "gpt-4"},
|
||||
"model_info": {"id": "global-id-1", "db_model": False},
|
||||
}
|
||||
router = MagicMock()
|
||||
router.model_list = [team_row, other_team_row, global_row]
|
||||
router.get_model_names.return_value = ["gpt-4"]
|
||||
router.get_model_access_groups.return_value = {}
|
||||
|
||||
prisma_client = MagicMock()
|
||||
caller_user_row = MagicMock()
|
||||
caller_user_row.teams = ["team-abc-123"]
|
||||
prisma_client.db.litellm_usertable.find_unique = AsyncMock(
|
||||
return_value=caller_user_row
|
||||
)
|
||||
|
||||
monkeypatch.setattr(ps, "user_model", None)
|
||||
monkeypatch.setattr(ps, "llm_model_list", router.model_list)
|
||||
monkeypatch.setattr(ps, "llm_router", router)
|
||||
monkeypatch.setattr(ps, "prisma_client", prisma_client)
|
||||
monkeypatch.setattr(
|
||||
ps, "_enrich_model_info_with_litellm_data", lambda model, **kw: model
|
||||
)
|
||||
|
||||
caller = UserAPIKeyAuth(
|
||||
user_id="user-1",
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
models=[],
|
||||
team_models=[],
|
||||
)
|
||||
resp = await ps.model_info_v1(user_api_key_dict=caller, litellm_model_id=None)
|
||||
|
||||
returned_ids = {m["model_info"]["id"] for m in resp["data"]}
|
||||
assert returned_ids == {"global-id-1", "byok-id-1"}
|
||||
assert "byok-id-other" not in returned_ids
|
||||
names = [m["model_name"] for m in resp["data"]]
|
||||
assert "team-claude-sonnet" in names
|
||||
assert "gpt-4" in names
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_model_info_v1_service_key_hides_all_team_byok(monkeypatch):
|
||||
"""A key without a resolvable user (e.g. CI/service token) sees only
|
||||
global deployments, never any team-scoped BYOK rows."""
|
||||
team_row = _team_row()
|
||||
other_team_row = _other_team_row()
|
||||
global_row = {
|
||||
"model_name": "gpt-4",
|
||||
"litellm_params": {"model": "gpt-4"},
|
||||
"model_info": {"id": "global-id-1", "db_model": False},
|
||||
}
|
||||
router = MagicMock()
|
||||
router.model_list = [team_row, other_team_row, global_row]
|
||||
router.get_model_names.return_value = ["gpt-4"]
|
||||
router.get_model_access_groups.return_value = {}
|
||||
|
||||
prisma_client = MagicMock()
|
||||
|
||||
monkeypatch.setattr(ps, "user_model", None)
|
||||
monkeypatch.setattr(ps, "llm_model_list", router.model_list)
|
||||
monkeypatch.setattr(ps, "llm_router", router)
|
||||
monkeypatch.setattr(ps, "prisma_client", prisma_client)
|
||||
monkeypatch.setattr(
|
||||
ps, "_enrich_model_info_with_litellm_data", lambda model, **kw: model
|
||||
)
|
||||
|
||||
caller = UserAPIKeyAuth(
|
||||
user_id=None,
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
team_id="team-abc-123",
|
||||
models=[],
|
||||
team_models=[],
|
||||
)
|
||||
resp = await ps.model_info_v1(user_api_key_dict=caller, litellm_model_id=None)
|
||||
|
||||
assert [m["model_info"]["id"] for m in resp["data"]] == ["global-id-1"]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue