mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
fix(proxy): let a key send managed file ids it may read and refuse s3 sources hidden in model_file_id_mapping
The s3 media gate refused every LiteLLM managed file id whose upload sits under the managed prefix, even when the key's own user or team created it. It now asks the managed files hook which unified ids the key may read and skips only the objects those ids embed. A forged or unowned id, or a proxy without the hook, still gets a 403. The gate also reads the provider file ids in a client-supplied model_file_id_mapping, so an alias can no longer point a non-admin request at an unlisted or managed object. Markdown documents (.md, .markdown) now map to Bedrock's md document format, so an s3 Markdown source without an explicit format no longer fails with an unknown media type.
This commit is contained in:
parent
005030c9ca
commit
37f085c966
7 changed files with 259 additions and 18 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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"] = {
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue