diff --git a/litellm/litellm_core_utils/cloud_storage_security.py b/litellm/litellm_core_utils/cloud_storage_security.py index bce209d5f6f..cf03033bf3c 100644 --- a/litellm/litellm_core_utils/cloud_storage_security.py +++ b/litellm/litellm_core_utils/cloud_storage_security.py @@ -3,6 +3,7 @@ import posixpath import re from collections.abc import Callable, Iterator, Mapping, Sequence from dataclasses import dataclass +from functools import partial from itertools import accumulate, chain, repeat, takewhile from typing import Final from urllib.parse import quote, unquote @@ -59,6 +60,7 @@ _GEMINI_MEDIA_PART_KEYS: Final = frozenset({"file_data", "fileData"}) _MEDIA_URI_KEYS: Final = frozenset({"url", "image_url", "file_id", "file_url", "file_data", "file_uri", "fileUri"}) _ENCODED_FILE_ID_PREFIXES: Final = ("file-", "batch_") _EMBEDDED_S3_URI: Final = re.compile(r"s3://[^;]*") +_MODEL_FILE_ID_MAPPING_KEY: Final = "model_file_id_mapping" def is_s3_uri(value: str) -> bool: @@ -81,6 +83,10 @@ def _s3_uris_embedded_in_file_id(file_id: str) -> Iterator[str]: return chain.from_iterable(map(_s3_uris_in, filter(None, map(_base64url_text, encodings)))) +def _with_embedded_s3_uris(file_id: str) -> Iterator[str]: + return chain((file_id,), _s3_uris_embedded_in_file_id(file_id)) + + @dataclass(frozen=True, slots=True) class _BodyNode: value: object @@ -106,19 +112,32 @@ def _next_level(level: tuple[_BodyNode, ...], _: None) -> tuple[_BodyNode, ...]: return tuple(chain.from_iterable(map(_child_nodes, level))) -def _body_nodes(request_body: Mapping[str, object]) -> Iterator[_BodyNode]: +def _body_nodes(request_body: object) -> Iterator[_BodyNode]: levels: Final = accumulate(repeat(None), _next_level, initial=(_BodyNode(request_body),)) return chain.from_iterable(takewhile(bool, levels)) -def _media_uri_candidates(node: _BodyNode) -> Iterator[str]: +def _string_leaf(node: _BodyNode) -> str | None: + return node.value if isinstance(node.value, str) else None + + +def _media_file_id(node: _BodyNode) -> str | None: + return _string_leaf(node) if node.inside_media_part and node.key == "file_id" else None + + +def _media_uri_candidates(trusted_file_ids: frozenset[str], node: _BodyNode) -> Iterator[str]: if not node.inside_media_part or not isinstance(node.value, str): return iter(()) if node.key == "file_id": - return chain((node.value,), _s3_uris_embedded_in_file_id(node.value)) + return iter((node.value,)) if node.value in trusted_file_ids else _with_embedded_s3_uris(node.value) return iter((node.value,) if node.key in _MEDIA_URI_KEYS else ()) +def _mapped_provider_file_ids(request_body: Mapping[str, object]) -> Iterator[str]: + mapping_leaves: Final = filter(None, map(_string_leaf, _body_nodes(request_body.get(_MODEL_FILE_ID_MAPPING_KEY)))) + return chain.from_iterable(map(_with_embedded_s3_uris, mapping_leaves)) + + def _s3_bucket(uri: str) -> str | None: return uri.removeprefix(_S3_SCHEME).partition("/")[0] if uri.startswith(_S3_SCHEME) else None @@ -143,9 +162,17 @@ def _refused_in_some_decoding( return bool(parts) and (not settled or not all(map(is_allowed, parts))) -def s3_media_uris(request_body: Mapping[str, object]) -> tuple[str, ...]: - media_uris: Final = chain.from_iterable(map(_media_uri_candidates, _body_nodes(request_body))) - return tuple(dict.fromkeys(filter(is_s3_uri, media_uris))) +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 s3_media_uris( + request_body: Mapping[str, object], trusted_file_ids: frozenset[str] = frozenset() +) -> tuple[str, ...]: + media_uris: Final = chain.from_iterable( + map(partial(_media_uri_candidates, trusted_file_ids), _body_nodes(request_body)) + ) + return tuple(dict.fromkeys(filter(is_s3_uri, chain(media_uris, _mapped_provider_file_ids(request_body))))) def is_managed_s3_uri(uri: str) -> bool: diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 3e8f7935219..3e546385d61 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -8,7 +8,7 @@ from dataclasses import dataclass from datetime import datetime, timezone from functools import lru_cache, partial from http import HTTPStatus -from itertools import chain +from itertools import chain, compress from types import MappingProxyType from typing import ( TYPE_CHECKING, @@ -65,6 +65,7 @@ from litellm.litellm_core_utils.bug_report import ( ) from litellm.litellm_core_utils.cloud_storage_security import ( is_managed_s3_uri, + media_file_ids, names_an_unlisted_s3_bucket, s3_media_uris, ) @@ -154,6 +155,10 @@ from litellm.proxy.guardrails.auto_router_compression import arm_pre_call as _ar from litellm.proxy.hooks.parallel_request_limiter_v3 import get_or_create_request_stash from litellm.proxy.hooks.proxy_track_cost_callback import failure_row_response_cost from litellm.proxy.native_compaction import with_proxy_compaction_executor +from litellm.proxy.openai_files_endpoints.common_utils import ( + ManagedResourceAccessChecker, + is_base64_encoded_unified_file_id, +) from litellm.proxy.route_llm_request import ( endpoint_unsupported_by_model, route_request, @@ -1044,22 +1049,45 @@ def _request_model(data: Mapping[str, object]) -> str | list[str] | None: return None -def refuse_s3_media_the_key_may_not_read( +async def _key_may_read_unified_file_id( + managed_files: ManagedResourceAccessChecker, user_api_key_dict: UserAPIKeyAuth, unified_file_id: str +) -> bool: + try: + return await managed_files.can_user_call_unified_file_id(unified_file_id, user_api_key_dict) + except HTTPException: + return False + + +async def _unified_file_ids_the_key_may_read( + data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth, managed_files: object +) -> 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) + ) + return frozenset(compress(unified_file_ids, readable)) + + +async def refuse_s3_media_the_key_may_not_read( data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth, general_settings: Mapping[str, object], + managed_files: object, ) -> None: if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN: return - s3_uris: Final = s3_media_uris(data) - if not s3_uris: + if not s3_media_uris(data): return + readable_file_ids: Final = await _unified_file_ids_the_key_may_read(data, user_api_key_dict, managed_files) + s3_uris: Final = s3_media_uris(data, readable_file_ids) managed_uris: Final = tuple(filter(is_managed_s3_uri, s3_uris)) if managed_uris: raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, detail=f"Raw s3:// URIs under LiteLLM-managed storage can only be sent by a proxy admin key " - f"({', '.join(managed_uris)}). Use the LiteLLM managed file id returned when the file was created.", + 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 @@ -2687,8 +2715,11 @@ class ProxyBaseLLMRequestProcessing: param="model", code=status.HTTP_400_BAD_REQUEST, ) - refuse_s3_media_the_key_may_not_read( - raw_request_fields, user_api_key_dict, _REQUEST_FIELDS.validate_python(general_settings) + await refuse_s3_media_the_key_may_not_read( + raw_request_fields, + user_api_key_dict, + _REQUEST_FIELDS.validate_python(general_settings), + proxy_logging_obj.get_proxy_hook("managed_files"), ) self.data = await add_litellm_data_to_request( data=self.data, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 3e07b634401..a1079456c61 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -19278,7 +19278,9 @@ async def async_queue_request( data = await request.json() data.pop("_litellm_strip_stream_usage", None) data.pop(STREAM_USAGE_INJECTED_FIELD, None) - refuse_s3_media_the_key_may_not_read(data, user_api_key_dict, general_settings) + await refuse_s3_media_the_key_may_not_read( + data, user_api_key_dict, general_settings, proxy_logging_obj.get_proxy_hook("managed_files") + ) # Include original request and headers in the data data["proxy_server_request"] = { diff --git a/litellm/types/files.py b/litellm/types/files.py index 259a836d9ad..083f5bf77d1 100644 --- a/litellm/types/files.py +++ b/litellm/types/files.py @@ -29,6 +29,7 @@ class FileType(Enum): JSON = "JSON" M4A = "M4A" M4V = "M4V" + MD = "MD" MOV = "MOV" MP3 = "MP3" MP4 = "MP4" @@ -75,6 +76,7 @@ FILE_EXTENSIONS: Final[Mapping[FileType, list[str]]] = MappingProxyType( FileType.JSON: ["json"], FileType.M4A: ["m4a"], FileType.M4V: ["m4v"], + FileType.MD: ["md", "markdown"], FileType.MOV: ["mov"], FileType.MP3: ["mp3"], FileType.MP4: ["mp4"], @@ -122,6 +124,7 @@ FILE_MIME_TYPES: Final[Mapping[FileType, str]] = MappingProxyType( FileType.JSON: "application/json", FileType.M4A: "audio/x-m4a", FileType.M4V: "video/x-m4v", + FileType.MD: "text/markdown", FileType.MOV: "video/quicktime", FileType.MP3: "audio/mpeg", FileType.MP4: "video/mp4", @@ -238,7 +241,7 @@ def is_audio_file_type(file_type): # Text -TEXT_FILE_TYPES: Final = {FileType.CSV, FileType.HTML, FileType.RTF, FileType.TXT} +TEXT_FILE_TYPES: Final = {FileType.CSV, FileType.HTML, FileType.MD, FileType.RTF, FileType.TXT} def is_text_file_type(file_type): diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index daf4a0e3e71..40c35420439 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -3422,6 +3422,16 @@ async def test_bedrock_converse_s3_pdf_is_a_document_named_after_its_uri(run_asy assert invoice_content[0]["document"]["name"] != report_name +@pytest.mark.parametrize("run_async", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize("uri", ["s3://media-bucket/docs/notes.md", "s3://media-bucket/docs/NOTES.markdown"]) +async def test_bedrock_converse_s3_markdown_without_a_format_is_an_md_document(uri: str, run_async: bool): + content: Final = await _converse_user_content({"type": "file", "file": {"file_id": uri}}, run_async) + + assert [{**block, "document": {**block["document"], "name": "-"}} for block in content] == [ + {"document": {"format": "md", "name": "-", "source": _s3_source(uri)}} + ] + + @pytest.mark.parametrize("run_async", [False, True], ids=["sync", "async"]) @pytest.mark.parametrize("uri", ["s3://media-bucket/uploads/7f3a9c", "s3://media-bucket/uploads/blob.unknownext"]) async def test_bedrock_converse_s3_media_without_a_known_extension_or_format_is_a_bad_request( diff --git a/tests/unit/litellm_core_utils/test_cloud_storage_security.py b/tests/unit/litellm_core_utils/test_cloud_storage_security.py index d39e235517c..15648503815 100644 --- a/tests/unit/litellm_core_utils/test_cloud_storage_security.py +++ b/tests/unit/litellm_core_utils/test_cloud_storage_security.py @@ -10,6 +10,7 @@ from litellm.litellm_core_utils.cloud_storage_security import ( is_managed_cloud_storage_uri, is_managed_s3_uri, names_an_unlisted_s3_bucket, + media_file_ids, s3_media_uris, ) @@ -113,6 +114,20 @@ def _gemini_turn(*parts: object) -> dict[str, object]: _responses_turn({"type": "input_file", "file_id": _UNIFIED_BATCH_OUTPUT_ID}), id="responses-unified-file-id", ), + pytest.param( + { + **_user_turn({"type": "file", "file": {"file_id": "alias"}}), + "model_file_id_mapping": {"alias": {"deployment-1": _MANAGED_UPLOAD}}, + }, + id="model-file-id-mapping", + ), + pytest.param( + { + **_responses_turn({"type": "input_file", "file_id": "alias"}), + "model_file_id_mapping": {"alias": {"deployment-1": _MODEL_EMBEDDED_UPLOAD_ID}}, + }, + id="model-file-id-mapping-to-a-model-embedded-file-id", + ), pytest.param(_responses_turn({"type": "input_image", "image_url": _MANAGED_UPLOAD}), id="responses-image"), pytest.param(_responses_turn({"type": "input_file", "file_url": _MANAGED_UPLOAD}), id="responses-file-url"), pytest.param(_responses_turn({"type": "input_file", "file_id": _MANAGED_UPLOAD}), id="responses-file-id"), @@ -340,6 +355,33 @@ def test_s3_media_uris_ignores_what_never_becomes_an_s3_media_source(request_bod assert s3_media_uris(request_body) == () +def test_s3_media_uris_skips_only_the_objects_a_trusted_file_id_embeds(): + listed_image: Final = "s3://media-bucket/dir/cat.png" + request_body: Final = _responses_turn( + {"type": "input_file", "file_id": _UNIFIED_BATCH_OUTPUT_ID}, + {"type": "input_file", "file_id": _MODEL_EMBEDDED_UPLOAD_ID}, + {"type": "input_image", "image_url": listed_image}, + ) + + assert s3_media_uris(request_body, frozenset({_UNIFIED_BATCH_OUTPUT_ID})) == (_MANAGED_UPLOAD, listed_image) + assert s3_media_uris(request_body, frozenset({_UNIFIED_BATCH_OUTPUT_ID, _MODEL_EMBEDDED_UPLOAD_ID})) == ( + listed_image, + ) + + +def test_media_file_ids_lists_the_file_ids_of_media_parts_only(): + request_body: Final = { + **_user_turn( + {"type": "file", "file": {"file_id": "file-a"}}, + {"type": "text", "text": "file-b", "file_id": "file-b"}, + {"type": "file", "file": {"file_id": "file-a"}}, + ), + "metadata": {"file_id": "file-c"}, + } + + assert media_file_ids(request_body) == ("file-a",) + + def test_s3_media_uris_names_each_s3_object_once_in_the_order_it_first_appears(): other_upload: Final = "s3://media-bucket/dir/cat.png" request_body: Final = _user_turn( diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index f92a33b3996..19c96250c9a 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -4,6 +4,7 @@ import copy import datetime import gzip import json +from base64 import urlsafe_b64encode from collections.abc import AsyncIterable, Awaitable, Coroutine, Mapping from types import MappingProxyType, SimpleNamespace from typing import AsyncGenerator, Callable, Final, Iterator, Literal, Optional, Sequence, TypeAlias, TypeVar @@ -33,7 +34,7 @@ from litellm.litellm_core_utils.bug_report import ( build_bug_report, ) from litellm.types.router import RouterRateLimitError, RouterRateLimitErrorBasic -from litellm.types.utils import CachingDetails, Delta, ModelResponseStream, StreamingChoices +from litellm.types.utils import CachingDetails, Delta, ModelResponseStream, SpecialEnums, StreamingChoices from litellm.constants import ( CLIENT_REQUESTED_MODEL_SCOPE_KEY, MAX_LITELLM_CALL_ID_LENGTH, @@ -260,6 +261,48 @@ _S3_MEDIA_SURFACES: Final = ( ) +def _chat_file(file_id: str) -> dict[str, object]: + return {"messages": [{"role": "user", "content": [{"type": "file", "file": {"file_id": file_id}}]}]} + + +def _responses_file(file_id: str) -> dict[str, object]: + return {"input": [{"role": "user", "content": [{"type": "input_file", "file_id": file_id}]}]} + + +_S3_FILE_ID_SURFACES: Final = ( + pytest.param("acompletion", "/v1/chat/completions", _chat_file, id="chat-completions"), + pytest.param("aresponses", "/v1/responses", _responses_file, id="responses"), +) +_UPLOADED_OBJECT: Final = "s3://media-bucket/litellm-bedrock-files/upload.pdf" +_UPLOADED_FILE_ID: Final = ( + urlsafe_b64encode( + SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( + "application/pdf", "u1", "nova", _UPLOADED_OBJECT, "deployment-1" + ).encode() + ) + .decode() + .rstrip("=") +) + + +class _ManagedFilesCreatedBy(CustomLogger): + def __init__(self, creators: Mapping[str, str]) -> None: + super().__init__() + self.creators: Final = creators + + async def can_user_call_unified_file_id( + self, unified_file_id: str, user_api_key_dict: ProxyUserAPIKeyAuth + ) -> bool: + if unified_file_id not in self.creators: + raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=f"File not found: {unified_file_id}") + return self.creators[unified_file_id] == user_api_key_dict.user_id + + async def can_user_call_unified_object_id( + self, unified_object_id: str, user_api_key_dict: ProxyUserAPIKeyAuth + ) -> bool: + return False + + _S3_OBJECT: Final = "s3://media-bucket/dir/cat.png" _HELLO: Final = {"role": "user", "content": "hello"} _S3_VALUES_OUTSIDE_MEDIA_PARTS: Final = ( @@ -402,9 +445,14 @@ async def _send_s3_media( body: dict[str, object], user_role: LitellmUserRoles, general_settings: dict[str, object], + managed_files: CustomLogger | None = None, ) -> object: from litellm.proxy.proxy_server import ProxyConfig as RealProxyConfig + proxy_logging_obj: Final = ProxyLogging(user_api_key_cache=UserApiKeyCache()) + if managed_files is not None: + proxy_logging_obj.proxy_hook_mapping["managed_files"] = managed_files + router: Final = Router( model_list=[ { @@ -424,9 +472,9 @@ async def _send_s3_media( } ), fastapi_response=Response(), - user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test", user_role=user_role), + user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test", user_id="user-a", user_role=user_role), route_type=route_type, - proxy_logging_obj=ProxyLogging(user_api_key_cache=UserApiKeyCache()), + proxy_logging_obj=proxy_logging_obj, general_settings=general_settings, proxy_config=RealProxyConfig(), llm_router=router, @@ -820,6 +868,84 @@ class TestProxyBaseLLMRequestProcessing: assert {key: data[key] for key in body} == body + @pytest.mark.asyncio + @pytest.mark.parametrize(("route_type", "path", "body_for"), _S3_FILE_ID_SURFACES) + async def test_non_admin_managed_file_id_its_user_created_passes( + self, route_type: _S3MediaRoute, path: str, body_for: Callable[[str], dict[str, object]] + ) -> None: + response: Final = await _send_s3_media( + route_type, + path, + body_for(_UPLOADED_FILE_ID), + LitellmUserRoles.INTERNAL_USER, + {}, + _ManagedFilesCreatedBy({_UPLOADED_FILE_ID: "user-a"}), + ) + + assert "a cat" in str(response) + + @pytest.mark.asyncio + @pytest.mark.parametrize(("route_type", "path", "body_for"), _S3_FILE_ID_SURFACES) + @pytest.mark.parametrize( + "managed_files", + [ + pytest.param(_ManagedFilesCreatedBy({_UPLOADED_FILE_ID: "user-b"}), id="another-user-created-it"), + pytest.param(_ManagedFilesCreatedBy({}), id="forged-id"), + pytest.param(None, id="no-managed-files-hook"), + ], + ) + async def test_non_admin_managed_file_id_the_key_may_not_read_is_refused( + self, + route_type: _S3MediaRoute, + path: str, + body_for: Callable[[str], dict[str, object]], + managed_files: CustomLogger | None, + ) -> None: + with pytest.raises(HTTPException) as refused: + await _send_s3_media( + route_type, path, body_for(_UPLOADED_FILE_ID), LitellmUserRoles.INTERNAL_USER, {}, managed_files + ) + + assert refused.value.status_code == status.HTTP_403_FORBIDDEN, refused.value.detail + assert _UPLOADED_OBJECT in refused.value.detail + assert "LiteLLM-managed storage" in refused.value.detail + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("mapped_source", "named_in_refusal"), + [ + pytest.param( + "s3://other-bucket/private.pdf", + "general_settings.allowed_client_side_s3_buckets", + id="unlisted-bucket", + ), + pytest.param( + "s3://media-bucket/litellm-bedrock-files/other-team.pdf", "LiteLLM-managed storage", id="managed-key" + ), + ], + ) + async def test_non_admin_s3_source_behind_a_model_file_id_mapping_is_refused_while_a_proxy_admin_key_passes( + self, mapped_source: str, named_in_refusal: str + ) -> None: + body: Final = {**_chat_file("alias"), "model_file_id_mapping": {"alias": {"deployment-1": mapped_source}}} + + with pytest.raises(HTTPException) as refused: + await _send_s3_media( + "acompletion", + "/v1/chat/completions", + body, + LitellmUserRoles.INTERNAL_USER, + {"allowed_client_side_s3_buckets": ["media-bucket"]}, + ) + admin_response: Final = await _send_s3_media( + "acompletion", "/v1/chat/completions", body, LitellmUserRoles.PROXY_ADMIN, {} + ) + + assert refused.value.status_code == status.HTTP_403_FORBIDDEN, refused.value.detail + assert mapped_source in refused.value.detail + assert named_in_refusal in refused.value.detail + assert "a cat" in str(admin_response) + @pytest.mark.asyncio async def test_non_admin_s3_media_refusal_names_a_repeated_object_once(self) -> None: with pytest.raises(HTTPException) as refused: