diff --git a/litellm/proxy/hooks/batch_rate_limiter.py b/litellm/proxy/hooks/batch_rate_limiter.py index 17f8bd277dd..5f214816578 100644 --- a/litellm/proxy/hooks/batch_rate_limiter.py +++ b/litellm/proxy/hooks/batch_rate_limiter.py @@ -1128,6 +1128,24 @@ class _PROXY_BatchRateLimiter(CustomLogger): }, ) + @staticmethod + async def _check_managed_file_ownership( + managed_files_obj: object, + file_id: str, + user_api_key_dict: UserAPIKeyAuth, + ) -> None: + """Deny before the input file is read, so a caller cannot make the proxy + download, count, or charge its own limits for another tenant's file.""" + from litellm.proxy.openai_files_endpoints.common_utils import ManagedResourceAccessChecker + + if not isinstance(managed_files_obj, ManagedResourceAccessChecker): + raise HTTPException(status_code=500, detail="Managed resource ownership validation is unavailable.") + if not await managed_files_obj.can_user_call_unified_file_id(file_id, user_api_key_dict): + raise HTTPException( + status_code=403, + detail=f"User {user_api_key_dict.user_id} does not have access to the file {file_id}", + ) + async def _fetch_managed_file_content( self, file_id: str, @@ -1170,6 +1188,8 @@ class _PROXY_BatchRateLimiter(CustomLogger): if llm_router is None: raise ValueError("llm_router not available. Cannot access managed files.") + await self._check_managed_file_ownership(managed_files_obj, file_id, user_api_key_dict) + # Use the managed files hook to get file content # This properly handles user permissions and file ownership file_content: Final = await managed_files_obj.afile_content( diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 6e226868918..a6e9bb7b3d5 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -122,6 +122,7 @@ from litellm.proxy.management_helpers.access_group_key_sync import ( sync_key_update_access_group_membership, ) from litellm.proxy.management_helpers.key_settings_audit import with_settings_updated_at +from litellm.proxy.management_helpers.managed_resource_owner_sync import repoint_key_owned_managed_resources from litellm.proxy.management_helpers.object_permission_utils import ( # noqa: F401 # legacy module exports ObjectPermissionUpsert, _set_object_permission, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export @@ -5957,6 +5958,13 @@ async def _execute_virtual_key_regeneration( await evict_and_broadcast(cache_keys=jwt_mapping_cache_keys, user_api_key_cache=user_api_key_cache) + await repoint_key_owned_managed_resources( + prisma_client=prisma_client, + existing_key_row=key_in_db, + previous_key_token=hashed_api_key, + new_key_token=new_token_hash, + ) + # After credential invalidation, so a failure here can never keep the old key alive. await sync_key_regeneration_access_group_membership( prisma_client=prisma_client, diff --git a/litellm/proxy/management_helpers/managed_resource_owner_sync.py b/litellm/proxy/management_helpers/managed_resource_owner_sync.py new file mode 100644 index 00000000000..4171205b52b --- /dev/null +++ b/litellm/proxy/management_helpers/managed_resource_owner_sync.py @@ -0,0 +1,44 @@ +"""Carry key-owned managed resources over to a regenerated key's new hash.""" + +from __future__ import annotations + +from typing import Final, Protocol + +from litellm.proxy._types import LiteLLM_VerificationToken +from litellm.proxy.db.routing_prisma_wrapper import writer_wrapper +from litellm.repositories.table_repositories import ManagedObjectRepository + +_KEY_OWNED_TABLES: Final = ( + "LiteLLM_ManagedFileTable", + "LiteLLM_ManagedObjectTable", + "LiteLLM_ManagedVectorStoreTable", +) + + +class _RawWriter(Protocol): + async def execute_raw(self, query: str, *args: str) -> int: ... + + +def _raw_writer(prisma_client: object) -> _RawWriter: + db: Final = ManagedObjectRepository(prisma_client).prisma_client.db # pyright: ignore[reportAny] # untyped Prisma client + return writer_wrapper(db) # pyright: ignore[reportAny, reportReturnType] # untyped Prisma client behind the pin + + +async def repoint_key_owned_managed_resources( + prisma_client: object, + existing_key_row: LiteLLM_VerificationToken, + previous_key_token: str, + new_key_token: str, +) -> None: + """A key with neither a user_id nor a team_id owns its managed files, batches and + vector stores as ``key:`` (``resolve_resource_owner_id``), so a + regenerated key would otherwise be denied its own resources.""" + if existing_key_row.user_id is not None or existing_key_row.team_id is not None: + return + writer: Final = _raw_writer(prisma_client) + for table in _KEY_OWNED_TABLES: + await writer.execute_raw( + f'UPDATE "{table}" SET "created_by" = $2 WHERE "created_by" = $1', + f"key:{previous_key_token}", + f"key:{new_key_token}", + ) diff --git a/tests/integration/management/test_managed_file_create_ownership.py b/tests/integration/management/test_managed_file_create_ownership.py index 3cce34b892c..a14f1fced19 100644 --- a/tests/integration/management/test_managed_file_create_ownership.py +++ b/tests/integration/management/test_managed_file_create_ownership.py @@ -10,7 +10,7 @@ import httpx import pytest from integration._support.client import JSON_OBJECT, Gateway, Scenario, string_value from integration._support.upstream import ScenarioHandle, delete_scenario, register_scenario -from integration.cost_calculation.cost_tracking_case import JsonResponse, RoutedResponse +from integration.cost_calculation.cost_tracking_case import JsonResponse, RoutedResponse, TextResponse from pydantic import JsonValue _FILE: Final = { @@ -198,3 +198,74 @@ def test_fine_tuning_create_from_own_managed_files_still_dispatches(gateway: Gat ) assert created.status_code == 200, created.text assert _upstream_creates(gateway, tenants.handle, "/fine_tuning/jobs") == 1 + + +def _upstream_calls(gateway: Gateway, handle: ScenarioHandle, method: str, suffix: str) -> int: + response: Final = httpx.get(f"{gateway.upstream_url}/__observations", timeout=15, trust_env=False) + assert response.status_code == 200, response.text + requests: Final = JSON_OBJECT.validate_json(response.content)["requests"] + assert isinstance(requests, list), response.text + return sum( + 1 + for request in requests + if isinstance(request, dict) + and request.get("method") == method + and handle.scenario_id in str(request.get("path")) + and str(request.get("path")).endswith(suffix) + ) + + +def test_batch_create_from_another_teams_file_is_403_before_the_rate_limiter_reads_it(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + scenario_id: Final = f"managed-file-create-ownership-{uuid.uuid4().hex}" + handle: Final = register_scenario(scenario_id, _routes()) + scenario.cleanups.callback(delete_scenario, handle) + model: Final = scenario.model(api_base=handle.api_base()) + line: Final = json.dumps( + { + "custom_id": "req-1", + "method": "POST", + "url": "/v1/chat/completions", + "body": {"model": model, "messages": [{"role": "user", "content": "ping"}], "max_tokens": 5}, + } + ) + register_scenario( + scenario_id, + RoutedResponse( + content_type="application/x-routed", + routes={ + **_routes().routes, + "GET /files/file-in-$REQUEST_ID/content": TextResponse( + content_type="application/jsonl", body=line + "\n" + ), + }, + ), + ) + team_a: Final = scenario.team(models=[model]) + team_b: Final = scenario.team(models=[model]) + owner: Final = scenario.key(user_id=scenario.user(user_role="internal_user"), team_id=team_a) + outsider: Final = scenario.key(user_id=scenario.user(user_role="internal_user"), team_id=team_b, tpm_limit=1) + owner_file_id: Final = _upload(scenario, owner, model, "batch") + denied: Final = _create_batch(gateway, outsider, model, owner_file_id) + assert denied.status_code == 403, denied.text + assert owner_file_id in denied.text, denied.text + assert _upstream_calls(gateway, handle, "GET", "/content") == 0 + + +def test_regenerated_key_without_user_or_team_keeps_its_managed_files(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + tenants: Final = _tenants(scenario, "batch") + original: Final = string_value(gateway.post("/key/generate", {})["key"]) + file_id: Final = _upload(scenario, original, tenants.model, "batch") + regenerated: Final = string_value(gateway.post("/key/regenerate", {"key": original})["key"]) + scenario.cleanups.callback(scenario.delete_key, regenerated) + batch: Final = _create_batch(gateway, regenerated, tenants.model, file_id) + retrieved: Final = gateway.request("GET", f"/v1/files/{file_id}", key=regenerated) + stranger: Final = _create_batch( + gateway, string_value(gateway.post("/key/generate", {})["key"]), tenants.model, file_id + ) + assert (batch.status_code, retrieved.status_code, stranger.status_code) == (200, 200, 403), ( + batch.text, + retrieved.text, + stranger.text, + ) diff --git a/tests/unit/proxy/hooks/test_batch_rate_limiter.py b/tests/unit/proxy/hooks/test_batch_rate_limiter.py index 2f71b177705..b4e972ea34f 100644 --- a/tests/unit/proxy/hooks/test_batch_rate_limiter.py +++ b/tests/unit/proxy/hooks/test_batch_rate_limiter.py @@ -11,6 +11,7 @@ import time from collections.abc import Iterator, Sequence from datetime import datetime, timezone from typing import Final, Literal +from unittest.mock import AsyncMock, MagicMock import httpx import pytest @@ -481,10 +482,9 @@ async def test_batch_rate_limiter_without_user_context(openai_files: respx.MockR assert usage_without_context.total_tokens == usage_with_context.total_tokens == _token_counter_total(rows) -@pytest.mark.asyncio -async def test_managed_file_is_read_through_the_managed_files_hook_with_user_context( - openai_files: respx.MockRouter, monkeypatch: pytest.MonkeyPatch -): +async def _install_managed_file_hook( + openai_files: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, rows: str +) -> tuple[respx.Route, respx.Route]: from litellm_enterprise.proxy.hooks.managed_files import _PROXY_LiteLLMManagedFiles from litellm import Router @@ -494,24 +494,22 @@ async def test_managed_file_is_read_through_the_managed_files_hook_with_user_con from litellm.proxy.utils import ProxyLogging assert is_base64_encoded_unified_file_id(_MANAGED_FILE_ID) - rows: Final = _batch_rows(("Test message for regression",)) provider_route: Final = _serve_file(openai_files, "file-provider-1", rows) - standard_route: Final = _serve_file(openai_files, "file-abc123", rows) unrouted_managed_read: Final = _serve_file(openai_files, _MANAGED_FILE_ID, rows) - file_cache: Final = InternalUsageCache(dual_cache=DualCache()) - await file_cache.async_set_cache( - key=_MANAGED_FILE_ID, - value=LiteLLM_ManagedFileTable( - unified_file_id=_MANAGED_FILE_ID, - model_mappings={"deployment-1": "file-provider-1"}, - flat_model_file_ids=["file-provider-1"], - created_by="test-user-regression", - ).model_dump(), - litellm_parent_otel_span=None, + row: Final = LiteLLM_ManagedFileTable( + unified_file_id=_MANAGED_FILE_ID, + model_mappings={"deployment-1": "file-provider-1"}, + flat_model_file_ids=["file-provider-1"], + created_by="test-user-regression", + team_id="team-owner", ) + file_cache: Final = InternalUsageCache(dual_cache=DualCache()) + await file_cache.async_set_cache(key=_MANAGED_FILE_ID, value=row.model_dump(), litellm_parent_otel_span=None) + prisma_client: Final = MagicMock() + prisma_client.db.litellm_managedfiletable.find_first = AsyncMock(return_value=row) proxy_logging: Final = ProxyLogging(user_api_key_cache=DualCache()) proxy_logging.proxy_hook_mapping["managed_files"] = _PROXY_LiteLLMManagedFiles( - internal_usage_cache=file_cache, prisma_client=None + internal_usage_cache=file_cache, prisma_client=prisma_client ) router: Final = Router( model_list=[ @@ -524,6 +522,16 @@ async def test_managed_file_is_read_through_the_managed_files_hook_with_user_con ) monkeypatch.setattr(proxy_server, "proxy_logging_obj", proxy_logging) monkeypatch.setattr(proxy_server, "llm_router", router) + return provider_route, unrouted_managed_read + + +@pytest.mark.asyncio +async def test_managed_file_is_read_through_the_managed_files_hook_with_user_context( + openai_files: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + rows: Final = _batch_rows(("Test message for regression",)) + provider_route, unrouted_managed_read = await _install_managed_file_hook(openai_files, monkeypatch, rows) + standard_route: Final = _serve_file(openai_files, "file-abc123", rows) _, _, batch_limiter = _make_limiters() user_api_key_dict: Final = UserAPIKeyAuth( api_key="test-key-regression", user_id="test-user-regression", tpm_limit=1000, rpm_limit=10 @@ -541,3 +549,25 @@ async def test_managed_file_is_read_through_the_managed_files_hook_with_user_con assert not unrouted_managed_read.called assert managed_usage.request_count == standard_usage.request_count == 1 assert managed_usage.total_tokens == standard_usage.total_tokens == _token_counter_total(rows) + + +@pytest.mark.asyncio +async def test_managed_file_of_another_tenant_is_denied_before_it_is_read( + openai_files: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + rows: Final = _batch_rows(("Test message for regression",)) + provider_route, unrouted_managed_read = await _install_managed_file_hook(openai_files, monkeypatch, rows) + _, _, batch_limiter = _make_limiters() + outsider: Final = UserAPIKeyAuth(api_key="test-key-outsider", user_id="outsider", team_id="team-other", tpm_limit=1) + + with pytest.raises(HTTPException) as denied: + await batch_limiter.async_pre_call_hook( + user_api_key_dict=outsider, + cache=DualCache(), + data=dict(_create_batch_data(_MANAGED_FILE_ID)), + call_type="acreate_batch", + ) + + assert denied.value.status_code == 403 + assert not provider_route.called + assert not unrouted_managed_read.called