diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 4fe5dc7dd3d..b8e0aa0d128 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2569,6 +2569,9 @@ class UserAPIKeyAuth( None # Expanded created_by user when expand=user is used ) end_user_object_permission: Optional[LiteLLM_ObjectPermissionTable] = None + # Team object_permission preloaded in auth (e.g. get_team_object) to avoid + # per-request object_permission fetches in downstream checks (vector stores, etc.) + team_object_permission: Optional[LiteLLM_ObjectPermissionTable] = None # Decoded upstream IdP claims (groups, roles, etc.) propagated by JWT auth machinery # and forwarded into outbound tokens by guardrails such as MCPJWTSigner. jwt_claims: Optional[Dict] = None diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index ff4f63593cb..9779a07b97b 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -867,6 +867,11 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 ), jwt_claims=jwt_claims, ) + valid_token.team_object_permission = ( + team_object.object_permission + if team_object is not None + else None + ) # Check if model has zero cost - if so, skip all budget checks model = get_model_from_request(request_data, route) @@ -1452,6 +1457,11 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 else: _team_obj = None + if _team_obj is not None: + valid_token.team_object_permission = _team_obj.object_permission + else: + valid_token.team_object_permission = None + await user_api_key_cache.async_set_cache( key=valid_token.team_id, value=_team_obj ) # save team table in cache - used for tpm/rpm limiting - tpm_rpm_limiter.py diff --git a/litellm/proxy/vector_store_endpoints/endpoints.py b/litellm/proxy/vector_store_endpoints/endpoints.py index b7e3d8de3d3..1fdfad8c96c 100644 --- a/litellm/proxy/vector_store_endpoints/endpoints.py +++ b/litellm/proxy/vector_store_endpoints/endpoints.py @@ -10,6 +10,7 @@ from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing from litellm.proxy.utils import jsonify_object +from litellm.proxy.vector_store_endpoints.utils import can_user_access_vector_store from litellm.types.vector_stores import IndexCreateRequest router = APIRouter() @@ -18,40 +19,25 @@ router = APIRouter() ######################################################## -def _check_vector_store_access( +async def _check_vector_store_access( vector_store: LiteLLM_ManagedVectorStore, user_api_key_dict: UserAPIKeyAuth, ) -> bool: """ - Check if the user has access to the vector store based on team membership. + Check if the user has access to the vector store. - Args: - vector_store: The vector store to check access for - user_api_key_dict: User API key authentication info - - Returns: - True if user has access, False otherwise - - Access rules: - - If vector store has no team_id, it's accessible to all (legacy behavior) - - If user's team_id matches the vector store's team_id, access is granted - - Otherwise, access is denied + Delegates to :func:`can_user_access_vector_store`, which honors: + - PROXY_ADMIN bypass + - legacy vector stores with no team_id + - key-level and team-level ``object_permission.vector_stores`` allowlists + - team_id match between key and store """ - vector_store_team_id = vector_store.get("team_id") - - # If vector store has no team_id, it's accessible to all (legacy behavior) - if vector_store_team_id is None: - return True - - # Check if user's team matches the vector store's team - user_team_id = user_api_key_dict.team_id - if user_team_id == vector_store_team_id: - return True - - return False + return await can_user_access_vector_store( + vector_store=vector_store, user_api_key_dict=user_api_key_dict + ) -def _update_request_data_with_litellm_managed_vector_store_registry( +async def _update_request_data_with_litellm_managed_vector_store_registry( data: Dict, vector_store_id: str, user_api_key_dict: Optional[UserAPIKeyAuth] = None, @@ -74,9 +60,8 @@ def _update_request_data_with_litellm_managed_vector_store_registry( ) ) if vector_store_to_run is not None: - # Check access control if user_api_key_dict is provided if user_api_key_dict is not None: - if not _check_vector_store_access( + if not await _check_vector_store_access( vector_store_to_run, user_api_key_dict ): raise HTTPException( @@ -140,7 +125,7 @@ async def vector_store_search( data["vector_store_id"] = vector_store_id # Check for legacy vector store registry (non-managed vector stores) - data = _update_request_data_with_litellm_managed_vector_store_registry( + data = await _update_request_data_with_litellm_managed_vector_store_registry( data=data, vector_store_id=vector_store_id, user_api_key_dict=user_api_key_dict ) @@ -322,7 +307,7 @@ async def vector_store_retrieve( data = {"vector_store_id": vector_store_id} - data = _update_request_data_with_litellm_managed_vector_store_registry( + data = await _update_request_data_with_litellm_managed_vector_store_registry( data=data, vector_store_id=vector_store_id, user_api_key_dict=user_api_key_dict ) @@ -462,7 +447,7 @@ async def vector_store_update( if "vector_store_id" not in data: data["vector_store_id"] = vector_store_id - data = _update_request_data_with_litellm_managed_vector_store_registry( + data = await _update_request_data_with_litellm_managed_vector_store_registry( data=data, vector_store_id=vector_store_id, user_api_key_dict=user_api_key_dict ) @@ -529,7 +514,7 @@ async def vector_store_delete( data = {"vector_store_id": vector_store_id} - data = _update_request_data_with_litellm_managed_vector_store_registry( + data = await _update_request_data_with_litellm_managed_vector_store_registry( data=data, vector_store_id=vector_store_id, user_api_key_dict=user_api_key_dict ) diff --git a/litellm/proxy/vector_store_endpoints/management_endpoints.py b/litellm/proxy/vector_store_endpoints/management_endpoints.py index cf579993660..fefa6cb4e94 100644 --- a/litellm/proxy/vector_store_endpoints/management_endpoints.py +++ b/litellm/proxy/vector_store_endpoints/management_endpoints.py @@ -25,6 +25,7 @@ from litellm.proxy._types import ( from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper from litellm.proxy.common_utils.rbac_utils import check_feature_access_for_user +from litellm.proxy.vector_store_endpoints.utils import can_user_access_vector_store from litellm.secret_managers.main import get_secret from litellm.types.vector_stores import ( LiteLLM_ManagedVectorStore, @@ -274,37 +275,22 @@ async def _resolve_embedding_config( ######################################################## # Helper Functions ######################################################## -def _check_vector_store_access( +async def _check_vector_store_access( vector_store: LiteLLM_ManagedVectorStore, user_api_key_dict: UserAPIKeyAuth, ) -> bool: """ - Check if the user has access to the vector store based on team membership. + Check if the user has access to the vector store. - Args: - vector_store: The vector store to check access for - user_api_key_dict: User API key authentication info - - Returns: - True if user has access, False otherwise - - Access rules: - - If vector store has no team_id, it's accessible to all (legacy behavior) - - If user's team_id matches the vector store's team_id, access is granted - - Otherwise, access is denied + Delegates to :func:`can_user_access_vector_store`, which honors: + - PROXY_ADMIN bypass + - legacy vector stores with no team_id + - key-level and team-level ``object_permission.vector_stores`` allowlists + - team_id match between key and store """ - vector_store_team_id = vector_store.get("team_id") - - # If vector store has no team_id, it's accessible to all (legacy behavior) - if vector_store_team_id is None: - return True - - # Check if user's team matches the vector store's team - user_team_id = user_api_key_dict.team_id - if user_team_id == vector_store_team_id: - return True - - return False + return await can_user_access_vector_store( + vector_store=vector_store, user_api_key_dict=user_api_key_dict + ) async def create_vector_store_in_db( @@ -565,12 +551,11 @@ async def list_vector_stores( vector_store_id=vector_store_id, updated_data=vector_store ) - # Filter vector stores based on team access - accessible_vector_stores = [ - vs - for vs in vector_store_map.values() - if _check_vector_store_access(vs, user_api_key_dict) - ] + # Filter vector stores based on access control + accessible_vector_stores = [] + for vs in vector_store_map.values(): + if await _check_vector_store_access(vs, user_api_key_dict): + accessible_vector_stores.append(vs) total_count = len(accessible_vector_stores) total_pages = (total_count + page_size - 1) // page_size @@ -647,7 +632,7 @@ async def delete_vector_store( ) # Check access control - if vector_store_to_check and not _check_vector_store_access( + if vector_store_to_check and not await _check_vector_store_access( vector_store_to_check, user_api_key_dict ): raise HTTPException( @@ -703,7 +688,9 @@ async def get_vector_store_info( ) if vector_store is not None: # Check access control - if not _check_vector_store_access(vector_store, user_api_key_dict): + if not await _check_vector_store_access( + vector_store, user_api_key_dict + ): raise HTTPException( status_code=403, detail="Access denied: You do not have permission to access this vector store", @@ -749,7 +736,7 @@ async def get_vector_store_info( # Check access control for DB vector store vector_store_dict = vector_store.model_dump() # type: ignore[attr-defined] vector_store_typed = LiteLLM_ManagedVectorStore(**vector_store_dict) - if not _check_vector_store_access(vector_store_typed, user_api_key_dict): + if not await _check_vector_store_access(vector_store_typed, user_api_key_dict): raise HTTPException( status_code=403, detail="Access denied: You do not have permission to access this vector store", diff --git a/litellm/proxy/vector_store_endpoints/utils.py b/litellm/proxy/vector_store_endpoints/utils.py index 5499abb19d5..061a8aaa240 100644 --- a/litellm/proxy/vector_store_endpoints/utils.py +++ b/litellm/proxy/vector_store_endpoints/utils.py @@ -2,11 +2,124 @@ from typing import Any, Dict, Literal, Optional from fastapi import HTTPException, Request -from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import ( + LiteLLM_ObjectPermissionTable, + LitellmUserRoles, + UserAPIKeyAuth, +) from litellm.types.utils import LlmProviders +from litellm.types.vector_stores import LiteLLM_ManagedVectorStore from litellm.utils import ProviderConfigManager +def _is_proxy_admin(user_api_key_dict: UserAPIKeyAuth) -> bool: + return ( + user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN + or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value + ) + + +def _object_permission_allows_vector_store( + object_permission: Optional[LiteLLM_ObjectPermissionTable], + vector_store_id: str, +) -> bool: + """Returns True if an object permission explicitly allowlists the vector store.""" + if object_permission is None: + return False + allowed = object_permission.vector_stores + if not allowed: + return False + return vector_store_id in allowed + + +async def _get_object_permission_for_id( + object_permission_id: Optional[str], +) -> Optional[LiteLLM_ObjectPermissionTable]: + """Load an object permission record by id, using the shared cache/DB helper.""" + if not object_permission_id: + return None + + from litellm.proxy.auth.auth_checks import get_object_permission + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + + if prisma_client is None: + return None + + try: + return await get_object_permission( + object_permission_id=object_permission_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + except Exception as e: + verbose_proxy_logger.debug( + "Failed to load object_permission id=%s: %s", + object_permission_id, + e, + ) + return None + + +async def can_user_access_vector_store( + vector_store: LiteLLM_ManagedVectorStore, + user_api_key_dict: UserAPIKeyAuth, +) -> bool: + """ + Returns True if the caller is allowed to access this managed vector store. + + Access is granted (first match wins) when any of the following is true: + 1. The caller's role is PROXY_ADMIN. + 2. The vector store has no team_id (legacy behavior - accessible to all). + 3. The caller's key-level object_permission.vector_stores explicitly lists + this vector store id. + 4. The caller's team-level object_permission.vector_stores explicitly lists + this vector store id. + 5. The caller's team_id matches the vector store's team_id. + + Otherwise access is denied. + """ + if _is_proxy_admin(user_api_key_dict): + return True + + vector_store_team_id = vector_store.get("team_id") + if vector_store_team_id is None: + return True + + vector_store_id = vector_store.get("vector_store_id") or "" + + key_object_permission = user_api_key_dict.object_permission + if key_object_permission is None: + key_object_permission = await _get_object_permission_for_id( + user_api_key_dict.object_permission_id + ) + if _object_permission_allows_vector_store(key_object_permission, vector_store_id): + return True + + team_object_permission: Optional[LiteLLM_ObjectPermissionTable] = ( + user_api_key_dict.team_object_permission + ) + if team_object_permission is None: + team_object_permission = await _get_object_permission_for_id( + user_api_key_dict.team_object_permission_id + ) + if _object_permission_allows_vector_store(team_object_permission, vector_store_id): + return True + + if ( + user_api_key_dict.team_id is not None + and user_api_key_dict.team_id == vector_store_team_id + ): + return True + + return False + + def _does_endpoint_match(endpoint_path: str, request_path: str) -> bool: if endpoint_path in request_path: return True diff --git a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_access_control.py b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_access_control.py index 46f0ddcd582..7d72121456a 100644 --- a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_access_control.py +++ b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_access_control.py @@ -11,14 +11,19 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import HTTPException -from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy._types import ( + LiteLLM_ObjectPermissionTable, + LitellmUserRoles, + UserAPIKeyAuth, +) from litellm.proxy.vector_store_endpoints.management_endpoints import ( _check_vector_store_access, ) from litellm.types.vector_stores import LiteLLM_ManagedVectorStore -def test_check_vector_store_access(): +@pytest.mark.asyncio +async def test_check_vector_store_access(): """Test core access control logic for team-based vector store access""" # Test 1: Legacy vector stores (no team_id) are accessible to all @@ -28,7 +33,7 @@ def test_check_vector_store_access(): "team_id": None, } user = UserAPIKeyAuth(team_id="team_456") - assert _check_vector_store_access(vector_store, user) is True + assert await _check_vector_store_access(vector_store, user) is True # Test 2: User can access their team's vector stores vector_store = { @@ -37,7 +42,7 @@ def test_check_vector_store_access(): "team_id": "team_456", } user = UserAPIKeyAuth(team_id="team_456") - assert _check_vector_store_access(vector_store, user) is True + assert await _check_vector_store_access(vector_store, user) is True # Test 3: User cannot access other teams' vector stores vector_store = { @@ -46,7 +51,57 @@ def test_check_vector_store_access(): "team_id": "team_456", } user = UserAPIKeyAuth(team_id="team_789") - assert _check_vector_store_access(vector_store, user) is False + assert await _check_vector_store_access(vector_store, user) is False + + +@pytest.mark.asyncio +async def test_check_vector_store_access_proxy_admin_bypass(): + """PROXY_ADMIN can access a vector store even if teams don't match.""" + vector_store: LiteLLM_ManagedVectorStore = { + "vector_store_id": "vs_team", + "custom_llm_provider": "openai", + "team_id": "team_456", + } + admin = UserAPIKeyAuth(team_id="team_999", user_role=LitellmUserRoles.PROXY_ADMIN) + assert await _check_vector_store_access(vector_store, admin) is True + + +@pytest.mark.asyncio +async def test_check_vector_store_access_key_object_permission_grants_access(): + """A key whose object_permission.vector_stores allowlists the store can access it + even if its team_id does not match the store's team_id.""" + vector_store: LiteLLM_ManagedVectorStore = { + "vector_store_id": "vs_explicit", + "custom_llm_provider": "openai", + "team_id": "team_456", + } + user = UserAPIKeyAuth( + team_id="team_789", + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id="op-1", + vector_stores=["vs_explicit"], + ), + ) + assert await _check_vector_store_access(vector_store, user) is True + + +@pytest.mark.asyncio +async def test_check_vector_store_access_key_object_permission_wrong_store_denied(): + """A key whose object_permission.vector_stores lists *other* stores is still denied + when the key has no other reason to access this store.""" + vector_store: LiteLLM_ManagedVectorStore = { + "vector_store_id": "vs_target", + "custom_llm_provider": "openai", + "team_id": "team_456", + } + user = UserAPIKeyAuth( + team_id="team_789", + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id="op-1", + vector_stores=["vs_other"], + ), + ) + assert await _check_vector_store_access(vector_store, user) is False @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py index c1dc0cba02c..44cc5cc4452 100644 --- a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py +++ b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py @@ -115,7 +115,8 @@ def test_router_vector_store_file_delete_passes_correct_args(): assert call_kwargs["custom_llm_provider"] == "openai" -def test_update_request_data_with_litellm_managed_vector_store_registry(): +@pytest.mark.asyncio +async def test_update_request_data_with_litellm_managed_vector_store_registry(): """ Test that _update_request_data_with_litellm_managed_vector_store_registry correctly updates request data with vector store registry information. @@ -139,7 +140,7 @@ def test_update_request_data_with_litellm_managed_vector_store_registry(): # Test with vector store registry with patch.object(litellm, "vector_store_registry", mock_registry): - result = _update_request_data_with_litellm_managed_vector_store_registry( + result = await _update_request_data_with_litellm_managed_vector_store_registry( data=data, vector_store_id=vector_store_id ) @@ -158,7 +159,7 @@ def test_update_request_data_with_litellm_managed_vector_store_registry(): # Test with no vector store registry with patch.object(litellm, "vector_store_registry", None): original_data = {"existing_key": "existing_value"} - result = _update_request_data_with_litellm_managed_vector_store_registry( + result = await _update_request_data_with_litellm_managed_vector_store_registry( data=original_data, vector_store_id=vector_store_id ) @@ -1686,10 +1687,26 @@ async def test_new_vector_store_auto_resolves_from_router(): ) +def _stub_user_api_key( + *, + team_id=None, + user_role=None, + object_permission=None, + object_permission_id=None, + team_object_permission_id=None, +): + user = UserAPIKeyAuth(team_id=team_id, user_role=user_role) + user.object_permission = object_permission + user.object_permission_id = object_permission_id + user.team_object_permission_id = team_object_permission_id + return user + + class TestCheckVectorStoreAccess: """Test suite for _check_vector_store_access function.""" - def test_access_granted_when_no_team_id(self): + @pytest.mark.asyncio + async def test_access_granted_when_no_team_id(self): """Test that access is granted when vector store has no team_id (legacy behavior).""" vector_store: LiteLLM_ManagedVectorStore = { "vector_store_id": "test-store", @@ -1697,13 +1714,12 @@ class TestCheckVectorStoreAccess: # No team_id field } - mock_user_api_key = MagicMock(spec=UserAPIKeyAuth) - mock_user_api_key.team_id = "team-123" - - result = _check_vector_store_access(vector_store, mock_user_api_key) + user = _stub_user_api_key(team_id="team-123") + result = await _check_vector_store_access(vector_store, user) assert result is True - def test_access_granted_when_team_ids_match(self): + @pytest.mark.asyncio + async def test_access_granted_when_team_ids_match(self): """Test that access is granted when user's team_id matches vector store's team_id.""" vector_store: LiteLLM_ManagedVectorStore = { "vector_store_id": "test-store", @@ -1711,13 +1727,12 @@ class TestCheckVectorStoreAccess: "team_id": "team-123", } - mock_user_api_key = MagicMock(spec=UserAPIKeyAuth) - mock_user_api_key.team_id = "team-123" - - result = _check_vector_store_access(vector_store, mock_user_api_key) + user = _stub_user_api_key(team_id="team-123") + result = await _check_vector_store_access(vector_store, user) assert result is True - def test_access_denied_when_team_ids_dont_match(self): + @pytest.mark.asyncio + async def test_access_denied_when_team_ids_dont_match(self): """Test that access is denied when user's team_id doesn't match vector store's team_id.""" vector_store: LiteLLM_ManagedVectorStore = { "vector_store_id": "test-store", @@ -1725,13 +1740,12 @@ class TestCheckVectorStoreAccess: "team_id": "team-123", } - mock_user_api_key = MagicMock(spec=UserAPIKeyAuth) - mock_user_api_key.team_id = "team-456" - - result = _check_vector_store_access(vector_store, mock_user_api_key) + user = _stub_user_api_key(team_id="team-456") + result = await _check_vector_store_access(vector_store, user) assert result is False - def test_access_denied_when_vector_store_has_team_id_but_user_doesnt(self): + @pytest.mark.asyncio + async def test_access_denied_when_vector_store_has_team_id_but_user_doesnt(self): """Test that access is denied when vector store has team_id but user doesn't.""" vector_store: LiteLLM_ManagedVectorStore = { "vector_store_id": "test-store", @@ -1739,10 +1753,8 @@ class TestCheckVectorStoreAccess: "team_id": "team-123", } - mock_user_api_key = MagicMock(spec=UserAPIKeyAuth) - mock_user_api_key.team_id = None - - result = _check_vector_store_access(vector_store, mock_user_api_key) + user = _stub_user_api_key(team_id=None) + result = await _check_vector_store_access(vector_store, user) assert result is False