From 9faf74ef8ee45224bbdf4c8718ef4536e499ec45 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 10 Nov 2026 17:32:29 -0800 Subject: [PATCH] fix(proxy): look up managed file ids only when they can change the s3 refusal (cherry picked from commit 0043e3020f5258d6542cad2031f70b2928b89101) --- .../cloud_storage_security.py | 4 + litellm/proxy/auth/auth_utils.py | 10 +- litellm/proxy/common_request_processing.py | 54 ++++++++--- .../proxy/test_common_request_processing.py | 94 +++++++++++++++++++ 4 files changed, 145 insertions(+), 17 deletions(-) diff --git a/litellm/litellm_core_utils/cloud_storage_security.py b/litellm/litellm_core_utils/cloud_storage_security.py index edf54fd375f..b88ce55ffa3 100644 --- a/litellm/litellm_core_utils/cloud_storage_security.py +++ b/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, ...]: diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 96450b7fd72..b2467457eb4 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/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, diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 4c87390bb5c..b7b80858ba2 100644 --- a/litellm/proxy/common_request_processing.py +++ b/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, diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 2922fca970c..27d56ade859 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/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"),