mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-05 02:41:56 +00:00
fix(vector_stores): run the full model grant check on caller-supplied model hints
The vector-store file routes accept a model hint through the ?model= query param and the x-litellm-model header. That hint was authorized with a hand rolled check that covered only the key allowlist and the team allowlist, so a key restricted by its project's model grant, a team-member restriction, or a key config still routed through the hinted deployment. Greptile flagged the gap as a P1 on the replacement PR. The hint now goes through the same authorize_model_for_key path the batches and files routes use, which runs can_key_call_resolved_model with every rule the proxy enforces elsewhere. Keys those extra rules deny now get a 403 on these routes. The two remaining behavioral differences are edge cases the old check tolerated: a key whose team_models is set without a team_id no longer runs the team allowlist, and a key with a config set skips the key allowlist, both matching the rest of the proxy. The regression test caches a project whose grant excludes the hinted model and asserts the request is refused before any deployment lookup. The two patch() calls on litellm.proxy.proxy_server carry a test-quality-ok reason because can_key_call_resolved_model reads prisma_client and user_api_key_cache through a lazy module import with no injection seam.
This commit is contained in:
parent
0feca8641f
commit
7133baa777
2 changed files with 40 additions and 21 deletions
|
|
@ -5,7 +5,6 @@ from fastapi.responses import ORJSONResponse
|
|||
|
||||
import litellm
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.auth.auth_checks import _can_object_call_model, can_key_call_model
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
|
||||
from litellm.proxy.common_utils.openai_endpoint_utils import (
|
||||
|
|
@ -14,6 +13,7 @@ from litellm.proxy.common_utils.openai_endpoint_utils import (
|
|||
get_custom_llm_provider_from_request_query,
|
||||
)
|
||||
from litellm.proxy.openai_files_endpoints.common_utils import (
|
||||
authorize_model_for_key,
|
||||
get_credentials_for_model,
|
||||
handle_model_based_routing,
|
||||
prepare_data_with_credentials,
|
||||
|
|
@ -212,26 +212,7 @@ async def _authorize_model_routing_hint(
|
|||
) -> None:
|
||||
if user_api_key_dict is None:
|
||||
return
|
||||
|
||||
key_models: Final = getattr(user_api_key_dict, "models", None)
|
||||
if not (isinstance(key_models, list) and "all-team-models" in key_models):
|
||||
await can_key_call_model(
|
||||
model=model,
|
||||
llm_model_list=None,
|
||||
valid_token=user_api_key_dict,
|
||||
llm_router=llm_router,
|
||||
)
|
||||
|
||||
team_models: Final = getattr(user_api_key_dict, "team_models", None)
|
||||
if isinstance(team_models, list) and len(team_models) > 0:
|
||||
_can_object_call_model(
|
||||
model=model,
|
||||
llm_router=llm_router,
|
||||
models=team_models,
|
||||
team_model_aliases=user_api_key_dict.team_model_aliases,
|
||||
team_id=user_api_key_dict.team_id,
|
||||
object_type="team",
|
||||
)
|
||||
await authorize_model_for_key(model_id=model, llm_router=llm_router, user_api_key_dict=user_api_key_dict)
|
||||
|
||||
|
||||
async def _update_request_data_with_model_routing_hint(
|
||||
|
|
|
|||
|
|
@ -609,6 +609,44 @@ async def test_vector_store_file_list_authorizes_model_query_param_before_creden
|
|||
llm_router.get_deployment_credentials_with_provider.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_vector_store_file_list_model_query_param_enforces_project_model_grant():
|
||||
from litellm.proxy._types import LiteLLM_ProjectTableCachedObj, LiteLLM_TeamTableCachedObj
|
||||
from litellm.proxy.auth.auth_checks import ProxyException
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache, project_cache_key
|
||||
|
||||
request = MagicMock(spec=Request)
|
||||
request.query_params = {"model": "team-openai"}
|
||||
request.headers = {}
|
||||
|
||||
llm_router = MagicMock()
|
||||
llm_router.model_group_alias = {}
|
||||
cache = UserApiKeyCache()
|
||||
await cache.async_set_cache(
|
||||
key="team_id:team-123",
|
||||
value=LiteLLM_TeamTableCachedObj(team_id="team-123", models=["team-openai"]),
|
||||
)
|
||||
await cache.async_set_cache(
|
||||
key=project_cache_key("proj-1"),
|
||||
value=LiteLLM_ProjectTableCachedObj(project_id="proj-1", models=["other-deployment"]),
|
||||
)
|
||||
user_api_key_dict = UserAPIKeyAuth(team_id="team-123", team_models=["team-openai"], project_id="proj-1")
|
||||
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), # test-quality-ok: proxy_server global, no seam
|
||||
patch("litellm.proxy.proxy_server.user_api_key_cache", cache), # test-quality-ok: proxy_server global, no seam
|
||||
):
|
||||
with pytest.raises(ProxyException):
|
||||
await _update_request_data_with_model_routing_hint(
|
||||
data={"vector_store_id": "vs_123"},
|
||||
request=request,
|
||||
llm_router=llm_router,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
|
||||
llm_router.get_deployment_credentials_with_provider.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_request_data_with_litellm_managed_vector_store_registry():
|
||||
"""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue