fix(proxy): look up managed file ids only when they can change the s3 refusal

(cherry picked from commit 0043e3020f)

litellm/litellm_core_utils/cloud_storage_security.py

@@ -177,6 +177,10 @@ def media_file_ids(request_body: Mapping[str, object]) -> tuple[str, ...]:
     return tuple(dict.fromkeys(filter(None, map(_media_file_id, _body_nodes(request_body)))))
 
 
+def file_id_embeds_s3_uri_matching(file_id: str, matches: Callable[[str], bool]) -> bool:
+    return any(map(matches, _s3_uris_embedded_in_file_id(file_id)))
+
+
 def s3_media_uris(
     request_body: Mapping[str, object], trusted_file_ids: frozenset[str] = frozenset()
 ) -> tuple[str, ...]:

litellm/proxy/auth/auth_utils.py

@@ -2,8 +2,8 @@ import importlib.util
 import os
 import re
 import sys
-from collections.abc import Collection, Iterator, Mapping
-from functools import lru_cache
+from collections.abc import Callable, Collection, Iterator, Mapping
+from functools import lru_cache, partial
 from itertools import accumulate, chain, repeat
 from logging import Logger
 from types import MappingProxyType
@@ -532,7 +532,7 @@ _BANNED_METADATA_PARAMS: Final[tuple[str, ...]] = tuple(
 )
 
 
-class ClientSideStorageAllowlist(BaseModel):
+class _ClientSideStorageAllowlist(BaseModel):
     model_config = ConfigDict(frozen=True)
 
     allowed_client_side_s3_buckets: tuple[str, ...] | None = None
@@ -541,6 +541,10 @@ class ClientSideStorageAllowlist(BaseModel):
         return bucket in (self.allowed_client_side_s3_buckets or ())
 
 
