mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
Merge a53109335f into 1c61c2606e
This commit is contained in:
commit
fdb94637c7
4 changed files with 155 additions and 22 deletions
|
|
@ -1,6 +1,7 @@
|
|||
import posixpath
|
||||
import re
|
||||
from collections.abc import Mapping, Sequence
|
||||
from itertools import accumulate, repeat
|
||||
from types import MappingProxyType
|
||||
from typing import Any, Final, cast
|
||||
from urllib.parse import quote, unquote
|
||||
|
|
@ -20,15 +21,33 @@ MANAGED_CLOUD_STORAGE_SCHEMES: Final = ("s3://", "gs://")
|
|||
_MAPPING_PROXY_TYPE: Final[type] = type(MappingProxyType({}))
|
||||
|
||||
|
||||
MAX_FILE_ID_DECODE_PASSES: Final = 8
|
||||
|
||||
|
||||
def _unquote_once(current: str, _: str) -> str:
|
||||
return unquote(current)
|
||||
|
||||
|
||||
def _decodings(value: str) -> tuple[str, ...]:
|
||||
return tuple(accumulate(repeat(value, MAX_FILE_ID_DECODE_PASSES + 2), _unquote_once))
|
||||
|
||||
|
||||
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``).
|
||||
True if file_id is a raw cloud-storage object URI (e.g. ``s3://bucket/key``),
|
||||
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 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._-]+")
|
||||
|
|
|
|||
|
|
@ -892,14 +892,10 @@ async def get_file_content(
|
|||
}
|
||||
)
|
||||
else:
|
||||
# A raw cloud-storage URI (s3://, gs://) supplied here would skip the
|
||||
# managed-file owner/team check that only runs for unified ids, letting
|
||||
# a caller read another tenant's object by its key. Such objects are only
|
||||
# reachable through their managed unified id.
|
||||
if is_managed_cloud_storage_uri(file_id):
|
||||
if is_managed_cloud_storage_uri(file_id) and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="Raw cloud storage file ids cannot be retrieved directly. Use the LiteLLM managed file id returned when the file was created.",
|
||||
status_code=403,
|
||||
detail="Raw cloud storage file ids can only be retrieved by a proxy admin key. Use the LiteLLM managed file id returned when the file was created.",
|
||||
)
|
||||
# Check for model-based credential routing
|
||||
(
|
||||
|
|
|
|||
|
|
@ -1,4 +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,
|
||||
)
|
||||
|
||||
|
|
@ -13,3 +19,34 @@ def test_is_managed_cloud_storage_uri_ignores_provider_and_unified_ids():
|
|||
assert not is_managed_cloud_storage_uri("file-abc123")
|
||||
assert not is_managed_cloud_storage_uri("bGl0ZWxsbV9wcm94eQ==")
|
||||
assert not is_managed_cloud_storage_uri("")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"file_id",
|
||||
(
|
||||
"gs%3A%2F%2Fbucket%2Flitellm-vertex-files%2Fprediction-model%2Fpredictions.jsonl",
|
||||
"gs%253A%252F%252Fbucket%252Flitellm-vertex-files%252Fprediction-model%252Fpredictions.jsonl",
|
||||
"s3%3A%2F%2Fbucket%2Flitellm-batch-outputs%2Fx.jsonl.out",
|
||||
"s3%253A%252F%252Fbucket%252Flitellm-batch-outputs%252Fx.jsonl.out",
|
||||
),
|
||||
)
|
||||
def test_is_managed_cloud_storage_uri_sees_through_percent_encoding(file_id: str):
|
||||
assert is_managed_cloud_storage_uri(file_id)
|
||||
|
||||
|
||||
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 is_managed_cloud_storage_uri(nested_plain_id)
|
||||
|
|
|
|||
|
|
@ -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 (
|
||||
|
|
@ -235,23 +235,104 @@ def test_invalid_purpose(mocker: MockerFixture, monkeypatch, llm_router: Router)
|
|||
assert error["param"] == "purpose"
|
||||
|
||||
|
||||
def test_get_file_content_rejects_raw_cloud_storage_uri(llm_router: Router):
|
||||
"""A raw s3:// file id must be rejected on the proxy content endpoint.
|
||||
RAW_GCS_OUTPUT_FILE_ID: Final = (
|
||||
"gs://my-gcs-bucket/litellm-vertex-files/publishers/google/models/gemini-3.8-flash/"
|
||||
"prediction-model-2026-09-10T20:27:18.178864Z/predictions.jsonl"
|
||||
)
|
||||
RAW_S3_OUTPUT_FILE_ID: Final = "s3://my-bucket/litellm-batch-outputs/job-123/input.jsonl.out"
|
||||
|
||||
Such an id is not a managed unified id, so it would otherwise skip the
|
||||
owner/team access check and let a caller read another tenant's batch output
|
||||
object by its key. Callers must use the managed unified file id.
|
||||
"""
|
||||
|
||||
def _raw_cloud_file_ids_in_every_encoding() -> tuple[tuple[str, str, str], ...]:
|
||||
from urllib.parse import quote
|
||||
|
||||
s3_file_id = "s3://my-bucket/litellm-batch-outputs/job-123/input.jsonl.out"
|
||||
response = client.get(
|
||||
f"/v1/files/{quote(s3_file_id, safe='')}/content?provider=bedrock",
|
||||
headers={"Authorization": "Bearer test-key"},
|
||||
return tuple(
|
||||
(encoded_id, raw_id, provider)
|
||||
for raw_id, provider in ((RAW_GCS_OUTPUT_FILE_ID, "vertex_ai"), (RAW_S3_OUTPUT_FILE_ID, "bedrock"))
|
||||
for encoded_id in (quote(raw_id, safe=""), quote(quote(raw_id, safe=""), safe=""))
|
||||
)
|
||||
|
||||
assert response.status_code == 400
|
||||
assert "managed file id" in response.json()["error"]["message"].lower()
|
||||
|
||||
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)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", llm_router)
|
||||
proxy_logging_obj.update_request_status = mocker.AsyncMock()
|
||||
proxy_logging_obj.post_call_failure_hook = mocker.AsyncMock()
|
||||
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
|
||||
api_key="test-key", user_role=user_role, user_id="test-user"
|
||||
)
|
||||
return proxy_logging_obj
|
||||
|
||||
|
||||
@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: pytest.MonkeyPatch,
|
||||
llm_router: Router,
|
||||
encoded_id: str,
|
||||
raw_id: str,
|
||||
provider: str,
|
||||
):
|
||||
import litellm.proxy.proxy_server as ps
|
||||
|
||||
_override_auth(monkeypatch, mocker, llm_router, LitellmUserRoles.INTERNAL_USER)
|
||||
afile_content = mocker.AsyncMock()
|
||||
monkeypatch.setattr(litellm, "afile_content", afile_content)
|
||||
|
||||
try:
|
||||
response = client.get(
|
||||
f"/v1/files/{encoded_id}/content?provider={provider}",
|
||||
headers={"Authorization": "Bearer test-key"},
|
||||
)
|
||||
finally:
|
||||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
|
||||
assert response.status_code == 403, response.text
|
||||
assert "proxy admin" in response.json()["error"]["message"]
|
||||
afile_content.assert_not_called()
|
||||
|
||||
|
||||
@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: pytest.MonkeyPatch,
|
||||
llm_router: Router,
|
||||
encoded_id: str,
|
||||
raw_id: str,
|
||||
provider: str,
|
||||
):
|
||||
import litellm.proxy.proxy_server as ps
|
||||
|
||||
_override_auth(monkeypatch, mocker, llm_router, LitellmUserRoles.PROXY_ADMIN)
|
||||
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", afile_content)
|
||||
|
||||
try:
|
||||
response = client.get(
|
||||
f"/v1/files/{encoded_id}/content?provider={provider}",
|
||||
headers={"Authorization": "Bearer test-key"},
|
||||
)
|
||||
finally:
|
||||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.content == b'{"custom_id": "request-1"}\n'
|
||||
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):
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue