From 7e49b4e2a016743e97a60777d49e67d0076bf867 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 30 May 2025 22:20:11 -0700 Subject: [PATCH] [Feat] Enforce Vector Store Access Controls on LiteLLM Auth (#11281) * fix LiteLLM_ObjectPermissionTable * fix include object_permission for list key * fix key list to inclue obj permissions * fix object permissions for vector stores on key info * add key edit view with vector stores * allow editing vector stores permissions * fixes obj permissions * feat: add obj permission on UI * fix: add object_permission:true * ui show org vector stores on org info * fix: show object permissions on /org/info * feat: allow updating obj permissions for keys * fixes: key object permissions * fixes: team object permissions * fixes: org object permissions * fix vector store selector for Orgs * feat: add auth checks for vector store permissions * feat: working auth checks for vector store permissions * test vector stores auth checks * Update litellm/proxy/_types.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * fix: linting --------- Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- litellm/proxy/_types.py | 33 +++- litellm/proxy/auth/auth_checks.py | 110 ++++++++++++ litellm/proxy/proxy_config.yaml | 11 -- .../proxy/auth/test_auth_checks.py | 158 +++++++++++++++++- 4 files changed, 298 insertions(+), 14 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index b9ff051da7..ffd399ae83 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1254,8 +1254,8 @@ class LiteLLM_ObjectPermissionTable(LiteLLMPydanticObjectBase): """Represents a LiteLLM_ObjectPermissionTable record""" object_permission_id: str - mcp_servers: List[str] - vector_stores: List[str] + mcp_servers: Optional[List[str]] = [] + vector_stores: Optional[List[str]] = [] class LiteLLM_TeamTable(TeamBase): @@ -2364,6 +2364,21 @@ class ProxyErrorTypes(str, enum.Enum): Team member permission error """ + key_vector_store_access_denied = "key_vector_store_access_denied" + """ + Key does not have access to the vector store + """ + + team_vector_store_access_denied = "team_vector_store_access_denied" + """ + Team does not have access to the vector store + """ + + org_vector_store_access_denied = "org_vector_store_access_denied" + """ + Organization does not have access to the vector store + """ + @classmethod def get_model_access_error_type_for_object( cls, object_type: Literal["key", "user", "team", "org"] @@ -2380,6 +2395,20 @@ class ProxyErrorTypes(str, enum.Enum): elif object_type == "org": return cls.org_model_access_denied + @classmethod + def get_vector_store_access_error_type_for_object( + cls, object_type: Literal["key", "team", "org"] + ) -> "ProxyErrorTypes": + """ + Get the vector store access error type for object_type + """ + if object_type == "key": + return cls.key_vector_store_access_denied + elif object_type == "team": + return cls.team_vector_store_access_denied + elif object_type == "org": + return cls.org_vector_store_access_denied + DB_CONNECTION_ERROR_TYPES = ( httpx.ConnectError, diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 5021247145..2cb6064989 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -32,6 +32,7 @@ from litellm.proxy._types import ( LiteLLM_EndUserTable, Litellm_EntityType, LiteLLM_JWTAuth, + LiteLLM_ObjectPermissionTable, LiteLLM_OrganizationMembershipTable, LiteLLM_OrganizationTable, LiteLLM_TeamTable, @@ -94,6 +95,7 @@ async def common_checks( 8. [OPTIONAL] If guardrails modified - is request allowed to change this 9. Check if request body is safe 10. [OPTIONAL] Organization checks - is user_object.organization_id is set, run these checks + 11. [OPTIONAL] Vector store checks - is the object allowed to access the vector store """ _model: Optional[Union[str, List[str]]] = get_model_from_request( request_body, route @@ -219,6 +221,13 @@ async def common_checks( valid_token=valid_token, ) + # 11. [OPTIONAL] Vector store checks - is the object allowed to access the vector store + await vector_store_access_check( + request_body=request_body, + team_object=team_object, + valid_token=valid_token, + ) + return True @@ -1576,3 +1585,104 @@ def _is_wildcard_pattern(allowed_model_pattern: str) -> bool: Checks if `*` is in the pattern. """ return "*" in allowed_model_pattern + + +async def vector_store_access_check( + request_body: dict, + team_object: Optional[LiteLLM_TeamTable], + valid_token: Optional[UserAPIKeyAuth], +): + """ + Checks if the object (key, team, org) has access to the vector store. + + Raises ProxyException if the object (key, team, org) cannot access the specific vector store. + """ + from litellm.proxy.proxy_server import prisma_client + + ######################################################### + # Get the vector store the user is trying to access + ######################################################### + if prisma_client is None: + verbose_proxy_logger.debug( + "Prisma client not found, skipping vector store access check" + ) + return True + + if litellm.vector_store_registry is None: + verbose_proxy_logger.debug( + "Vector store registry not found, skipping vector store access check" + ) + return True + + vector_store_ids_to_run = litellm.vector_store_registry.get_vector_store_ids_to_run( + non_default_params=request_body, tools=request_body.get("tools", None) + ) + if vector_store_ids_to_run is None: + verbose_proxy_logger.debug( + "Vector store to run not found, skipping vector store access check" + ) + return True + + ######################################################### + # Check if the object (key, team, org) has access to the vector store + ######################################################### + # Check if the key can access the vector store + if valid_token is not None and valid_token.object_permission_id is not None: + key_object_permission = ( + await prisma_client.db.litellm_objectpermissiontable.find_unique( + where={"object_permission_id": valid_token.object_permission_id}, + ) + ) + if key_object_permission is not None: + _can_object_call_vector_stores( + object_type="key", + vector_store_ids_to_run=vector_store_ids_to_run, + object_permissions=key_object_permission, + ) + + # Check if the team can access the vector store + if team_object is not None and team_object.object_permission_id is not None: + team_object_permission = ( + await prisma_client.db.litellm_objectpermissiontable.find_unique( + where={"object_permission_id": team_object.object_permission_id}, + ) + ) + if team_object_permission is not None: + _can_object_call_vector_stores( + object_type="team", + vector_store_ids_to_run=vector_store_ids_to_run, + object_permissions=team_object_permission, + ) + return True + + +def _can_object_call_vector_stores( + object_type: Literal["key", "team", "org"], + vector_store_ids_to_run: List[str], + object_permissions: Optional[LiteLLM_ObjectPermissionTable], +): + """ + Raises ProxyException if the object (key, team, org) cannot access the specific vector store. + """ + if object_permissions is None: + return True + + if object_permissions.vector_stores is None: + return True + + # If length is 0, then the object has access to all vector stores. + if len(object_permissions.vector_stores) == 0: + return True + + for vector_store_id in vector_store_ids_to_run: + if vector_store_id not in object_permissions.vector_stores: + raise ProxyException( + message=f"User not allowed to access vector store. Tried to access {vector_store_id}. Only allowed to access {object_permissions.vector_stores}", + type=ProxyErrorTypes.get_vector_store_access_error_type_for_object( + object_type + ), + param="vector_store", + code=status.HTTP_401_UNAUTHORIZED, + ) + + return True diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index af3dc2b908..ded872e0ce 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -3,16 +3,5 @@ model_list: litellm_params: model: anthropic/* - general_settings: store_prompts_in_spend_logs: true - -guardrails: - - guardrail_name: "bedrock-pre-guard" - litellm_params: - guardrail: bedrock # supported values: "aporia", "bedrock", "lakera" - mode: "post_call" - guardrailIdentifier: wf0hkdb5x07f # your guardrail ID on bedrock - guardrailVersion: "DRAFT" # your guardrail version on bedrock - default_on: true - diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 97c952b63f..9b7af1859e 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -14,11 +14,21 @@ import pytest import litellm from litellm.proxy._types import ( + LiteLLM_ObjectPermissionTable, + LiteLLM_TeamTable, LiteLLM_UserTable, LitellmUserRoles, + ProxyErrorTypes, + ProxyException, SSOUserDefinedValues, + UserAPIKeyAuth, +) +from litellm.proxy.auth.auth_checks import ( + ExperimentalUIJWTToken, + _can_object_call_vector_stores, + get_user_object, + vector_store_access_check, ) -from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken, get_user_object from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper from litellm.utils import get_utc_datetime @@ -178,3 +188,149 @@ async def test_default_internal_user_params_with_get_user_object(monkeypatch): assert creation_args["models"] == ["gpt-4", "claude-3-opus"] assert creation_args["max_budget"] == 200.0 assert creation_args["user_role"] == "internal_user" + + +# Vector Store Auth Check Tests + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "prisma_client,vector_store_registry,expected_result", + [ + (None, MagicMock(), True), # No prisma client + (MagicMock(), None, True), # No vector store registry + (MagicMock(), MagicMock(), True), # No vector stores to run + ], +) +async def test_vector_store_access_check_early_returns( + prisma_client, vector_store_registry, expected_result +): + """Test vector_store_access_check returns True for early exit conditions""" + request_body = {"messages": [{"role": "user", "content": "test"}]} + + if vector_store_registry: + vector_store_registry.get_vector_store_ids_to_run.return_value = None + + with patch("litellm.proxy.proxy_server.prisma_client", prisma_client), patch( + "litellm.vector_store_registry", vector_store_registry + ): + result = await vector_store_access_check( + request_body=request_body, + team_object=None, + valid_token=None, + ) + + assert result == expected_result + + +@pytest.mark.parametrize( + "object_permissions,vector_store_ids,should_raise,error_type", + [ + (None, ["store-1"], False, None), # None permissions - should pass + ( + {"vector_stores": []}, + ["store-1"], + False, + None, + ), # Empty vector_stores - should pass (access to all) + ( + {"vector_stores": ["store-1", "store-2"]}, + ["store-1"], + False, + None, + ), # Has access + ( + {"vector_stores": ["store-1", "store-2"]}, + ["store-3"], + True, + ProxyErrorTypes.key_vector_store_access_denied, + ), # No access + ( + {"vector_stores": ["store-1"]}, + ["store-1", "store-3"], + True, + ProxyErrorTypes.team_vector_store_access_denied, + ), # Partial access + ], +) +def test_can_object_call_vector_stores_scenarios( + object_permissions, vector_store_ids, should_raise, error_type +): + """Test _can_object_call_vector_stores with various permission scenarios""" + # Convert dict to object if not None + if object_permissions is not None: + mock_permissions = MagicMock() + mock_permissions.vector_stores = object_permissions["vector_stores"] + object_permissions = mock_permissions + + object_type = ( + "key" + if error_type == ProxyErrorTypes.key_vector_store_access_denied + else "team" + ) + + if should_raise: + with pytest.raises(ProxyException) as exc_info: + _can_object_call_vector_stores( + object_type=object_type, + vector_store_ids_to_run=vector_store_ids, + object_permissions=object_permissions, + ) + assert exc_info.value.type == error_type + else: + result = _can_object_call_vector_stores( + object_type=object_type, + vector_store_ids_to_run=vector_store_ids, + object_permissions=object_permissions, + ) + assert result is True + + +@pytest.mark.asyncio +async def test_vector_store_access_check_with_permissions(): + """Test vector_store_access_check with actual permission checking""" + request_body = {"tools": [{"type": "function", "function": {"name": "test"}}]} + + # Test with valid token that has access + valid_token = UserAPIKeyAuth( + token="test-token", + object_permission_id="perm-123", + models=["gpt-4"], + max_budget=100.0, + ) + + mock_prisma_client = MagicMock() + mock_permissions = MagicMock() + mock_permissions.vector_stores = ["store-1", "store-2"] + mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock( + return_value=mock_permissions + ) + + mock_vector_store_registry = MagicMock() + mock_vector_store_registry.get_vector_store_ids_to_run.return_value = ["store-1"] + + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), patch( + "litellm.vector_store_registry", mock_vector_store_registry + ): + result = await vector_store_access_check( + request_body=request_body, + team_object=None, + valid_token=valid_token, + ) + + assert result is True + + # Test with denied access + mock_vector_store_registry.get_vector_store_ids_to_run.return_value = ["store-3"] + + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), patch( + "litellm.vector_store_registry", mock_vector_store_registry + ): + with pytest.raises(ProxyException) as exc_info: + await vector_store_access_check( + request_body=request_body, + team_object=None, + valid_token=valid_token, + ) + + assert exc_info.value.type == ProxyErrorTypes.key_vector_store_access_denied