diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 8270314faf1..f87400371b3 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -2,8 +2,9 @@ import importlib.util import os import re import sys -from collections.abc import Collection, Iterator, Mapping, Sequence -from functools import lru_cache +from collections.abc import Callable, Collection, Iterator, Mapping, Sequence +from dataclasses import dataclass +from functools import cache, lru_cache from itertools import accumulate, chain, repeat from logging import Logger from types import MappingProxyType @@ -546,16 +547,63 @@ def _opt_in_model_groups( return routed | frozenset(registry.resolve(name, kind="alias") or name for name in fallbacks) +_MAX_OPT_IN_FALLBACK_MODELS: Final = 100 + + +@dataclass(frozen=True, slots=True) +class _OptInGroups: + every: tuple[str, ...] + vertex_target: tuple[str, ...] + + +def _deployment_id(deployment: Mapping[str, object]) -> str | None: + match deployment.get("model_info"): + case {"id": str(deployment_id)} if deployment_id: + return deployment_id + case _: + return None + + +def _same_deployments_key(model_group: str, llm_router: Router) -> frozenset[tuple[str | None, str | None]] | str: + deployments: Final = _model_group_deployments(model_group, llm_router) + key: Final = tuple( + (_deployment_id(deployment), _deployment_provider(deployment["litellm_params"])) for deployment in deployments + ) + return frozenset(key) if key and all(deployment_id for deployment_id, _ in key) else model_group + + +def _distinct_opt_in_groups(model_groups: frozenset[str], llm_router: Router | None) -> tuple[str, ...]: + if llm_router is None: + return tuple(sorted(model_groups)) + return tuple( + {_same_deployments_key(model_group, llm_router): model_group for model_group in sorted(model_groups)}.values() + ) + + +def _resolve_opt_in_groups( + model: str, also_requested_models: tuple[str, ...], llm_router: Router | None +) -> _OptInGroups: + if len(frozenset(also_requested_models)) > _MAX_OPT_IN_FALLBACK_MODELS: + raise ValueError( + f"Rejected Request: the request names more than {_MAX_OPT_IN_FALLBACK_MODELS} fallback models, " + "too many to check against configurable_clientside_auth_params." + ) + every: Final = _distinct_opt_in_groups(_opt_in_model_groups(model, also_requested_models, llm_router), llm_router) + return _OptInGroups( + every=every, + vertex_target=tuple(group for group in every if _model_group_reads_vertex_target(group, llm_router)), + ) + + def _check_banned_params( body: Mapping[str, object], general_settings: dict, llm_router: Router | None, - model: str, *, manages_deployments: bool = False, mcp_server_fields_allowed: bool = False, banned_params: tuple[str, ...] = _BANNED_REQUEST_BODY_PARAMS, - also_requested_models: tuple[str, ...] = (), + opt_in_groups: Callable[[], _OptInGroups], ) -> None: """Raise ``ValueError`` if ``body`` carries a banned param without admin opt-in. @@ -585,8 +633,9 @@ def _check_banned_params( llm_router=llm_router, ) is True - for model_group in _opt_in_model_groups(model, also_requested_models, llm_router) - if param not in VERTEX_TARGET_PARAMS or _model_group_reads_vertex_target(model_group, llm_router) + for model_group in ( + opt_in_groups().vertex_target if param in VERTEX_TARGET_PARAMS else opt_in_groups().every + ) ): # Per-param opt-in: only THIS param is permitted by the # deployment's ``configurable_clientside_auth_params``. Skip @@ -780,14 +829,14 @@ def is_request_body_safe( *request_fallback_model_names(root_body), *((route_model,) if route_model else ()), ) + opt_in_groups: Final = cache(lambda: _resolve_opt_in_groups(model, also_requested_models, llm_router)) _check_banned_params( root_body, general_settings, llm_router, - model, manages_deployments=manages_deployments, mcp_server_fields_allowed=route_manages_mcp_servers(route), - also_requested_models=also_requested_models, + opt_in_groups=opt_in_groups, ) for nested_key in _NESTED_CONFIG_KEYS: nested = _coerce_metadata_to_dict(request_body.get(nested_key)) @@ -796,9 +845,8 @@ def is_request_body_safe( nested, general_settings, llm_router, - model, manages_deployments=manages_deployments, - also_requested_models=also_requested_models, + opt_in_groups=opt_in_groups, ) for metadata_key in _NESTED_METADATA_KEYS: metadata = _coerce_metadata_to_dict(request_body.get(metadata_key)) @@ -807,20 +855,18 @@ def is_request_body_safe( metadata, general_settings, llm_router, - model, manages_deployments=manages_deployments, banned_params=_BANNED_METADATA_PARAMS, - also_requested_models=also_requested_models, + opt_in_groups=opt_in_groups, ) if any(isinstance(key, str) and key.startswith(f"{metadata_key}[") for key in request_body): _check_banned_params( extract_nested_form_metadata(form_data=request_body, prefix=f"{metadata_key}["), general_settings, llm_router, - model, manages_deployments=manages_deployments, banned_params=_BANNED_METADATA_PARAMS, - also_requested_models=also_requested_models, + opt_in_groups=opt_in_groups, ) for target in iter_request_fallback_targets(request_body): if isinstance(target, dict): @@ -828,9 +874,8 @@ def is_request_body_safe( target, general_settings, llm_router, - model, manages_deployments=manages_deployments, - also_requested_models=also_requested_models, + opt_in_groups=opt_in_groups, ) target_model = target.get("model") if isinstance(target_model, str): @@ -848,10 +893,9 @@ def is_request_body_safe( litellm_params_metadata, general_settings, llm_router, - model, manages_deployments=manages_deployments, banned_params=_BANNED_METADATA_PARAMS, - also_requested_models=also_requested_models, + opt_in_groups=opt_in_groups, ) for user_config_params in _user_config_litellm_params(request_body): # pyright: ignore[reportUnknownArgumentType] # request_body is an untyped dict reject_api_key_secret_reference(user_config_params) diff --git a/litellm/proxy/rag_endpoints/endpoints.py b/litellm/proxy/rag_endpoints/endpoints.py index aee958fae3d..e105afda33c 100644 --- a/litellm/proxy/rag_endpoints/endpoints.py +++ b/litellm/proxy/rag_endpoints/endpoints.py @@ -234,6 +234,14 @@ def _corpus_on_configured_project(vector_store_config: Mapping[str, object]) -> ) +async def _authorize_rewritten_corpus( + requested: Mapping[str, object], rewritten: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth +) -> None: + rewritten_id: Final = rewritten.get("vector_store_id") + if isinstance(rewritten_id, str) and rewritten_id != requested.get("vector_store_id"): + await assert_user_can_access_vector_store_id(vector_store_id=rewritten_id, user_api_key_dict=user_api_key_dict) + + def _vertex_target_error(vector_store_config: Mapping[str, object]) -> str | None: corpus_name: Final = vector_store_config.get("vector_store_id") corpus: Final = _clean_vertex_corpus_name(corpus_name) @@ -793,6 +801,7 @@ async def rag_ingest( requested_base, user_api_key_dict, general_settings_view() ) and not await _admin_set_store_target(managed_store, prisma_client) base: Final = _corpus_on_configured_project(requested_base) if vertex_target_checked else requested_base + await _authorize_rewritten_corpus(requested_base, base, user_api_key_dict) vertex_target_error: Final = _vertex_target_error(base) if vertex_target_checked else None if vertex_target_error is not None: raise HTTPException(status_code=400, detail={"error": vertex_target_error}) diff --git a/tests/unit/proxy/auth/test_auth_utils.py b/tests/unit/proxy/auth/test_auth_utils.py index 68b2326ac64..4db586645f1 100644 --- a/tests/unit/proxy/auth/test_auth_utils.py +++ b/tests/unit/proxy/auth/test_auth_utils.py @@ -4,6 +4,7 @@ Unit tests for auth_utils functions related to rate limiting and customer ID ext import base64 import logging +import time from collections.abc import Callable from typing import Final, Optional from unittest.mock import MagicMock, patch @@ -5265,6 +5266,39 @@ class TestIsRequestBodySafeBlocksVertexTargetOverrides: with pytest.raises(ValueError, match=r"vertex_location|api_base"): is_request_body_safe(request_body, {}, router, model, route="/v1/chat/completions") + def test_many_fallback_entries_on_one_wildcard_are_checked_in_linear_time(self) -> None: + router: Final = Router(model_list=[{"model_name": "*", "litellm_params": {"model": "openai/*", "api_key": "k"}}]) + request_body: Final = { + "model": "gpt-4o-mini", + "fallbacks": [{"model": f"model-{index % 100}", "vertex_location": f"r{index}"} for index in range(5000)], + } + + started: Final = time.perf_counter() + assert is_request_body_safe(request_body, {}, router, "gpt-4o-mini", route="/v1/chat/completions") is True + assert time.perf_counter() - started < 10 + + def test_too_many_fallback_models_for_an_opt_in_check_are_rejected(self) -> None: + router: Final = Router(model_list=[{"model_name": "*", "litellm_params": {"model": "openai/*", "api_key": "k"}}]) + request_body: Final = { + "model": "gpt-4o-mini", + "vertex_location": "us-central1", + "fallbacks": [f"model-{index}" for index in range(101)], + } + + with pytest.raises(ValueError, match="more than 100 fallback models"): + is_request_body_safe(request_body, {}, router, "gpt-4o-mini", route="/v1/chat/completions") + + def test_one_wildcard_serving_two_providers_checks_the_google_one(self) -> None: + router: Final = Router(model_list=[{"model_name": "*", "litellm_params": {"model": "*"}}]) + request_body: Final = { + "model": "vertex_ai/gemini-2.5-flash", + "vertex_location": "us-central1", + "fallbacks": ["xai/grok-3"], + } + + with pytest.raises(ValueError, match="vertex_location"): + is_request_body_safe(request_body, {}, router, "vertex_ai/gemini-2.5-flash", route="/v1/chat/completions") + @pytest.mark.parametrize( "model, is_allowed", [("public-alias", True), ("gemini-optin", False)], ids=["alias", "aliased_target"] ) diff --git a/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py b/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py index 620dc13aacb..3d92a812ef4 100644 --- a/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py +++ b/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py @@ -1143,6 +1143,41 @@ _PROXY_CORPUS_PREFIX: Final = f"projects/{_PROXY_VERTEX_PROJECT}/locations/{_PRO _FOREIGN_CORPUS: Final = f"projects/other-project/locations/{_PROXY_VERTEX_LOCATION}/ragCorpora/1" +@pytest.mark.usefixtures("proxy_vertex_target") +def test_rag_ingest_refuses_another_teams_store_named_by_a_project_number(client_internal_user): + other_teams_store: Final = LiteLLM_ManagedVectorStore( + vector_store_id=f"{_PROXY_CORPUS_PREFIX}/1", + custom_llm_provider="vertex_ai", + litellm_params={}, + user_id="other-user", + team_id="other-team", + ) + registry: Final = MagicMock() + registry.get_litellm_managed_vector_store_from_registry.side_effect = lambda vector_store_id: ( + other_teams_store if vector_store_id == other_teams_store["vector_store_id"] else None + ) + with ( + patch( # test-quality-ok: aingest is the endpoint's downstream boundary; the test asserts it is never called + "litellm.proxy.rag_endpoints.endpoints.litellm.aingest", + new=AsyncMock(return_value={"vector_store_id": f"{_PROXY_CORPUS_PREFIX}/1", "file_id": "file_123"}), + ) as mock_aingest, + patch.object(litellm, "vector_store_registry", registry), # test-quality-ok: seeds the saved stores by id + _patched_prisma_client(None), + ): + response = client_internal_user.post( + "/v1/rag/ingest", + **_ingest_form( + { + "custom_llm_provider": "vertex_ai", + "vector_store_id": f"projects/0/locations/{_PROXY_VERTEX_LOCATION}/ragCorpora/1", + } + ), + ) + + assert response.status_code == 403, response.json() + mock_aingest.assert_not_awaited() + + @pytest.mark.parametrize( "vector_store_id", [