+def client_side_s3_bucket_check(general_settings: Mapping[str, object]) -> Callable[[str], bool]:
+    return _ClientSideStorageAllowlist.model_validate(general_settings).admits_s3_bucket
+
+
 def _check_banned_params(
     body: Mapping[str, object],
     general_settings: dict,

litellm/proxy/common_request_processing.py

@@ -3,10 +3,10 @@ import contextlib
 import json
 import logging
 import math
-from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Coroutine, Mapping, Sequence
+from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Coroutine, Iterable, Mapping, Sequence
 from datetime import datetime
-from functools import lru_cache
-from itertools import chain, compress
+from functools import lru_cache, partial
+from itertools import chain
 from types import MappingProxyType
 from typing import (
     TYPE_CHECKING,
@@ -59,6 +59,7 @@ from litellm.litellm_core_utils.bug_report import (
     strip_bug_report_notice,
 )
 from litellm.litellm_core_utils.cloud_storage_security import (
+    file_id_embeds_s3_uri_matching,
     is_managed_s3_uri,
     media_file_ids,
     names_an_unlisted_s3_bucket,
@@ -100,7 +101,7 @@ from litellm.proxy.auth.auth_checks import (
     request_skips_budget_checks,
     tag_max_budget_check_for_tags,
 )
-from litellm.proxy.auth.auth_utils import ClientSideStorageAllowlist, check_response_size_is_safe, get_request_route
+from litellm.proxy.auth.auth_utils import check_response_size_is_safe, client_side_s3_bucket_check, get_request_route
 from litellm.proxy.bug_report_config import build_proxy_bug_report
 from litellm.proxy.common_utils.callback_utils import (
     get_logging_caching_headers,
@@ -874,16 +875,41 @@ async def _key_may_read_unified_file_id(
         return False
 
 
+async def _leading_readable_file_ids(
+    managed_files: ManagedResourceAccessChecker, user_api_key_dict: UserAPIKeyAuth, unified_file_ids: Iterable[str]
+) -> AsyncIterator[str]:
+    for file_id in unified_file_ids:
+        if not await _key_may_read_unified_file_id(managed_files, user_api_key_dict, file_id):
+            return
+        yield file_id
+
+
 async def _unified_file_ids_the_key_may_read(
-    data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth, managed_files: object
+    data: Mapping[str, object],
+    user_api_key_dict: UserAPIKeyAuth,
+    managed_files: object,
+    is_refused: Callable[[str], bool],
 ) -> frozenset[str]:
     if not isinstance(managed_files, ManagedResourceAccessChecker):
         return frozenset()
-    unified_file_ids: Final = tuple(filter(is_base64_encoded_unified_file_id, media_file_ids(data)))
-    readable: Final = await asyncio.gather(
-        *(_key_may_read_unified_file_id(managed_files, user_api_key_dict, file_id) for file_id in unified_file_ids)
+    file_ids_hiding_a_refused_uri: Final = filter(
+        partial(file_id_embeds_s3_uri_matching, matches=is_refused),
+        filter(is_base64_encoded_unified_file_id, media_file_ids(data)),
     )
-    return frozenset(compress(unified_file_ids, readable))
+    return frozenset(
+        [
+            file_id
+            async for file_id in _leading_readable_file_ids(
+                managed_files, user_api_key_dict, file_ids_hiding_a_refused_uri
+            )
+        ]
+    )
+
+
+def _unlisted_s3_uri_check(general_settings: Mapping[str, object]) -> Callable[[str], bool]:
+    if general_settings.get("allow_client_side_credentials") is True:
+        return lambda _: False
+    return partial(names_an_unlisted_s3_bucket, is_listed=client_side_s3_bucket_check(general_settings))
 
 
 async def refuse_s3_media_the_key_may_not_read(
@@ -894,7 +920,10 @@ async def refuse_s3_media_the_key_may_not_read(
 ) -> None:
     if not s3_media_uris(data) or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN:
         return
-    readable_file_ids: Final = await _unified_file_ids_the_key_may_read(data, user_api_key_dict, managed_files)
+    is_unlisted: Final = _unlisted_s3_uri_check(general_settings)
+    readable_file_ids: Final = await _unified_file_ids_the_key_may_read(
+        data, user_api_key_dict, managed_files, lambda uri: is_managed_s3_uri(uri) or is_unlisted(uri)
+    )
     s3_uris: Final = s3_media_uris(data, readable_file_ids)
     managed_uris: Final = tuple(filter(is_managed_s3_uri, s3_uris))
     if managed_uris:
@@ -903,10 +932,7 @@ async def refuse_s3_media_the_key_may_not_read(
             detail=f"Raw s3:// URIs under LiteLLM-managed storage can only be sent by a proxy admin key "
             f"({', '.join(managed_uris)}). Use a LiteLLM managed file id that this key's user or team created.",
         )
-    if general_settings.get("allow_client_side_credentials") is True:
-        return
-    is_listed: Final = ClientSideStorageAllowlist.model_validate(general_settings).admits_s3_bucket
-    unlisted_uris: Final = tuple(uri for uri in s3_uris if names_an_unlisted_s3_bucket(uri, is_listed))
+    unlisted_uris: Final = tuple(filter(is_unlisted, s3_uris))
     if unlisted_uris:
         raise HTTPException(
             status_code=status.HTTP_403_FORBIDDEN,

tests/unit/proxy/test_common_request_processing.py

@@ -283,6 +283,30 @@ class _ManagedFilesCreatedBy(CustomLogger):
         return False
 
 
+class _ManagedFilesCountingLookups(_ManagedFilesCreatedBy):
+    def __init__(self, creators: Mapping[str, str]) -> None:
+        super().__init__(creators)
+        self.looked_up: Final[list[str]] = []  # mutable-ok: the fake records the lookups the gate makes
+
+    async def can_user_call_unified_file_id(
+        self, unified_file_id: str, user_api_key_dict: ProxyUserAPIKeyAuth
+    ) -> bool:
+        self.looked_up.append(unified_file_id)
+        return await super().can_user_call_unified_file_id(unified_file_id, user_api_key_dict)
+
+
+def _unified_file_id_for(s3_uri: str) -> str:
+    managed_id: Final = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format(
+        "application/pdf", "u1", "nova", s3_uri, "deployment-1"
+    )
+    return urlsafe_b64encode(managed_id.encode()).decode().rstrip("=")
+
+
+def _chat_files(file_ids: Sequence[str]) -> dict[str, object]:
+    parts: Final = [{"type": "file", "file": {"file_id": file_id}} for file_id in file_ids]
+    return {"messages": [{"role": "user", "content": parts}]}
+
+
 _S3_OBJECT: Final = "s3://media-bucket/dir/cat.png"
 _HELLO: Final = {"role": "user", "content": "hello"}
 _S3_VALUES_OUTSIDE_MEDIA_PARTS: Final = (
@@ -941,6 +965,76 @@ class TestProxyBaseLLMRequestProcessing:
         assert _UPLOADED_OBJECT in refused.value.detail
         assert "LiteLLM-managed storage" in refused.value.detail
 
+    @pytest.mark.asyncio
+    async def test_non_admin_forged_managed_file_ids_stop_at_the_first_failed_lookup(self) -> None:
+        forged_ids: Final = tuple(
+            _unified_file_id_for(f"s3://media-bucket/litellm-bedrock-files/forged-{index}.pdf") for index in range(2000)
+        )
+        managed_files: Final = _ManagedFilesCountingLookups({})
+
+        with pytest.raises(HTTPException) as refused:
+            await _send_s3_media(
+                "acompletion",
+                "/v1/chat/completions",
+                _chat_files(forged_ids),
+                LitellmUserRoles.INTERNAL_USER,
+                {},
+                managed_files,
+            )
+
+        assert refused.value.status_code == status.HTTP_403_FORBIDDEN, refused.value.detail
+        assert managed_files.looked_up == [forged_ids[0]]
+
+    @pytest.mark.asyncio
+    async def test_non_admin_file_ids_that_cannot_change_the_outcome_cost_no_lookup(self) -> None:
+        listed_object_id: Final = _unified_file_id_for("s3://media-bucket/dir/report.pdf")
+        id_without_s3_uri: Final = _unified_file_id_for("gs://other/report.pdf")
+        managed_files: Final = _ManagedFilesCountingLookups({_UPLOADED_FILE_ID: "user-a"})
+
+        response: Final = await _send_s3_media(
+            "acompletion",
+            "/v1/chat/completions",
+            _chat_files((listed_object_id, id_without_s3_uri, _UPLOADED_FILE_ID)),
+            LitellmUserRoles.INTERNAL_USER,
+            {"allowed_client_side_s3_buckets": ["media-bucket"]},
+            managed_files,
+        )
+
+        assert "a cat" in str(response)
+        assert managed_files.looked_up == [_UPLOADED_FILE_ID]
+
+    @pytest.mark.asyncio
+    async def test_non_admin_managed_file_id_its_user_created_passes_from_an_unlisted_bucket(self) -> None:
+        owned_id: Final = _unified_file_id_for("s3://other-bucket/report.pdf")
+
+        response: Final = await _send_s3_media(
+            "acompletion",
+            "/v1/chat/completions",
+            _chat_files((owned_id,)),
+            LitellmUserRoles.INTERNAL_USER,
+            {},
+            _ManagedFilesCreatedBy({owned_id: "user-a"}),
+        )
+
+        assert "a cat" in str(response)
+
+    @pytest.mark.asyncio
+    async def test_non_admin_unlisted_file_id_costs_no_lookup_when_client_side_credentials_are_allowed(self) -> None:
+        unlisted_object_id: Final = _unified_file_id_for("s3://other-bucket/report.pdf")
+        managed_files: Final = _ManagedFilesCountingLookups({})
+
+        response: Final = await _send_s3_media(
+            "acompletion",
+            "/v1/chat/completions",
+            _chat_files((unlisted_object_id,)),
+            LitellmUserRoles.INTERNAL_USER,
+            {"allow_client_side_credentials": True},
+            managed_files,
+        )
+
+        assert "a cat" in str(response)
+        assert managed_files.looked_up == []
+
     @pytest.mark.asyncio
     @pytest.mark.parametrize(
         ("mapped_source", "named_in_refusal"),