From cf2cbab426e676464ca4824808af16e053cec41d Mon Sep 17 00:00:00 2001 From: shivam Date: Sat, 16 May 2026 14:40:34 -0700 Subject: [PATCH] fix(proxy): bound BYOK search DB fetch in /v2/model/info Previously the DB-side search OR'd a JSON-path predicate `{model_info: {path: [team_public_model_name], string_contains: ""}}` to compensate for Prisma's case-sensitive JSON `string_contains` on Postgres. That predicate matches every row that has any `team_public_model_name` set, so any authenticated caller could force a full BYOK-table read with `/v2/model/info?search=x` regardless of page size. Drop the JSON-path branch. The DB query now does a bounded `model_name contains ` lookup. BYOK rows that are loaded into the router are still searchable by their `team_public_model_name` via the router-side filter; only the rare edge case of a BYOK row that exists only in the DB (router sync failed) loses display-name search, which is an acceptable trade-off given the DoS surface. Co-Authored-By: Claude Opus 4.7 (1M context) --- litellm/proxy/proxy_server.py | 71 ++++++---------- tests/test_litellm/proxy/test_proxy_server.py | 81 +------------------ 2 files changed, 30 insertions(+), 122 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 9fddae3c7ee..e09ab2fe934 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -11020,30 +11020,24 @@ async def _apply_search_filter_to_models( # Only query database if prisma_client is available if prisma_client is not None: try: - # Prisma's JSON path `string_contains` is case-sensitive in - # Postgres (it doesn't accept the `mode: insensitive` flag the - # way column-level string filters do), so the BYOK branch - # below can't match mixed-case stored names like - # "Claude Sonnet" against a lowercased search term. Widen the - # JSON branch to "row has a team_public_model_name set" - # (`string_contains: ""` matches any string at the path) and - # filter case-insensitively in Python below so behavior - # matches the router-side path in `_model_matches_search`. + # Match the persisted internal `model_name` only. Earlier + # iterations of this code also OR'd a JSON-path match on + # `model_info.team_public_model_name` so BYOK rows present only + # in the DB (not the router) were searchable by display name. + # That branch had to fall back to `string_contains: ""` because + # Prisma's JSON `string_contains` is case-sensitive on + # Postgres, which made the predicate match every row with any + # `team_public_model_name` set — an authenticated user could + # force a full BYOK-table read with `/v2/model/info?search=x`. + # BYOK rows that are actually loaded into the router are + # already searchable via the router-side path above (which + # checks `team_public_model_name`), so we drop the unbounded + # JSON branch here to keep the DB cost bounded by `search`. db_where_condition: Dict[str, Any] = { - "OR": [ - { - "model_name": { - "contains": search_lower, - "mode": "insensitive", - } - }, - { - "model_info": { - "path": ["team_public_model_name"], - "string_contains": "", - } - }, - ] + "model_name": { + "contains": search_lower, + "mode": "insensitive", + } } # Exclude models already in router if we have any if db_model_ids_in_router: @@ -11051,38 +11045,25 @@ async def _apply_search_filter_to_models( "not": {"in": list(db_model_ids_in_router)} } - # Fetch all candidates and filter in Python. We can't trust a - # DB-level count because the BYOK branch is over-broad — it - # returns every row with a team_public_model_name regardless - # of whether it matches the search term. db_models_raw = await prisma_client.db.litellm_proxymodeltable.find_many( where=db_where_condition, ) - def _db_row_matches_search(db_model: Any) -> bool: - info = ( - db_model.model_info if isinstance(db_model.model_info, dict) else {} + # Scope BYOK rows to the caller's allowed teams so non-admin + # callers can't enumerate other teams' BYOK metadata via + # `/v2/model/info?search=...`. + matching_db_rows = [ + m + for m in db_models_raw + if not _is_byok_outside_caller_teams( + m.model_info if isinstance(m.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() - ) - - matching_db_rows = [m for m in db_models_raw if _db_row_matches_search(m)] + ] db_models_total_count = len(matching_db_rows) # Calculate total count for search results search_total_count = router_models_count + db_models_total_count - # Decrypt matching rows. Done after the in-Python filter so we - # don't decrypt BYOK rows we're going to throw away. for db_model in matching_db_rows: decrypted_models = proxy_config.decrypt_model_list_from_db([db_model]) if decrypted_models: diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 491bd172554..21584fb2ff6 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -1225,87 +1225,14 @@ async def test_apply_search_filter_matches_team_public_model_name(): assert filtered == [] -@pytest.mark.asyncio -async def test_apply_search_filter_matches_db_byok_case_insensitively(): - """ - Regression test: BYOK rows that only exist in the DB (not in the - in-memory router) must still match search case-insensitively against - their stored `team_public_model_name`. Prisma's JSON path - `string_contains` is case-sensitive in Postgres, so a search like - "claude" must still match a stored value of "Claude Sonnet". - """ - from litellm.proxy.proxy_server import _apply_search_filter_to_models - - # Stored team_public_model_name uses mixed case; lowercased search - # would never match it via Prisma's case-sensitive JSON filter. - byok_db_row = MagicMock() - byok_db_row.model_id = "byok-db-only" - byok_db_row.model_name = "model_name_team-xyz_internal" - byok_db_row.model_info = { - "id": "byok-db-only", - "team_id": "team-xyz", - "team_public_model_name": "Claude Sonnet 4.6", - "db_model": True, - } - # Decoy row with team_public_model_name set but not matching the - # search — verifies the Python filter prunes the over-broad DB query. - decoy_db_row = MagicMock() - decoy_db_row.model_id = "decoy-db-only" - decoy_db_row.model_name = "model_name_team-xyz_decoy" - decoy_db_row.model_info = { - "id": "decoy-db-only", - "team_id": "team-xyz", - "team_public_model_name": "Some Gemini Variant", - "db_model": True, - } - - prisma_client = MagicMock() - prisma_client.db.litellm_proxymodeltable.find_many = AsyncMock( - return_value=[byok_db_row, decoy_db_row] - ) - - proxy_config = MagicMock() - - # decrypt_model_list_from_db echoes back a router-shaped dict; mock - # it so we can identify which DB row(s) survived the Python filter. - def _fake_decrypt(rows): - return [ - { - "model_name": r.model_name, - "model_info": r.model_info, - "litellm_params": {"model": "claude-sonnet"}, - } - for r in rows - ] - - proxy_config.decrypt_model_list_from_db = _fake_decrypt - - filtered, total_count = await _apply_search_filter_to_models( - all_models=[], - search="claude", - prisma_client=prisma_client, - proxy_config=proxy_config, - ) - - filtered_ids = {m["model_info"]["id"] for m in filtered} - assert ( - "byok-db-only" in filtered_ids - ), "mixed-case team_public_model_name must match lowercased search" - assert ( - "decoy-db-only" not in filtered_ids - ), "non-matching BYOK row fetched by over-broad query must be filtered out" - 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 teams the caller is not a member of. Even with a bounded + `model_name`-contains DB query, a non-admin caller could otherwise + see other teams' BYOK rows that happen to match by internal name. + The post-fetch team scope drops those. """ from litellm.proxy.proxy_server import _apply_search_filter_to_models