fix(proxy): cap the decode passes the raw cloud storage file id guard makes

Any authenticated key could send a file id wrapped in tens of thousands of percent-encoding layers and the guard would unwrap every one of them, each pass walking the whole id, so a 120 KB request line held a worker for about five seconds and a 2 MB one for minutes. The guard now unwraps at most eight layers, and an id still changing after that is treated as a raw cloud storage id, so it stays admin-only instead of being decoded any further

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._-]+")

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)

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):