diff --git a/litellm/litellm_core_utils/cloud_storage_security.py b/litellm/litellm_core_utils/cloud_storage_security.py index 2d7eaac676f..87c586459a5 100644 --- a/litellm/litellm_core_utils/cloud_storage_security.py +++ b/litellm/litellm_core_utils/cloud_storage_security.py @@ -1,7 +1,7 @@ import posixpath import re from collections.abc import Mapping, Sequence -from itertools import accumulate, pairwise, repeat +from itertools import accumulate, repeat from types import MappingProxyType from typing import Any, Final, cast from urllib.parse import quote, unquote @@ -21,21 +21,29 @@ MANAGED_CLOUD_STORAGE_SCHEMES: Final = ("s3://", "gs://") _MAPPING_PROXY_TYPE: Final[type] = type(MappingProxyType({})) -def _fully_unquoted(value: str) -> str: - decodings: Final = accumulate(repeat(value), lambda current, _: unquote(current)) - return next(current for current, following in pairwise(decodings) if current == following) +MAX_FILE_ID_DECODE_PASSES: Final = 8 + + +def _decodings(value: str) -> tuple[str, ...]: + return tuple(accumulate(repeat(value, MAX_FILE_ID_DECODE_PASSES + 2), lambda current, _: unquote(current))) def is_managed_cloud_storage_uri(file_id: str) -> bool: """ True if file_id is a raw cloud-storage object URI (e.g. ``s3://bucket/key``), - however many times it was percent-encoded on the way in. + up to ``MAX_FILE_ID_DECODE_PASSES`` layers of percent-encoding deep, or an id + encoded deeper than that, which no client produces and which is treated as raw + rather than decoded any further. These are internal provider artifacts. On the multi-tenant proxy they must be retrieved through their managed unified file id so owner/team access is enforced; a raw URI supplied by a caller bypasses that check. """ - return isinstance(file_id, str) and _fully_unquoted(file_id).startswith(MANAGED_CLOUD_STORAGE_SCHEMES) + if not isinstance(file_id, str): + return False + decodings: Final = _decodings(file_id) + settled: Final = decodings[-1] == decodings[-2] + return decodings[-1].startswith(MANAGED_CLOUD_STORAGE_SCHEMES) or not settled _SAFE_OBJECT_COMPONENT_PATTERN: Final = re.compile(r"[^A-Za-z0-9._-]+") diff --git a/tests/test_litellm/litellm_core_utils/test_cloud_storage_security.py b/tests/test_litellm/litellm_core_utils/test_cloud_storage_security.py index db1dba12965..ef012361f0c 100644 --- a/tests/test_litellm/litellm_core_utils/test_cloud_storage_security.py +++ b/tests/test_litellm/litellm_core_utils/test_cloud_storage_security.py @@ -1,6 +1,10 @@ +from functools import reduce +from urllib.parse import quote + import pytest from litellm.litellm_core_utils.cloud_storage_security import ( + MAX_FILE_ID_DECODE_PASSES, is_managed_cloud_storage_uri, ) @@ -30,9 +34,19 @@ def test_is_managed_cloud_storage_uri_sees_through_percent_encoding(file_id: str assert is_managed_cloud_storage_uri(file_id) -def test_is_managed_cloud_storage_uri_survives_an_id_encoded_thousands_of_times(): +def _quoted_times(value: str, times: int) -> str: + return reduce(lambda current, _: quote(current, safe=""), range(times), value) + + +def test_is_managed_cloud_storage_uri_decodes_up_to_the_pass_cap(): + assert is_managed_cloud_storage_uri(_quoted_times("gs://bucket/x", MAX_FILE_ID_DECODE_PASSES)) + assert not is_managed_cloud_storage_uri(_quoted_times("file-abc/x", MAX_FILE_ID_DECODE_PASSES)) + + +def test_is_managed_cloud_storage_uri_fails_closed_on_an_id_encoded_past_the_pass_cap(): nested_gs_id = "gs" + "%" + "25" * 5000 + "3A//bucket/x" nested_plain_id = "%" + "25" * 5000 + assert is_managed_cloud_storage_uri(_quoted_times("file-abc/x", MAX_FILE_ID_DECODE_PASSES + 1)) assert is_managed_cloud_storage_uri(nested_gs_id) - assert not is_managed_cloud_storage_uri(nested_plain_id) + assert is_managed_cloud_storage_uri(nested_plain_id) diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py index bdf614528b7..386c43ae435 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py @@ -12,7 +12,7 @@ from pytest_mock import MockerFixture import litellm from litellm import Router from litellm.files.types import FileContentStreamingResult -from litellm.proxy._types import LiteLLM_UserTableFiltered, UserAPIKeyAuth +from litellm.proxy._types import LiteLLM_UserTableFiltered, LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.hooks import get_proxy_hook from litellm.proxy.management_endpoints.internal_user_endpoints import ui_view_users from litellm.proxy.openai_files_endpoints.file_content_streaming_handler import ( @@ -252,7 +252,9 @@ def _raw_cloud_file_ids_in_every_encoding() -> tuple[tuple[str, str, str], ...]: ) -def _override_auth(monkeypatch, mocker: MockerFixture, llm_router: Router, user_role) -> ProxyLogging: +def _override_auth( + monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture, llm_router: Router, user_role: LitellmUserRoles +) -> ProxyLogging: import litellm.proxy.proxy_server as ps proxy_logging_obj = setup_proxy_logging_object(monkeypatch, llm_router) @@ -269,10 +271,14 @@ def _override_auth(monkeypatch, mocker: MockerFixture, llm_router: Router, user_ @pytest.mark.parametrize(("encoded_id", "raw_id", "provider"), _raw_cloud_file_ids_in_every_encoding()) def test_get_file_content_answers_403_for_a_raw_cloud_id_from_a_non_admin_key( - mocker: MockerFixture, monkeypatch, llm_router: Router, encoded_id: str, raw_id: str, provider: str + mocker: MockerFixture, + monkeypatch: pytest.MonkeyPatch, + llm_router: Router, + encoded_id: str, + raw_id: str, + provider: str, ): import litellm.proxy.proxy_server as ps - from litellm.proxy._types import LitellmUserRoles _override_auth(monkeypatch, mocker, llm_router, LitellmUserRoles.INTERNAL_USER) afile_content = mocker.AsyncMock() @@ -293,25 +299,26 @@ def test_get_file_content_answers_403_for_a_raw_cloud_id_from_a_non_admin_key( @pytest.mark.parametrize(("encoded_id", "raw_id", "provider"), _raw_cloud_file_ids_in_every_encoding()) def test_get_file_content_forwards_a_raw_cloud_id_from_a_proxy_admin_key( - mocker: MockerFixture, monkeypatch, llm_router: Router, encoded_id: str, raw_id: str, provider: str + mocker: MockerFixture, + monkeypatch: pytest.MonkeyPatch, + llm_router: Router, + encoded_id: str, + raw_id: str, + provider: str, ): import litellm.proxy.proxy_server as ps - from litellm.proxy._types import LitellmUserRoles _override_auth(monkeypatch, mocker, llm_router, LitellmUserRoles.PROXY_ADMIN) - captured_kwargs: dict = {} - - async def _mock_afile_content(**kwargs): - captured_kwargs.update(kwargs) - return HttpxBinaryResponseContent( + afile_content = mocker.AsyncMock( + return_value=HttpxBinaryResponseContent( response=httpx.Response( status_code=200, content=b'{"custom_id": "request-1"}\n', headers={"content-type": "application/octet-stream"}, ) ) - - monkeypatch.setattr(litellm, "afile_content", _mock_afile_content) + ) + monkeypatch.setattr(litellm, "afile_content", afile_content) try: response = client.get( @@ -323,8 +330,9 @@ def test_get_file_content_forwards_a_raw_cloud_id_from_a_proxy_admin_key( assert response.status_code == 200, response.text assert response.content == b'{"custom_id": "request-1"}\n' - assert captured_kwargs["custom_llm_provider"] == provider - assert captured_kwargs["file_id"] == raw_id + afile_content.assert_called_once() + assert afile_content.call_args.kwargs["custom_llm_provider"] == provider + assert afile_content.call_args.kwargs["file_id"] == raw_id def test_mock_create_audio_file(mocker: MockerFixture, monkeypatch, llm_router: Router):