scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-10 03:28:53 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 85

fix(proxy): scope BYOK search results to caller's accessible teams

Browse source 
`_apply_search_filter_to_models` was widened to fetch every row with a
`team_public_model_name` set so case-insensitive search could match
mixed-case stored names. `/v2/model/info` is reachable by non-admin
keys though, and the helper ran before `include_team_models` / `teamId`
filtering — so a non-admin caller could search a common substring like
"claude" and see BYOK rows belonging to teams they're not a member of.

Resolve the caller's team membership once (admin → no scoping, else
their `user_row.teams`) and drop BYOK rows (those with
`model_info.team_id` set) outside that scope on both the router-side
matches and the over-broad DB query, before display-name matching.
Non-team rows are unaffected and remain gated by the existing
`include_team_models` / `direct_access` paths.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
shivam 2026-05-16 14:00:32 -07:00
parent ce87a78459
commit 7e45bede3b
No known key found for this signature in database
2 changed files with 201 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

80
litellm/proxy/proxy_server.py
View file

@ -10889,11 +10889,48 @@ def _enrich_model_info_with_litellm_data(
return model
async def _get_caller_byok_team_scope(
user_api_key_dict: Optional[UserAPIKeyAuth],
prisma_client: Optional[Any],
) -> Optional[Set[str]]:
"""
Return the team IDs whose BYOK rows the caller is allowed to see via
`/v2/model/info` search results.
`None` means "no scoping" — used for admins and for callers/paths that
have already been scoped upstream (or in tests that supply their own
pre-filtered input set). A returned set (possibly empty) means BYOK rows
must have `model_info.team_id` ∈ that set, otherwise they belong to a
team the caller is not a member of and must be dropped.
"""
if user_api_key_dict is None or prisma_client is None:
return None
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN:
return None
user_id = user_api_key_dict.user_id
if user_id is None:
return set()
try:
user_row = await prisma_client.db.litellm_usertable.find_unique(
where={"user_id": user_id}
)
except Exception:
verbose_proxy_logger.exception(
"Failed to look up caller teams while scoping BYOK search; "
"defaulting to no team access."
)
return set()
if user_row is None:
return set()
return set(user_row.teams or [])
async def _apply_search_filter_to_models(
all_models: List[Dict[str, Any]],
search: str,
prisma_client: Optional[Any],
proxy_config: Any,
user_api_key_dict: Optional[UserAPIKeyAuth] = None,
) -> Tuple[List[Dict[str, Any]], Optional[int]]:
"""
Apply search filter to models, querying database for additional matching models.
@ -10903,6 +10940,10 @@ async def _apply_search_filter_to_models(
search: Search term (case-insensitive)
prisma_client: Prisma client for database queries
proxy_config: Proxy config for decrypting models
user_api_key_dict: Caller identity used to scope BYOK matches to
teams the caller belongs to. When omitted (None), no team
scoping is applied — pass it from request handlers that expose
this function to non-admin callers.
Returns:
Tuple of (filtered_models, total_count). total_count is None if not searching.
@ -10912,6 +10953,22 @@ async def _apply_search_filter_to_models(
search_lower = search.lower().strip()
allowed_team_ids = await _get_caller_byok_team_scope(
user_api_key_dict=user_api_key_dict,
prisma_client=prisma_client,
)
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
def _model_matches_search(m: Dict[str, Any]) -> bool:
# Team BYOK models persist an internal `model_name`
# (e.g. `model_name_{team_id}_{uuid}`) and expose the user-facing
@ -10924,8 +10981,16 @@ async def _apply_search_filter_to_models(
) or ""
return search_lower in team_public_model_name.lower()
# Filter models in router by search term
filtered_router_models = [m for m in all_models if _model_matches_search(m)]
# Filter models in router by search term, dropping BYOK rows that
# belong to teams the caller is not a member of so search can't leak
# other teams' models when the request omits `include_team_models` /
# `teamId`.
filtered_router_models = [
m
for m in all_models
if _model_matches_search(m)
and not _is_byok_outside_caller_teams(m.get("model_info") or {})
]
# Separate filtered models into config vs db models, and track db model IDs
filtered_config_models = []
@ -10992,11 +11057,17 @@ async def _apply_search_filter_to_models(
)
def _db_row_matches_search(db_model: Any) -> bool:
if search_lower in (db_model.model_name or "").lower():
return True
info = (
db_model.model_info if isinstance(db_model.model_info, dict) else {}
)
# Scope BYOK rows to the caller's teams before applying
# the display-name match, so the over-broad
# `string_contains: ""` JSON branch can't leak other
# teams' models into search results.
if _is_byok_outside_caller_teams(info):
return False
if search_lower in (db_model.model_name or "").lower():
return True
return (
search_lower in (info.get("team_public_model_name") or "").lower()
)
@ -11491,6 +11562,7 @@ async def model_info_v2(
search=search or "",
prisma_client=prisma_client,
proxy_config=proxy_config,
user_api_key_dict=user_api_key_dict,
)
if user_models_only:

125
tests/test_litellm/proxy/test_proxy_server.py
View file

@ -1297,6 +1297,131 @@ async def test_apply_search_filter_matches_db_byok_case_insensitively():
assert total_count == 1
@pytest.mark.asyncio
async def test_apply_search_filter_scopes_byok_to_caller_teams():
"""
Regression test: `/v2/model/info?search=...` must not leak BYOK rows
from teams the caller is not a member of. The new DB branch fetches
every row with a `team_public_model_name` set (the JSON
`string_contains: ""` widening), so without team scoping a non-admin
user could search for a common substring like "claude" and see other
teams' BYOK models.
"""
from litellm.proxy.proxy_server import _apply_search_filter_to_models
# In-router BYOK rows: one in the caller's team, one in someone else's.
caller_team_byok = {
"model_name": "model_name_team-mine_internal",
"litellm_params": {"model": "claude-sonnet"},
"model_info": {
"id": "byok-mine",
"team_id": "team-mine",
"team_public_model_name": "claude-sonnet-prod",
"db_model": True,
},
}
other_team_byok = {
"model_name": "model_name_team-other_internal",
"litellm_params": {"model": "claude-sonnet"},
"model_info": {
"id": "byok-other",
"team_id": "team-other",
"team_public_model_name": "claude-sonnet-staging",
"db_model": True,
},
}
# Non-team row stays in the router-side result regardless of teams.
public_model = {
"model_name": "claude-public",
"litellm_params": {"model": "claude-sonnet"},
"model_info": {"id": "public-id", "db_model": False},
}
# DB-only BYOK rows fetched by the over-broad JSON branch.
db_caller_row = MagicMock()
db_caller_row.model_id = "byok-db-mine"
db_caller_row.model_name = "model_name_team-mine_db"
db_caller_row.model_info = {
"id": "byok-db-mine",
"team_id": "team-mine",
"team_public_model_name": "Claude DB Mine",
"db_model": True,
}
db_other_row = MagicMock()
db_other_row.model_id = "byok-db-other"
db_other_row.model_name = "model_name_team-other_db"
db_other_row.model_info = {
"id": "byok-db-other",
"team_id": "team-other",
"team_public_model_name": "Claude DB Other",
"db_model": True,
}
prisma_client = MagicMock()
prisma_client.db.litellm_proxymodeltable.find_many = AsyncMock(
return_value=[db_caller_row, db_other_row]
)
caller_user_row = MagicMock()
caller_user_row.teams = ["team-mine"]
prisma_client.db.litellm_usertable.find_unique = AsyncMock(
return_value=caller_user_row
)
proxy_config = MagicMock()
proxy_config.decrypt_model_list_from_db = lambda rows: [
{
"model_name": r.model_name,
"model_info": r.model_info,
"litellm_params": {"model": "claude-sonnet"},
}
for r in rows
]
non_admin = MagicMock(spec=UserAPIKeyAuth)
non_admin.user_role = LitellmUserRoles.INTERNAL_USER
non_admin.user_id = "user-mine"
filtered, total_count = await _apply_search_filter_to_models(
all_models=[caller_team_byok, other_team_byok, public_model],
search="claude",
prisma_client=prisma_client,
proxy_config=proxy_config,
user_api_key_dict=non_admin,
)
filtered_ids = {m["model_info"]["id"] for m in filtered}
assert "byok-mine" in filtered_ids
assert "byok-db-mine" in filtered_ids
assert "public-id" in filtered_ids
assert "byok-other" not in filtered_ids, (
"router-side BYOK from another team must be dropped from search "
"when caller doesn't belong to that team"
)
assert "byok-db-other" not in filtered_ids, (
"DB-only BYOK from another team must be dropped from search when "
"caller doesn't belong to that team"
)
# total_count is router_models_count (3: caller_team_byok, public_model
# would be counted; other_team_byok is dropped — 2) + 1 db row = 3.
assert total_count == 3
# Admins keep the un-scoped view across teams.
admin = MagicMock(spec=UserAPIKeyAuth)
admin.user_role = LitellmUserRoles.PROXY_ADMIN
admin.user_id = "admin-1"
filtered_admin, _ = await _apply_search_filter_to_models(
all_models=[caller_team_byok, other_team_byok, public_model],
search="claude",
prisma_client=prisma_client,
proxy_config=proxy_config,
user_api_key_dict=admin,
)
admin_ids = {m["model_info"]["id"] for m in filtered_admin}
assert "byok-other" in admin_ids
assert "byok-db-other" in admin_ids
@pytest.mark.asyncio
async def test_filter_models_by_team_id_excludes_viewer_direct_access():
"""