diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 3a6fa5f51da..45be6ad7aa2 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -8,11 +8,14 @@ import sys import threading import time import weakref -from collections.abc import AsyncIterable, Callable, Iterable, Mapping +import zlib +from collections.abc import AsyncGenerator, AsyncIterable, Callable, Iterable, Iterator, Mapping, Sequence +from contextlib import aclosing +from dataclasses import dataclass from http.cookiejar import CookieJar, DefaultCookiePolicy from io import BytesIO from types import MappingProxyType -from typing import TYPE_CHECKING, Any, ClassVar, Final, NoReturn, Optional, TypeAlias, TypedDict, TypeVar +from typing import TYPE_CHECKING, Any, ClassVar, Final, NoReturn, Optional, Protocol, TypeAlias, TypedDict, TypeVar import certifi import httpx @@ -638,6 +641,15 @@ class HTTPResponseLimitError(ValueError): pass +class HTTPResponseEncodingError(HTTPResponseLimitError): + pass + + +_WIRE_BODY_HEADERS: Final = frozenset({"content-encoding", "content-length"}) +_ENCODED_OVERHEAD_DIVISOR: Final = 4096 +_ENCODED_OVERHEAD_BYTES: Final = 64 + + class MaskedHTTPStatusError(httpx.HTTPStatusError): def __init__(self, original_error, message: str | None = None, text: str | None = None): # Create a new error with the masked URL @@ -690,6 +702,223 @@ class MaskedHTTPStatusError(httpx.HTTPStatusError): self.status_code = original_error.response.status_code +def _headers_of_the_decoded_body(headers: httpx.Headers) -> httpx.Headers: + return httpx.Headers( + tuple((name, value) for name, value in headers.multi_items() if name not in _WIRE_BODY_HEADERS) + ) + + +def _encodings_of(headers: httpx.Headers) -> tuple[str, ...]: + declared: Final = header_value(headers, "content-encoding") + names: Final = tuple( + name.strip().lower() for name in ("identity" if declared is None else declared).split(",") if name.strip() + ) + return tuple(name for name in names if name != "identity") or ("identity",) + + +def _encoded_byte_limit(max_bytes: int) -> int: + return max_bytes + max_bytes // _ENCODED_OVERHEAD_DIVISOR + _ENCODED_OVERHEAD_BYTES + + +def _wire_byte_limit(headers: httpx.Headers, max_bytes: int) -> int: + if _encodings_of(headers) == ("identity",): + return max_bytes + return _encoded_byte_limit(max_bytes) + + +async def _wire_bounded(response: httpx.Response, limit: int) -> AsyncGenerator[bytes, None]: + async for chunk in response.aiter_raw(): + if response.num_bytes_downloaded > limit: + raise HTTPResponseLimitError("Response exceeds the configured size limit") + yield chunk + + +def _already_read_within(response: httpx.Response, max_bytes: int) -> httpx.Response: + if len(response.content) > max_bytes: + raise HTTPResponseLimitError("Response exceeds the configured size limit") + return httpx.Response( + response.status_code, + headers=_headers_of_the_decoded_body(response.headers), + content=response.content, + request=response.request, + ) + + +class _BoundedDecoder(Protocol): + def decode(self, data: bytes, max_output: int) -> bytes: ... + + def finish(self) -> None: + """Raise httpx.DecodingError when the wire ended before the compressed stream did.""" + + +class _IdentityDecoder: + def decode(self, data: bytes, max_output: int) -> bytes: + return data + + def finish(self) -> None: + return None + + +class _ZlibDecoder: + def __init__(self, wbits: int) -> None: + self._inflate = zlib.decompressobj(wbits) + self._raw_fallback = wbits == zlib.MAX_WBITS + + def _inflate_once(self, data: bytes, max_output: int) -> bytes: + try: + return self._inflate.decompress(data, max_output) + except zlib.error as zlib_wrapped: + if not self._raw_fallback: + raise httpx.DecodingError(str(zlib_wrapped)) from zlib_wrapped + self._inflate = zlib.decompressobj(-zlib.MAX_WBITS) + try: + return self._inflate.decompress(data, max_output) + except zlib.error as raw_deflate: + raise httpx.DecodingError(str(raw_deflate)) from raw_deflate + + def decode(self, data: bytes, max_output: int) -> bytes: + def drained() -> Iterator[bytes]: + produced = 0 # rebind-ok: running total of the bytes yielded so far + pending = data # rebind-ok: first call feeds the wire bytes, later calls feed the unconsumed tail + while produced < max_output and pending: + piece: bytes = self._inflate_once(pending, max_output - produced) + self._raw_fallback = False + produced += len(piece) + pending = self._inflate.unconsumed_tail + yield piece + + return b"".join(drained()) + + def finish(self) -> None: + if not self._inflate.eof: + raise httpx.DecodingError("Compressed response ended before the end of the stream") + + +class _BrotliInflater(Protocol): + def process(self, data: bytes, output_buffer_limit: int) -> bytes: ... + + def is_finished(self) -> bool: ... + + +@dataclass(frozen=True, slots=True) +class _Brotli: + new_inflater: Callable[[], _BrotliInflater] + error: type[Exception] + + +def _installed_brotli() -> _Brotli | None: + try: + from brotli import ( # pyright: ignore[reportMissingTypeStubs, reportUnknownVariableType] # no stubs + Decompressor, + error, + ) + except ImportError: + return None + return _Brotli(new_inflater=Decompressor, error=error) # pyright: ignore[reportUnknownArgumentType] # brotli has no stubs + + +_BROTLI: Final = _installed_brotli() + + +class _BrotliDecoder: + def __init__(self, brotli: _Brotli) -> None: + self._inflate: Final = brotli.new_inflater() + self._error: Final = brotli.error + + def decode(self, data: bytes, max_output: int) -> bytes: + def drained() -> Iterator[bytes]: + produced = 0 # rebind-ok: running total of the bytes yielded so far + pending = data # rebind-ok: first call feeds the wire bytes, later calls drain buffered output + while produced < max_output: + try: + piece: bytes = self._inflate.process(pending, output_buffer_limit=max_output - produced) + except self._error as exc: + raise httpx.DecodingError(str(exc)) from exc + produced += len(piece) + pending = b"" + yield piece + if not piece or self._inflate.is_finished(): + return + + return b"".join(drained()) + + def finish(self) -> None: + if not self._inflate.is_finished(): + raise httpx.DecodingError("Compressed response ended before the end of the stream") + + +class _CappedStage: + """Intermediate stage of a stacked encoding whose total output over the whole response is capped.""" + + def __init__(self, stage: _BoundedDecoder, cap: int) -> None: + self._stage: Final = stage + self._cap: Final = cap + self._produced = 0 + + def decode(self, data: bytes) -> bytes: + allowance: Final = self._cap - self._produced + out: Final = self._stage.decode(data, allowance + 1) + self._produced += len(out) + if len(out) > allowance: + raise HTTPResponseLimitError("Response exceeds the configured size limit") + return out + + def finish(self) -> None: + self._stage.finish() + + +class _ChainedDecoder: + """Undo stacked Content-Encoding values, last applied first, bounding every intermediate stage.""" + + def __init__(self, stages: Sequence[_BoundedDecoder], max_output: int) -> None: + *outer, last = stages + cap: Final = _encoded_byte_limit(max_output) + self._outer: Final = tuple(_CappedStage(stage, cap) for stage in outer) + self._last: Final = last + + def decode(self, data: bytes, max_output: int) -> bytes: + passed = data # rebind-ok: each stage's output feeds the next stage + for stage in self._outer: + passed = stage.decode(passed) + return self._last.decode(passed, max_output) + + def finish(self) -> None: + for stage in self._outer: + stage.finish() + self._last.finish() + + +def _bounded_decoder(headers: httpx.Headers, max_bytes: int) -> _BoundedDecoder: + brotli: Final = _BROTLI + factories: Final[Mapping[str, Callable[[], _BoundedDecoder]]] = MappingProxyType( + { + "identity": _IdentityDecoder, + "gzip": lambda: _ZlibDecoder(zlib.MAX_WBITS | 16), + "x-gzip": lambda: _ZlibDecoder(zlib.MAX_WBITS | 16), + "deflate": lambda: _ZlibDecoder(zlib.MAX_WBITS), + **({} if brotli is None else {"br": lambda: _BrotliDecoder(brotli)}), + } + ) + encodings: Final = _encodings_of(headers) + if any(name not in factories for name in encodings): + raise HTTPResponseEncodingError( + "Response size limits require an identity, gzip, deflate, or br encoded response" + ) + stages: Final = tuple(factories[name]() for name in reversed(encodings)) + return stages[0] if len(stages) == 1 else _ChainedDecoder(stages, max_bytes) + + +async def _decoded_within(response: httpx.Response, wire: AsyncGenerator[bytes, None], max_bytes: int) -> bytes: + decoder: Final = _bounded_decoder(response.headers, max_bytes) + with BytesIO() as body: + async for chunk in wire: + body.write(decoder.decode(chunk, max_bytes + 1 - body.tell())) + if body.tell() > max_bytes: + raise HTTPResponseLimitError("Response exceeds the configured size limit") + decoder.finish() + return body.getvalue() + + class AsyncHTTPHandler: def __init__( self, @@ -913,19 +1142,18 @@ class AsyncHTTPHandler: ) if response.is_redirect or response.is_error: return httpx.Response(response.status_code, headers=response.headers, request=response.request) - content_encoding: Final = header_value(response.headers, "content-encoding") - if ("identity" if content_encoding is None else content_encoding).lower() != "identity": - raise HTTPResponseLimitError("Response size limits require an uncompressed response") + if response.is_stream_consumed: + return _already_read_within(response, max_bytes) + wire_limit: Final = _wire_byte_limit(response.headers, max_bytes) content_length: Final = header_value(response.headers, "content-length") - if int("0" if content_length is None else content_length) > max_bytes: + if int("0" if content_length is None else content_length) > wire_limit: raise HTTPResponseLimitError("Response exceeds the configured size limit") - with BytesIO() as body: - async for chunk in response.aiter_bytes(chunk_size=65536): - if body.tell() + len(chunk) > max_bytes: - raise HTTPResponseLimitError("Response exceeds the configured size limit") - body.write(chunk) + async with aclosing(_wire_bounded(response, wire_limit)) as wire: return httpx.Response( - response.status_code, headers=response.headers, content=body.getvalue(), request=response.request + response.status_code, + headers=_headers_of_the_decoded_body(response.headers), + content=await _decoded_within(response, wire, max_bytes), + request=response.request, ) finally: await response.aclose() diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index d3ac264518e..7f5c60fb504 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -87,6 +87,7 @@ from litellm.types.proxy.carried_budget_state import ( UserBudgetSnapshot as UserBudgetSnapshot, # noqa: PLC0414 # public re-export ) from litellm.types.proxy.control_plane_endpoints import WorkerRegistryEntry +from litellm.types.proxy.rag_ingest import RagIngestSettings from litellm.types.proxy.request_tag_mappings import RequestTagMapping from litellm.types.proxy.skill_scanning import SkillScannerConfig from litellm.types.proxy.spend_capture_rate import SpendCaptureRateCheckSettings @@ -2662,6 +2663,15 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): None, description="override user_api_key_auth with your own auth script - https://docs.litellm.ai/docs/proxy/virtual_keys#custom-auth", ) + rag_ingest: RagIngestSettings | None = Field( + None, + description=( + "controls run on every upload that lands in a vector store (/v1/rag/ingest and /v1/files with purpose " + "assistants or user_data); rag_ingest.malware_scanner names the scanner instance as . " + "the way custom_auth does, and unset runs the EICAR test scanner; rag_ingest.files_api_controls: false lets " + "the /v1/files uploads through unchecked - https://docs.litellm.ai/docs/rag_ingest" + ), + ) max_parallel_requests: int | None = Field( None, description="maximum parallel requests for each api key", diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index e613d744fb0..0c78242129a 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -106,6 +106,12 @@ from litellm.proxy.openai_files_endpoints.general_upload_validation import ( coerce_optional_str_list_setting, raise_upload_validation_failure, ) +from litellm.proxy.rag_endpoints.upload_security import ( + MAX_UPLOAD_SIZE_BYTES, + MalwareScanner, + RejectedUpload, + validate_upload, +) from litellm.proxy.utils import PrismaClient, ProxyLogging, is_known_model from litellm.repositories.managed_file_repository import ManagedFileRepository from litellm.router import Router @@ -117,6 +123,7 @@ from litellm.types.llms.openai import ( OpenAIFilesPurpose, ) from litellm.types.proxy.auth.user_api_key_auth import user_api_key_has_admin_view +from litellm.types.proxy.rag_ingest import RagIngestSettings from litellm.types.router import Deployment from litellm.types.utils import LlmProviders @@ -542,6 +549,45 @@ async def route_create_file( return response +_VECTOR_STORE_UPLOAD_PURPOSES: Final[frozenset[str]] = frozenset({"assistants", "user_data"}) + + +def _vector_store_upload_controls_apply(purpose: str, settings: RagIngestSettings) -> bool: + return purpose in _VECTOR_STORE_UPLOAD_PURPOSES and settings.files_api_controls + + +def _upload_read_limit_bytes(max_file_size_mb: int | None, controls_apply: bool) -> int | None: + control_limit: Final = MAX_UPLOAD_SIZE_BYTES + 1 if controls_apply else None + if max_file_size_mb is None or max_file_size_mb <= 0: + return control_limit + admin_limit: Final = max_file_size_mb * MB + 1 + return admin_limit if control_limit is None else min(admin_limit, control_limit) + + +def _in_memory_upload(file_source: bytes | BinaryIO) -> bytes: + if isinstance(file_source, bytes): + return file_source + raise ProxyException( + message="Uploads with purpose assistants or user_data must be read into memory before the upload controls run.", + type=ProxyErrorTypes.internal_server_error.value, + param="file", + code=500, + ) + + +async def _reject_unsafe_vector_store_upload(content: bytes, scanner: MalwareScanner) -> None: + validation: Final = await asyncio.to_thread( + validate_upload, content=content, scanner=scanner, max_size_bytes=MAX_UPLOAD_SIZE_BYTES + ) + if isinstance(validation, RejectedUpload): + raise ProxyException( + message=f"{validation.message} Rejection reason: {validation.reason.value}.", + type="invalid_request_error", + param="file", + code=400, + ) + + @router.post( "/{provider}/v1/files", dependencies=[Depends(user_api_key_auth)], @@ -596,6 +642,7 @@ async def create_file( version, ) + rag_ingest_controls: Final = proxy_config.rag_ingest_controls data: dict = {} # Spools this request owns. Starlette owns the upload handle; anything the guardrail scan # opens is ours, and a batch upload that fails after the scan would otherwise hold the @@ -607,18 +654,21 @@ async def create_file( raise_upload_validation_failure(unsafe_filename_failure) max_file_size_mb: Final = coerce_optional_int_setting(general_settings.get("max_file_size_mb")) + controls_apply: Final = _vector_store_upload_controls_apply(purpose, rag_ingest_controls.settings) + read_limit_bytes: Final = _upload_read_limit_bytes(max_file_size_mb, controls_apply) # Batch uploads can be gigabytes. Starlette has already spooled the upload # to disk, so stream from that handle instead of reading it into memory. - # Other uploads stay in-memory bytes, bounded to max_file_size_mb (plus one - # byte, to still tell "exactly at the limit" from "over it") when it is set, - # so an oversized upload cannot be read to completion before it is rejected. + # Other uploads stay in-memory bytes, bounded to the smaller of max_file_size_mb + # and the vector store cap (plus one byte, to still tell "exactly at the limit" + # from "over it") when either applies, so an oversized upload cannot be read to + # completion before it is rejected. file_source: bytes | BinaryIO if purpose == "batch": await file.seek(0) file_source = file.file - elif max_file_size_mb is not None and max_file_size_mb > 0: - file_source = await file.read(max_file_size_mb * MB + 1) + elif read_limit_bytes is not None: + file_source = await file.read(read_limit_bytes) else: file_source = await file.read() custom_llm_provider = ( @@ -671,6 +721,9 @@ async def create_file( if blocked_extension_failure is not None: raise_upload_validation_failure(blocked_extension_failure) + if controls_apply: + await _reject_unsafe_vector_store_upload(_in_memory_upload(file_source), rag_ingest_controls.scanner) + if passthrough: _validate_passthrough_upload( purpose=purpose, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index c1af5d6f980..3a4ee321783 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -950,6 +950,11 @@ from litellm.proxy.prometheus_cleanup import mark_dead_workers, mark_worker_exit from litellm.proxy.public_endpoints import router as public_endpoints_router from litellm.proxy.public_endpoints.public_v1 import router as public_v1_router from litellm.proxy.rag_endpoints.endpoints import router as rag_router +from litellm.proxy.rag_endpoints.upload_security import ( + RagIngestConfigError, + RagIngestControls, + load_rag_ingest_controls, +) from litellm.proxy.rerank_endpoints.endpoints import router as rerank_router from litellm.proxy.response_api_endpoints.endpoints import router as response_router from litellm.proxy.route_llm_request import route_request @@ -5175,6 +5180,14 @@ _DB_OVERLAY_REMOTE_MODULE_LIST_FIELDS: Final[dict[str, tuple[str, ...]]] = { "audit_log_callbacks", ), } +_DB_OVERLAY_REMOTE_MODULE_NESTED_STR_FIELDS: Final[Mapping[str, tuple[tuple[str, str], ...]]] = MappingProxyType( + { + "general_settings": ( + ("litellm_jwtauth", "custom_validate"), + ("rag_ingest", "malware_scanner"), + ), + } +) def _is_remote_module_url(value: object) -> bool: @@ -5281,16 +5294,17 @@ def _scrub_db_overlay_remote_module_loads(section: str, db_value: object) -> obj if isinstance(lp, dict): _scrub_guardrail_inner(lp) - # ``general_settings.litellm_jwtauth.custom_validate`` is a nested - # string field. - if section == "general_settings": - jwt: Final = sanitized.get("litellm_jwtauth") - if isinstance(jwt, dict) and _is_remote_module_url(jwt.get("custom_validate")): + for parent, nested_field in _DB_OVERLAY_REMOTE_MODULE_NESTED_STR_FIELDS.get(section, ()): + if isinstance(nested := sanitized.get(parent), dict) and _is_remote_module_url(nested.get(nested_field)): verbose_proxy_logger.warning( - "Refused remote-URL custom_validate from DB-overlay general_settings.litellm_jwtauth: %r", - jwt.get("custom_validate"), + "Refused remote-URL %s from DB-overlay %s.%s: %r", + nested_field, + section, + parent, + nested.get(nested_field), ) - jwt["custom_validate"] = None + nested[nested_field] = None + if section == "general_settings": # ``pass_through_endpoints`` is a list of dicts whose ``target`` # is passed through ``create_pass_through_route`` → # ``get_instance_fn``. A DB-overlay ``target: "s3://attacker/m.i"`` @@ -5941,6 +5955,7 @@ class ProxyConfig: def __init__(self) -> None: self.config: Mapping[str, object] = MappingProxyType({}) + self.rag_ingest_controls: RagIngestControls = RagIngestControls() self.auto_router_db_catalog: tuple[AutoRouterCatalogEntry, ...] | None = None self._last_semantic_filter_config: dict[str, object] | None = None self._last_websearch_interception_config: dict[str, object] | None = None @@ -7249,6 +7264,16 @@ class ProxyConfig: TypeAdapter(MCPDisabledCapabilities).validate_python(disabled_mcp_capabilities) + match load_rag_ingest_controls( + typed_general_settings.get("rag_ingest"), + config_file_path=config_file_path, + load_instance=get_instance_fn, + ): + case RagIngestConfigError(message=message): + raise ValueError(message) + case RagIngestControls() as rag_ingest_controls: + self.rag_ingest_controls = rag_ingest_controls + if os.getenv("NUM_WORKERS", "1") != "1" and redis_usage_cache is None: warn_login_counters_are_per_worker(os.getenv("NUM_WORKERS", "1")) if declared_proxy_ranges(general_settings) is None: @@ -7372,7 +7397,6 @@ class ProxyConfig: custom_auth_configured=custom_auth is not None, run_common_checks=bool(general_settings.get("custom_auth_run_common_checks", False)), ) - log_once_if_budget_reservation_disabled( disabled=general_settings.get("disable_budget_reservation") is True, ) diff --git a/litellm/proxy/rag_endpoints/endpoints.py b/litellm/proxy/rag_endpoints/endpoints.py index 2d81dc32181..1c0e24826c0 100644 --- a/litellm/proxy/rag_endpoints/endpoints.py +++ b/litellm/proxy/rag_endpoints/endpoints.py @@ -6,12 +6,14 @@ Provides: - /rag/query: RAG query pipeline (Search -> Rerank -> LLM Completion) """ +import asyncio import base64 import json -from collections.abc import Mapping +from collections.abc import Awaitable, Callable, Mapping from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, TypeAlias +import httpx import orjson from fastapi import APIRouter, Depends, HTTPException, Request, Response from fastapi.responses import ORJSONResponse, StreamingResponse @@ -24,6 +26,12 @@ from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import ( LiteLLM_ManagedVectorStore, ) +from litellm.litellm_core_utils.url_utils import async_safe_get +from litellm.llms.custom_httpx.http_handler import ( + HTTPResponseEncodingError, + HTTPResponseLimitError, + get_async_httpx_client, +) from litellm.proxy._types import * from litellm.proxy.auth.auth_utils import is_request_body_safe from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth @@ -46,9 +54,9 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( ) from litellm.proxy.rag_endpoints.upload_security import ( MAX_UPLOAD_SIZE_BYTES, - EicarTestMalwareScanner, MalwareScanner, RejectedUpload, + RejectionReason, validate_upload, ) from litellm.proxy.vector_store_endpoints.endpoints import ( @@ -61,6 +69,7 @@ from litellm.proxy.vector_store_endpoints.utils import ( ) from litellm.rag.main import get_ingestion_class from litellm.repositories.table_repositories import ManagedVectorStoresRepository +from litellm.types.llms.custom_http import httpxSpecialProvider from litellm.types.utils import ModelResponse if TYPE_CHECKING: @@ -471,17 +480,60 @@ async def _save_vector_store_to_db_from_rag_ingest( verbose_proxy_logger.exception("Failed to save vector store %s to database: %s", vector_store_id, db_error) -def _secure_uploaded_file( - file_data: tuple[str, bytes, str], - scanner: MalwareScanner, -) -> tuple[str, bytes, str]: - validation: Final = validate_upload(content=file_data[1], scanner=scanner) +async def _secure_uploaded_file(content: bytes, scanner: MalwareScanner) -> tuple[str, bytes, str]: + validation: Final = await asyncio.to_thread(validate_upload, content=content, scanner=scanner) if isinstance(validation, RejectedUpload): raise HTTPException( status_code=400, detail={"error": validation.message, "reason": validation.reason.value}, ) - return validation.safe_filename, file_data[1], validation.content_type + return validation.safe_filename, content, validation.content_type + + +UrlFetcher: TypeAlias = Callable[[str], Awaitable[httpx.Response]] + +_FILE_URL_FETCH_TIMEOUT: Final = httpx.Timeout(connect=10.0, read=60.0, write=10.0, pool=10.0) + + +async def fetch_file_url(url: str) -> httpx.Response: + client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.RAG) + return await async_safe_get(client, url, timeout=_FILE_URL_FETCH_TIMEOUT, max_response_bytes=MAX_UPLOAD_SIZE_BYTES) + + +async def _download_file_url(file_url: str, fetch_url: UrlFetcher) -> bytes: + try: + response: Final = await fetch_url(file_url) + except HTTPResponseEncodingError as e: + raise HTTPException(status_code=400, detail={"error": f"Could not fetch file_url: {e}"}) from e + except HTTPResponseLimitError as e: + raise HTTPException( + status_code=400, + detail={"error": f"Could not fetch file_url: {e}", "reason": RejectionReason.FILE_TOO_LARGE.value}, + ) from e + except (ValueError, httpx.HTTPError) as e: + raise HTTPException(status_code=400, detail={"error": f"Could not fetch file_url: {e}"}) from e + if not response.is_success: + raise HTTPException( + status_code=400, + detail={"error": f"Could not fetch file_url: the server answered HTTP {response.status_code}"}, + ) + return response.content + + +async def fetch_and_secure_file_url( + file_url: str, + *, + scanner: MalwareScanner, + fetch_url: UrlFetcher, +) -> tuple[str, bytes, str]: + return await _secure_uploaded_file(await _download_file_url(file_url, fetch_url), scanner) + + +def _optional_str_field(body: Mapping[str, object], name: str) -> str | None: + value: Final = body.get(name) + if value is None or isinstance(value, str): + return value + raise HTTPException(status_code=400, detail={"error": f"'{name}' must be a string"}) async def parse_rag_ingest_request( @@ -498,7 +550,9 @@ async def parse_rag_ingest_request( Uploaded file bytes are validated against the vector-store upload controls (size limit, format allowlist with content inspection, archive rejection, and the injected malware scanner) and given a server-generated filename - before they are returned. + before they are returned. A ``file_url`` is returned as given; the route + fetches it through :func:`fetch_and_secure_file_url` once the rest of the + request has been validated. Returns: Tuple of (ingest_options, file_data, file_url, file_id) @@ -526,15 +580,15 @@ async def parse_rag_ingest_request( if request_json_str: request_data: Final = orjson.loads(request_json_str) ingest_options = request_data.get("ingest_options", {}) - file_url = request_data.get("file_url") - file_id = request_data.get("file_id") + file_url = _optional_str_field(request_data, "file_url") # pyright: ignore[reportUnknownArgumentType] # orjson.loads is untyped + file_id = _optional_str_field(request_data, "file_id") # pyright: ignore[reportUnknownArgumentType] # orjson.loads is untyped else: # JSON body data: Final = await read_request_body(request) ingest_options = data.get("ingest_options", {}) - file_url = data.get("file_url") - file_id = data.get("file_id") + file_url = _optional_str_field(data, "file_url") # pyright: ignore[reportUnknownArgumentType] # read_request_body returns a bare dict + file_id = _optional_str_field(data, "file_id") # pyright: ignore[reportUnknownArgumentType] # read_request_body returns a bare dict # Handle base64-encoded file in JSON body file_obj = data.get("file") @@ -560,10 +614,6 @@ async def parse_rag_ingest_request( detail={"error": "Must provide file, file_url, or file_id"}, ) - secured_file_data: Final[tuple[str, bytes, str] | None] = ( - _secure_uploaded_file(file_data, scanner) if file_data is not None else None - ) - if "vector_store" not in ingest_options: raise HTTPException( status_code=400, @@ -605,6 +655,8 @@ async def parse_rag_ingest_request( }, ) + secured_file_data: Final = await _secure_uploaded_file(file_data[1], scanner) if file_data is not None else None + return ingest_options, secured_file_data, file_url, file_id @@ -666,11 +718,10 @@ async def rag_ingest( version, ) + scanner: Final = proxy_config.rag_ingest_controls.scanner try: # Parse request - ingest_options, file_data, file_url, file_id = await parse_rag_ingest_request( - request, scanner=EicarTestMalwareScanner() - ) + ingest_options, file_data, file_url, file_id = await parse_rag_ingest_request(request, scanner=scanner) resolved_stores: Final = await _authorize_nested_vector_store_ids( payload=ingest_options, @@ -721,6 +772,12 @@ async def rag_ingest( detail={"error": provider_error}, ) + uploaded_file_data: Final = ( + await fetch_and_secure_file_url(file_url, scanner=scanner, fetch_url=fetch_file_url) + if file_data is None and file_url is not None + else file_data + ) + verbose_proxy_logger.debug( "RAG Ingest - options: %s, custom_llm_provider: %s", ingest_options, @@ -730,8 +787,7 @@ async def rag_ingest( # Call ingest response: Final = await litellm.aingest( ingest_options=merged_ingest_options, - file_data=file_data, - file_url=file_url, + file_data=uploaded_file_data, file_id=file_id, router=llm_router, **request_data, diff --git a/litellm/proxy/rag_endpoints/upload_security.py b/litellm/proxy/rag_endpoints/upload_security.py index f0318f2f709..8a34a876a48 100644 --- a/litellm/proxy/rag_endpoints/upload_security.py +++ b/litellm/proxy/rag_endpoints/upload_security.py @@ -11,16 +11,21 @@ server-generated filename so the client-controlled name never reaches storage. from __future__ import annotations import uuid -from collections.abc import Mapping -from dataclasses import dataclass +from collections.abc import Callable, Mapping +from dataclasses import dataclass, field from enum import Enum from types import MappingProxyType from typing import Final, Protocol, TypeAlias, runtime_checkable +from pydantic import ValidationError from typing_extensions import assert_never +from litellm.types.proxy.rag_ingest import RagIngestSettings + MAX_UPLOAD_SIZE_BYTES: Final = 512 * 1024 * 1024 +MALWARE_SCANNER_OPTION: Final = "general_settings.rag_ingest.malware_scanner" + EICAR_TEST_SIGNATURE: Final = b"X5O!P%@AP[4\\PZX54(P^)7CC)7}$EICAR-STANDARD-ANTIVIRUS-TEST-FILE!$H+H*" _ARCHIVE_MAGIC_PREFIXES: Final[tuple[bytes, ...]] = ( @@ -91,6 +96,13 @@ class ScanResult: @runtime_checkable class MalwareScanner(Protocol): + """One shared instance scans every upload, from worker threads running at the same time. + + ``scan`` must therefore be safe to call concurrently and must return + :class:`ScanResult` with the ``error`` verdict instead of raising when the + engine is unavailable. + """ + def scan(self, content: bytes) -> ScanResult: ... @@ -110,6 +122,69 @@ class EicarTestMalwareScanner: return ScanResult(verdict=ScanVerdict.CLEAN) +@dataclass(frozen=True, slots=True) +class RagIngestConfigError: + message: str + + +InstanceLoader: TypeAlias = Callable[[str, str | None], object] + + +def parse_rag_ingest_settings(rag_ingest: object) -> RagIngestSettings | RagIngestConfigError: + if rag_ingest is None: + return RagIngestSettings() + try: + return RagIngestSettings.model_validate(rag_ingest) + except ValidationError as e: + return RagIngestConfigError(f"general_settings.rag_ingest is invalid: {e}") + + +def resolve_malware_scanner( + settings: RagIngestSettings, + *, + config_file_path: str | None, + load_instance: InstanceLoader, +) -> MalwareScanner | RagIngestConfigError: + if settings.malware_scanner is None: + return EicarTestMalwareScanner() + try: + loaded: Final = load_instance(settings.malware_scanner, config_file_path) + except Exception as e: # noqa: BLE001 # importing the operator's scanner module can raise anything + return RagIngestConfigError(f"{MALWARE_SCANNER_OPTION}={settings.malware_scanner!r} could not be loaded: {e}") + if isinstance(loaded, type): + return RagIngestConfigError( + f"{MALWARE_SCANNER_OPTION}={settings.malware_scanner!r} names the class {loaded.__qualname__}; " + "name an instance of it instead" + ) + if isinstance(loaded, MalwareScanner) and callable(loaded.scan): + return loaded + return RagIngestConfigError( + f"{MALWARE_SCANNER_OPTION}={settings.malware_scanner!r} resolved to a {type(loaded).__qualname__}, " + "which has no scan(content: bytes) -> ScanResult method" + ) + + +@dataclass(frozen=True, slots=True) +class RagIngestControls: + settings: RagIngestSettings = field(default_factory=RagIngestSettings) + scanner: MalwareScanner = field(default_factory=EicarTestMalwareScanner) + + +def load_rag_ingest_controls( + rag_ingest: object, + *, + config_file_path: str | None, + load_instance: InstanceLoader, +) -> RagIngestControls | RagIngestConfigError: + settings: Final = parse_rag_ingest_settings(rag_ingest) + if isinstance(settings, RagIngestConfigError): + return settings + scanner: Final = resolve_malware_scanner(settings, config_file_path=config_file_path, load_instance=load_instance) + if isinstance(scanner, RagIngestConfigError): + return scanner + return RagIngestControls(settings=settings, scanner=scanner) + + @dataclass(frozen=True, slots=True) class AllowedContent: format: DetectedFormat diff --git a/litellm/types/proxy/rag_ingest.py b/litellm/types/proxy/rag_ingest.py new file mode 100644 index 00000000000..01ebf070bf6 --- /dev/null +++ b/litellm/types/proxy/rag_ingest.py @@ -0,0 +1,26 @@ +"""``general_settings.rag_ingest``: the controls run on bytes before they land in a vector store.""" + +from pydantic import BaseModel, ConfigDict, Field + + +class RagIngestSettings(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + malware_scanner: str | None = Field( + default=None, + description=( + "The scanner every vector-store upload goes through, as . where the module sits " + "next to config.yaml or is importable, resolved the way custom_auth is. The instance must expose " + "scan(content: bytes) -> ScanResult and be safe to call from several threads at once. Unset runs " + "the EICAR test scanner, which flags only the EICAR test file" + ), + ) + files_api_controls: bool = Field( + default=True, + description=( + "Whether /v1/files uploads with purpose assistants or user_data run the upload controls (format, size " + "and malware checks) before reaching the provider. false lets those uploads through unchecked, for a " + "deployment that sends the provider a format the controls reject on that route, such as DOCX or " + "images. /v1/rag/ingest runs the controls whatever this is set to" + ), + ) diff --git a/tests/unit/llms/custom_httpx/test_http_handler.py b/tests/unit/llms/custom_httpx/test_http_handler.py index 3e583bf8dff..8a535184b78 100644 --- a/tests/unit/llms/custom_httpx/test_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_http_handler.py @@ -1,12 +1,15 @@ import asyncio, importlib, json, time import gc +import gzip import io import logging import os import pathlib import ssl import threading +import tracemalloc import weakref +import zlib from collections.abc import AsyncIterator, Callable, Mapping from concurrent.futures import ThreadPoolExecutor from contextlib import asynccontextmanager @@ -30,6 +33,7 @@ from litellm.llms.custom_httpx.http_handler import ( _CLIENT_REFCOUNT_WHEN_HANDLER_IS_SOLE_REFERRER, AsyncHTTPHandler, HTTPHandler, + HTTPResponseLimitError, MaskedHTTPStatusError, get_httpx_client, get_shared_realtime_ssl_context, @@ -1935,6 +1939,266 @@ async def test_bounded_get_stops_redirect_loops(respx_mock, monkeypatch): assert route.call_count == 11 +@pytest.mark.asyncio +async def test_bounded_get_decodes_a_compressed_body_and_reports_the_decoded_length(respx_mock, monkeypatch): + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + document = b"benign document text\n" * 40 + compressed = gzip.compress(document) + respx_mock.get("https://cdn.example/notes.txt").respond( + 200, content=compressed, headers={"content-encoding": "gzip", "content-length": str(len(compressed))} + ) + handler = AsyncHTTPHandler() + try: + response = await handler.get("https://cdn.example/notes.txt", max_response_bytes=len(document)) + finally: + await handler.close() + assert response.content == document + assert "content-encoding" not in response.headers + assert response.headers["content-length"] == str(len(document)) + + +@pytest.mark.asyncio +async def test_bounded_get_reads_a_compressed_body_whose_wire_length_exceeds_the_cap(respx_mock, monkeypatch): + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + document = bytes(range(256)) + compressed = gzip.compress(document) + assert len(compressed) > len(document) + respx_mock.get("https://cdn.example/blob.bin").respond( + 200, content=compressed, headers={"content-encoding": "gzip", "content-length": str(len(compressed))} + ) + handler = AsyncHTTPHandler() + try: + response = await handler.get("https://cdn.example/blob.bin", max_response_bytes=len(document)) + finally: + await handler.close() + assert response.content == document + + +@pytest.mark.asyncio +async def test_bounded_get_caps_the_decoded_size_of_a_compressed_body(respx_mock, monkeypatch): + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + compressed = gzip.compress(b"0" * 200_000) + assert len(compressed) < 50_000 + respx_mock.get("https://cdn.example/bomb.txt").respond( + 200, content=compressed, headers={"content-encoding": "gzip", "content-length": str(len(compressed))} + ) + handler = AsyncHTTPHandler() + try: + with pytest.raises(HTTPResponseLimitError, match="size limit"): + await handler.get("https://cdn.example/bomb.txt", max_response_bytes=50_000) + finally: + await handler.close() + + +@pytest.mark.asyncio +async def test_bounded_get_caps_the_wire_bytes_of_a_compressed_body_that_decodes_to_nothing(respx_mock, monkeypatch): + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + padded = gzip.compress(b"hello") + b"\x00" * 100_000 + respx_mock.get("https://cdn.example/padded.txt").respond(200, content=padded, headers={"content-encoding": "gzip"}) + handler = AsyncHTTPHandler() + try: + with pytest.raises(HTTPResponseLimitError, match="size limit"): + await handler.get("https://cdn.example/padded.txt", max_response_bytes=1024) + finally: + await handler.close() + + +@pytest.mark.asyncio +async def test_bounded_get_serves_a_body_the_transport_already_read(monkeypatch): + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + card: Final = {"data": [{"id": "served", "max_model_len": 8192}]} + + def respond(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, json=card) + + handler = AsyncHTTPHandler() + handler.client = httpx.AsyncClient(transport=httpx.MockTransport(respond)) + try: + response = await handler.get("https://one.test/v1/models", max_response_bytes=2 * 1024 * 1024) + finally: + await handler.close() + assert response.json() == card + assert response.headers["content-length"] == str(len(response.content)) + + +@pytest.mark.asyncio +async def test_bounded_get_bounds_an_already_read_body_on_its_decoded_length(monkeypatch): + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + document: Final = b"0" * 4096 + + def respond(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, content=gzip.compress(document), headers={"content-encoding": "gzip"}) + + handler = AsyncHTTPHandler() + handler.client = httpx.AsyncClient(transport=httpx.MockTransport(respond)) + try: + served = await handler.get("https://cdn.example/notes.txt", max_response_bytes=len(document)) + assert served.content == document + assert "content-encoding" not in served.headers + assert served.headers["content-length"] == str(len(document)) + with pytest.raises(HTTPResponseLimitError, match="size limit"): + await handler.get("https://cdn.example/notes.txt", max_response_bytes=len(document) - 1) + finally: + await handler.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("encoding", "compress"), + [("gzip", gzip.compress), ("deflate", zlib.compress), ("br", None)], +) +async def test_bounded_get_never_inflates_a_compressed_body_past_the_cap( + respx_mock, monkeypatch, encoding: str, compress: Callable[[bytes], bytes] | None +): + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + inflated: Final = 64 * 1024 * 1024 + cap: Final = 1024 * 1024 + bomb: Final = (compress or pytest.importorskip("brotli").compress)(b"\0" * inflated) + assert len(bomb) < cap + + class WireStream(httpx.AsyncByteStream): + async def __aiter__(self): + yield bomb + + respx_mock.get("https://cdn.example/bomb").respond(200, stream=WireStream(), headers={"content-encoding": encoding}) + handler = AsyncHTTPHandler() + tracemalloc.start() + try: + with pytest.raises(HTTPResponseLimitError, match="size limit"): + await handler.get("https://cdn.example/bomb", max_response_bytes=cap) + _, peak = tracemalloc.get_traced_memory() + finally: + tracemalloc.stop() + await handler.close() + assert peak < 4 * cap, f"decoder materialized {peak} bytes for a {cap} byte cap" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("encoding", "compress"), + [("gzip", gzip.compress), ("deflate", zlib.compress), ("br", None)], +) +async def test_bounded_get_decodes_a_compressed_body_under_the_cap_and_rejects_a_truncated_one( + respx_mock, monkeypatch, encoding: str, compress: Callable[[bytes], bytes] | None +): + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + document: Final = b"line %d of the notes\n" * 2000 % tuple(range(2000)) + wire: Final = (compress or pytest.importorskip("brotli").compress)(document) + for name, body in (("whole", wire), ("cut", wire[: len(wire) // 2])): + respx_mock.get(f"https://cdn.example/{name}").respond(200, content=body, headers={"content-encoding": encoding}) + handler = AsyncHTTPHandler() + try: + served = await handler.get("https://cdn.example/whole", max_response_bytes=len(document)) + assert served.content == document + with pytest.raises(httpx.DecodingError, match="ended before the end of the stream"): + await handler.get("https://cdn.example/cut", max_response_bytes=len(document)) + finally: + await handler.close() + + +@pytest.mark.asyncio +async def test_bounded_get_undoes_stacked_encodings_and_still_caps_the_decoded_size(respx_mock, monkeypatch): + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + document: Final = b"line %d of the notes\n" * 2000 % tuple(range(2000)) + stacked: Final = zlib.compress(gzip.compress(document)) + bomb: Final = gzip.compress(gzip.compress(b"\0" * (64 * 1024 * 1024))) + assert len(bomb) < 1024 + + class WireStream(httpx.AsyncByteStream): + def __init__(self, body: bytes) -> None: + self.body: Final = body + + async def __aiter__(self): + yield self.body + + for name, body, encoding in (("doc", stacked, "gzip, deflate"), ("bomb", bomb, "gzip, gzip")): + respx_mock.get(f"https://cdn.example/{name}").respond( + 200, stream=WireStream(body), headers={"content-encoding": encoding} + ) + handler = AsyncHTTPHandler() + tracemalloc.start() + try: + served = await handler.get("https://cdn.example/doc", max_response_bytes=len(document)) + assert served.content == document + with pytest.raises(HTTPResponseLimitError, match="size limit"): + await handler.get("https://cdn.example/bomb", max_response_bytes=1024 * 1024) + _, peak = tracemalloc.get_traced_memory() + finally: + tracemalloc.stop() + await handler.close() + assert peak < 4 * 1024 * 1024, f"stacked decoder materialized {peak} bytes for a 1 MiB cap" + + +@pytest.mark.asyncio +async def test_bounded_get_counts_stacked_intermediate_bytes_across_wire_chunks(respx_mock, monkeypatch): + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + cap: Final = 1024 * 1024 + outer: Final = zlib.compressobj(wbits=zlib.MAX_WBITS | 16) + empty_stored_block: Final = b"\x00\x00\x00\xff\xff" + hollow: Final = zlib.compress(b"")[:2] + empty_stored_block * (8 * 180 * 1024) + zlib.compress(b"")[2:] + assert zlib.decompress(hollow) == b"" + segment: Final = len(hollow) // 8 + 1 + assert segment < cap < len(hollow) + wire_chunks: Final = tuple( + outer.compress(hollow[start : start + segment]) + outer.flush(zlib.Z_SYNC_FLUSH) + for start in range(0, len(hollow), segment) + ) + (outer.flush(),) + + class WireStream(httpx.AsyncByteStream): + async def __aiter__(self): + for chunk in wire_chunks: + yield chunk + + respx_mock.get("https://cdn.example/hollow").respond( + 200, stream=WireStream(), headers={"content-encoding": "deflate, gzip"} + ) + handler = AsyncHTTPHandler() + try: + with pytest.raises(HTTPResponseLimitError, match="size limit"): + await handler.get("https://cdn.example/hollow", max_response_bytes=cap) + finally: + await handler.close() + + +@pytest.mark.asyncio +async def test_bounded_get_refuses_an_encoding_it_cannot_decode_under_the_cap(respx_mock, monkeypatch): + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + + class WireStream(httpx.AsyncByteStream): + async def __aiter__(self): + yield b"\x28\xb5\x2f\xfd" + + respx_mock.get("https://cdn.example/zst").respond(200, stream=WireStream(), headers={"content-encoding": "zstd"}) + handler = AsyncHTTPHandler() + try: + with pytest.raises(HTTPResponseLimitError, match="identity, gzip, deflate, or br"): + await handler.get("https://cdn.example/zst", max_response_bytes=1024) + finally: + await handler.close() + + +@pytest.mark.asyncio +async def test_bounded_get_rejects_a_declared_compressed_length_over_the_wire_limit(respx_mock, monkeypatch): + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + consumed = [] + + class RecordingStream(httpx.AsyncByteStream): + async def __aiter__(self): + consumed.append(True) + yield b"" + + respx_mock.get("https://cdn.example/big.gz").respond( + 200, headers={"content-encoding": "gzip", "content-length": "1000000"}, stream=RecordingStream() + ) + handler = AsyncHTTPHandler() + try: + with pytest.raises(HTTPResponseLimitError, match="size limit"): + await handler.get("https://cdn.example/big.gz", max_response_bytes=1024) + finally: + await handler.close() + assert consumed == [] + + @pytest.mark.asyncio async def test_bounded_get_closes_stream_on_cancellation(respx_mock, monkeypatch): monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") diff --git a/tests/unit/proxy/mcp_server/test_openapi_to_mcp_generator.py b/tests/unit/proxy/mcp_server/test_openapi_to_mcp_generator.py index adb185e4242..89226247cd3 100644 --- a/tests/unit/proxy/mcp_server/test_openapi_to_mcp_generator.py +++ b/tests/unit/proxy/mcp_server/test_openapi_to_mcp_generator.py @@ -1567,7 +1567,14 @@ class TestBoundedOpenAPISpecLoading: assert route.calls[0].request.headers["accept-encoding"] == "identity" @pytest.mark.asyncio - @pytest.mark.parametrize("headers", [{"content-length": "1000000"}, {"content-encoding": "gzip"}]) + @pytest.mark.parametrize( + "headers", + [ + {"content-length": "1000000"}, + {"content-length": "1000000", "content-encoding": "identity"}, + {"content-length": "1000000", "content-encoding": "gzip"}, + ], + ) async def test_unsafe_response_headers_reject_before_reading(self, respx_mock, monkeypatch, headers): import httpx @@ -1592,6 +1599,36 @@ class TestBoundedOpenAPISpecLoading: await load_openapi_spec_async("https://93.184.216.34/spec.json", max_bytes=12) assert closed == [True] + @pytest.mark.asyncio + async def test_compressed_spec_loads_when_its_decoded_size_fits(self, respx_mock, monkeypatch): + import gzip + + from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import load_openapi_spec_async + + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + compressed = gzip.compress(b'{"paths":{}}') + assert len(compressed) > 12 + respx_mock.get("https://93.184.216.34/spec.json").respond( + 200, content=compressed, headers={"content-encoding": "gzip"} + ) + assert await load_openapi_spec_async("https://93.184.216.34/spec.json", max_bytes=12) == {"paths": {}} + + @pytest.mark.asyncio + async def test_compressed_spec_is_capped_on_its_decoded_size(self, respx_mock, monkeypatch): + import gzip + + from litellm.llms.custom_httpx.http_handler import HTTPResponseLimitError + from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import load_openapi_spec_async + + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + compressed = gzip.compress(b'{"paths":{},"pad":"' + b"x" * 100_000 + b'"}') + assert len(compressed) < 12_000 + respx_mock.get("https://93.184.216.34/spec.json").respond( + 200, content=compressed, headers={"content-encoding": "gzip"} + ) + with pytest.raises(HTTPResponseLimitError): + await load_openapi_spec_async("https://93.184.216.34/spec.json", max_bytes=12_000) + @pytest.mark.asyncio async def test_chunked_response_is_bounded_and_closed(self, respx_mock, monkeypatch): import httpx diff --git a/tests/unit/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/unit/proxy/openai_files_endpoint/test_files_endpoint.py index ff63e37bcd8..7aa11e2abcf 100644 --- a/tests/unit/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/unit/proxy/openai_files_endpoint/test_files_endpoint.py @@ -21,7 +21,9 @@ from litellm.proxy.management_endpoints.internal_user_endpoints import ui_view_u from litellm.proxy.openai_files_endpoints.file_content_streaming_handler import ( FileContentStreamingHandler, ) +from litellm.proxy.openai_files_endpoints.general_upload_validation import MB from litellm.proxy.proxy_server import app +from litellm.proxy.rag_endpoints.upload_security import MAX_UPLOAD_SIZE_BYTES from litellm.types.llms.openai import ( FileListPage, HttpxBinaryResponseContent, @@ -5336,7 +5338,7 @@ def test_create_file_blocked_extension_unset_allows_everything(monkeypatch, llm_ try: response = client.post( "/v1/files", - files={"file": ("payload.exe", b"MZ\x90\x00", "application/octet-stream")}, + files={"file": ("payload.exe", b"plain text under an executable name\n", "application/octet-stream")}, data={"purpose": "user_data"}, headers={"Authorization": "Bearer test-key"}, ) @@ -6306,6 +6308,205 @@ def test_create_file_passthrough_fails_closed_when_guardrails_would_scan_the_bat assert forwarded_calls == [] +EICAR_BYTES: Final = b"X5O!P%@AP[4\\PZX54(P^)7CC)7}$EICAR-STANDARD-ANTIVIRUS-TEST-FILE!$H+H*" + + +def _post_file_as_admin(monkeypatch, llm_router: Router, *, purpose: str, name: str, content: bytes, ctype: str): + import litellm.proxy.proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.openai_files_endpoints import files_endpoints as fe + + 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) + setup_proxy_logging_object(monkeypatch, llm_router) + captured: dict = {} + + async def fake_route_create_file(*, _create_file_request, **kwargs): + captured["file"] = _create_file_request["file"] + return OpenAIFileObject( + id="file-clean", + object="file", + bytes=len(content), + created_at=1234567890, + filename=name, + purpose=purpose, + status="uploaded", + ) + + monkeypatch.setattr(fe, "route_create_file", fake_route_create_file) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test-user" + ) + try: + resp = client.post( + "/v1/files", + files={"file": (name, content, ctype)}, + data={"purpose": purpose}, + headers={"Authorization": "Bearer test-key"}, + ) + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + return resp, captured + + +@pytest.mark.parametrize("purpose", ["assistants", "user_data"]) +@pytest.mark.parametrize( + "name, content, ctype, reason", + [ + ("notes.txt", EICAR_BYTES, "text/plain", "malware_detected"), + ("setup.txt", b"#!/bin/sh\nrm -rf /\n", "text/plain", "executable_not_allowed"), + ("report.pdf", b"PK\x03\x04\x14\x00\x00\x00payload", "application/pdf", "archive_not_allowed"), + ("empty.txt", b"", "text/plain", "empty_file"), + ], +) +def test_create_file_for_vector_store_purposes_runs_the_upload_controls( + monkeypatch, llm_router: Router, purpose, name, content, ctype, reason +): + resp, captured = _post_file_as_admin( + monkeypatch, llm_router, purpose=purpose, name=name, content=content, ctype=ctype + ) + assert resp.status_code == 400, resp.text + error = resp.json()["error"] + assert error["type"] == "invalid_request_error" + assert error["param"] == "file" + assert f"Rejection reason: {reason}." in error["message"] + assert captured == {} + + +@pytest.mark.parametrize("purpose", ["assistants", "user_data"]) +@pytest.mark.parametrize( + "name, content, ctype", + [ + ("Quarterly Notes (final).txt", b"benign document text\n", "text/plain"), + ("handbook.pdf", b"%PDF-1.7\n1 0 obj<<>>endobj\n", "application/pdf"), + ], +) +def test_create_file_for_vector_store_purposes_keeps_the_callers_filename_on_clean_uploads( + monkeypatch, llm_router: Router, purpose, name, content, ctype +): + resp, captured = _post_file_as_admin( + monkeypatch, llm_router, purpose=purpose, name=name, content=content, ctype=ctype + ) + assert resp.status_code == 200, resp.text + assert captured["file"][0] == name + assert captured["file"][1] == content + + +def test_create_file_for_vector_store_purposes_uses_the_configured_scanner(monkeypatch, llm_router: Router): + from litellm.proxy.rag_endpoints.upload_security import RagIngestControls, ScanResult, ScanVerdict + + class _InfectedScanner: + def scan(self, content: bytes) -> ScanResult: + return ScanResult(verdict=ScanVerdict.INFECTED, signature="Custom.Sig") + + monkeypatch.setattr( + "litellm.proxy.proxy_server.proxy_config.rag_ingest_controls", RagIngestControls(scanner=_InfectedScanner()) + ) + resp, captured = _post_file_as_admin( + monkeypatch, llm_router, purpose="assistants", name="notes.txt", content=b"benign text\n", ctype="text/plain" + ) + assert resp.status_code == 400, resp.text + assert "Custom.Sig" in resp.json()["error"]["message"] + assert captured == {} + + +@pytest.mark.parametrize( + "purpose, name, content, ctype", + [ + ("batch", "batch.jsonl", VALID_BATCH_LINE, "application/jsonl"), + ("fine-tune", "train.jsonl", VALID_BATCH_LINE, "application/jsonl"), + ("vision", "photo.png", b"\x89PNG\r\n\x1a\n" + b"\x00" * 16, "image/png"), + ], +) +def test_create_file_other_purposes_skip_the_vector_store_upload_controls( + monkeypatch, llm_router: Router, purpose, name, content, ctype +): + from litellm.proxy.rag_endpoints.upload_security import RagIngestControls, ScanResult, ScanVerdict + + class _InfectedScanner: + def scan(self, content: bytes) -> ScanResult: + return ScanResult(verdict=ScanVerdict.INFECTED, signature="Custom.Sig") + + monkeypatch.setattr( + "litellm.proxy.proxy_server.proxy_config.rag_ingest_controls", RagIngestControls(scanner=_InfectedScanner()) + ) + resp, captured = _post_file_as_admin( + monkeypatch, llm_router, purpose=purpose, name=name, content=content, ctype=ctype + ) + assert resp.status_code == 200, resp.text + assert captured["file"][0] == name + + +def test_create_file_for_vector_store_purposes_lets_uploads_through_when_files_api_controls_is_off( + monkeypatch, llm_router: Router +): + from litellm.proxy.rag_endpoints.upload_security import RagIngestControls, ScanResult, ScanVerdict + from litellm.types.proxy.rag_ingest import RagIngestSettings + + class _InfectedScanner: + def scan(self, content: bytes) -> ScanResult: + return ScanResult(verdict=ScanVerdict.INFECTED, signature="Custom.Sig") + + monkeypatch.setattr( + "litellm.proxy.proxy_server.proxy_config.rag_ingest_controls", + RagIngestControls(settings=RagIngestSettings(files_api_controls=False), scanner=_InfectedScanner()), + ) + zip_bytes: Final = b"PK\x03\x04\x14\x00\x00\x00payload" + resp, captured = _post_file_as_admin( + monkeypatch, llm_router, purpose="assistants", name="deck.zip", content=zip_bytes, ctype="application/zip" + ) + assert resp.status_code == 200, resp.text + assert captured["file"][0] == "deck.zip" + assert captured["file"][1] == zip_bytes + + +def test_create_file_for_vector_store_purposes_reads_at_most_the_cap_plus_one_byte(monkeypatch, llm_router: Router): + from litellm.proxy.openai_files_endpoints import files_endpoints as fe + + monkeypatch.setattr(fe, "MAX_UPLOAD_SIZE_BYTES", 10) + resp, captured = _post_file_as_admin( + monkeypatch, llm_router, purpose="assistants", name="notes.txt", content=b"x" * 20, ctype="text/plain" + ) + assert resp.status_code == 400, resp.text + message = resp.json()["error"]["message"] + assert "Uploaded file is 11 bytes" in message + assert "Rejection reason: file_too_large." in message + assert captured == {} + + +@pytest.mark.parametrize( + "max_file_size_mb, controls_apply, expected", + [ + (None, True, MAX_UPLOAD_SIZE_BYTES + 1), + (0, True, MAX_UPLOAD_SIZE_BYTES + 1), + (1, True, MB + 1), + (1024, True, MAX_UPLOAD_SIZE_BYTES + 1), + (1, False, MB + 1), + (None, False, None), + ], +) +def test_upload_read_limit_is_the_smaller_of_the_admin_limit_and_the_vector_store_cap( + max_file_size_mb, controls_apply, expected +): + from litellm.proxy.openai_files_endpoints.files_endpoints import _upload_read_limit_bytes + + assert _upload_read_limit_bytes(max_file_size_mb, controls_apply) == expected + + +def test_vector_store_upload_controls_refuse_a_source_that_was_not_read_into_memory(): + from io import BytesIO + + from litellm.proxy._types import ProxyException + from litellm.proxy.openai_files_endpoints.files_endpoints import _in_memory_upload + + assert _in_memory_upload(b"in memory") == b"in memory" + with pytest.raises(ProxyException) as raised: + _in_memory_upload(BytesIO(b"spooled to disk")) + assert raised.value.code == "500" + assert raised.value.param == "file" + + def _callers_files_on_every_deployment() -> list[MagicMock]: return [ _managed_file_row("unified-gpt-newer", model_mappings={"gpt-3.5-turbo-id": "file-gpt-newer"}), diff --git a/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py b/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py index bf4b5d5d77a..d2e096668ed 100644 --- a/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py +++ b/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py @@ -2,22 +2,33 @@ Tests for RAG proxy endpoints. """ +import base64 import io import json from collections.abc import Iterator, Mapping, Sequence +from dataclasses import dataclass from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest import respx +from fastapi import HTTPException from fastapi.testclient import TestClient import litellm from litellm import Router +from litellm.llms.custom_httpx.http_handler import HTTPResponseLimitError from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.proxy_server import app +from litellm.proxy.rag_endpoints.upload_security import ( + MAX_UPLOAD_SIZE_BYTES, + EicarTestMalwareScanner, + RagIngestControls, + ScanResult, + ScanVerdict, +) from litellm.types.vector_stores import LiteLLM_ManagedVectorStore from litellm.vector_stores.vector_store_registry import VectorStoreRegistry @@ -229,7 +240,9 @@ class TestRagIngestSSRFBlocked: error_text = detail.get("error", "") if isinstance(detail, dict) else str(detail) assert field in error_text, f"Error should name the offending field: {error_text}" - def test_clean_bedrock_ingest_options_not_rejected(self, client_internal_user): + def test_clean_bedrock_ingest_options_not_rejected(self, client_internal_user, monkeypatch): + fetch_url, _seen = _fetcher_answering(200, b"%PDF-1.7\n1 0 obj<<>>endobj\n", content_type="application/pdf") + monkeypatch.setattr("litellm.proxy.rag_endpoints.endpoints.fetch_file_url", fetch_url) with patch( "litellm.proxy.rag_endpoints.endpoints.litellm.aingest", new_callable=AsyncMock, @@ -486,7 +499,12 @@ def test_rag_ingest_uses_the_forwarded_provider_key_over_team_provider_credentia @pytest.mark.parametrize("ingest_request", [_ingest_form, _ingest_json]) -def test_rag_ingest_drop_params_drops_a_forwarded_key_that_cannot_be_an_http_header(client_team_a, ingest_request): +def test_rag_ingest_drop_params_drops_a_forwarded_key_that_cannot_be_an_http_header( + client_team_a, ingest_request, monkeypatch +): + monkeypatch.setattr( + "litellm.proxy.rag_endpoints.endpoints.fetch_file_url", _fetcher_answering(200, b"test content")[0] + ) aingest_patch, registry_patch = _patched_ingest_boundary( None, {"status": "completed", "vector_store_id": "vs_new", "file_id": "file_123"}, @@ -512,7 +530,12 @@ def test_rag_ingest_drop_params_drops_a_forwarded_key_that_cannot_be_an_http_hea @pytest.mark.parametrize("unread_request_field", ["not json", json.dumps({"drop_params": False})]) -def test_rag_ingest_json_body_ignores_a_request_field_it_does_not_read(client_team_a, unread_request_field): +def test_rag_ingest_json_body_ignores_a_request_field_it_does_not_read( + client_team_a, unread_request_field, monkeypatch +): + monkeypatch.setattr( + "litellm.proxy.rag_endpoints.endpoints.fetch_file_url", _fetcher_answering(200, b"test content")[0] + ) aingest_patch, registry_patch = _patched_ingest_boundary( None, {"status": "completed", "vector_store_id": "vs_new", "file_id": "file_123"}, @@ -1764,6 +1787,156 @@ def _multipart_ingest_request(*, filename: str, content: bytes, content_type: st return Request(scope, receive) +async def _never_fetch(url: str): + raise AssertionError(f"no fetch expected, got {url}") + + +@dataclass(frozen=True) +class _FixedScanner: + verdict: ScanVerdict + + def scan(self, content: bytes) -> ScanResult: + return ScanResult(verdict=self.verdict, signature="Custom.Sig") + + +def _fetcher_answering(status_code: int, content: bytes, content_type: str = "text/plain"): + seen: list[str] = [] + + async def fetch_url(url: str) -> httpx.Response: + seen.append(url) + return httpx.Response(status_code, content=content, headers={"content-type": content_type}) + + return fetch_url, seen + + +FILE_URL_BODY = { + "file_url": "https://docs.example.com/reports/q3.txt", + "ingest_options": json.loads(INGEST_REQUEST)["ingest_options"], +} + + +class TestFileUrlUploadControls: + """A ``file_url`` is fetched by the proxy and goes through the same controls as an inline upload.""" + + async def _fetch(self, fetch_url, scanner=None): + from litellm.proxy.rag_endpoints.endpoints import fetch_and_secure_file_url + + return await fetch_and_secure_file_url( + FILE_URL_BODY["file_url"], + scanner=scanner or EicarTestMalwareScanner(), + fetch_url=fetch_url, + ) + + async def _expect_rejection(self, fetch_url, reason: str | None, scanner=None) -> dict: + with pytest.raises(HTTPException) as excinfo: + await self._fetch(fetch_url, scanner=scanner) + assert excinfo.value.status_code == 400 + detail = excinfo.value.detail + assert detail.get("reason") == reason, detail + return detail + + async def test_eicar_behind_a_url_is_blocked(self): + fetch_url, seen = _fetcher_answering(200, EICAR.encode()) + await self._expect_rejection(fetch_url, "malware_detected") + assert seen == [FILE_URL_BODY["file_url"]] + + async def test_shebang_behind_a_url_is_blocked(self): + fetch_url, _seen = _fetcher_answering(200, b"#!/bin/sh\nrm -rf /\n") + await self._expect_rejection(fetch_url, "executable_not_allowed") + + async def test_oversize_url_is_blocked_before_it_is_buffered(self): + async def fetch_url(url: str) -> httpx.Response: + raise HTTPResponseLimitError(f"Response exceeded {MAX_UPLOAD_SIZE_BYTES} bytes") + + detail = await self._expect_rejection(fetch_url, "file_too_large") + assert str(MAX_UPLOAD_SIZE_BYTES) in detail["error"] + + @pytest.mark.parametrize("status_code", [301, 403, 404, 500]) + async def test_non_2xx_fetch_is_a_client_error_naming_the_status(self, status_code): + fetch_url, _seen = _fetcher_answering(status_code, b"benign document text") + detail = await self._expect_rejection(fetch_url, None) + assert f"HTTP {status_code}" in detail["error"] + + async def test_ssrf_or_transport_failure_is_a_client_error(self): + async def fetch_url(url: str) -> httpx.Response: + raise ValueError("Blocked private address") + + detail = await self._expect_rejection(fetch_url, None) + assert "Blocked private address" in detail["error"] + + async def test_clean_text_behind_a_url_is_handed_on_as_bytes(self): + fetch_url, _seen = _fetcher_answering(200, b"benign document text\n", content_type="text/x-whatever") + server_filename, content_bytes, content_type = await self._fetch(fetch_url) + assert content_bytes == b"benign document text\n" + assert server_filename.endswith(".txt") and "/" not in server_filename + assert content_type == "text/plain" + + async def test_url_bytes_go_through_the_injected_scanner(self): + fetch_url, _seen = _fetcher_answering(200, b"benign document text\n") + await self._expect_rejection(fetch_url, "malware_detected", scanner=_FixedScanner(ScanVerdict.INFECTED)) + + def test_route_prefers_an_inline_file_and_never_fetches_the_url(self, client_internal_user, monkeypatch): + monkeypatch.setattr("litellm.proxy.rag_endpoints.endpoints.fetch_file_url", _never_fetch) + inline = {"filename": "inline.txt", "content": base64.b64encode(b"inline text").decode()} + with ( + patch( # test-quality-ok: aingest is the endpoint's downstream boundary; the test asserts the forwarded bytes + "litellm.proxy.rag_endpoints.endpoints.litellm.aingest", + new_callable=AsyncMock, + return_value={"vector_store_id": "vs_new", "file_id": "file_123"}, + ) as mock_aingest + ): + response = client_internal_user.post("/v1/rag/ingest", json={**FILE_URL_BODY, "file": inline}) + assert response.status_code == 200, response.text + assert mock_aingest.await_args.kwargs["file_data"][1] == b"inline text" + + @pytest.mark.parametrize("field", [{"file_url": ["https://docs.example.com/q3.txt"]}, {"file_id": {"id": 1}}]) + def test_route_rejects_a_non_string_file_reference(self, client_internal_user, monkeypatch, field): + monkeypatch.setattr("litellm.proxy.rag_endpoints.endpoints.fetch_file_url", _never_fetch) + body = {"ingest_options": FILE_URL_BODY["ingest_options"], **field} + response = client_internal_user.post("/v1/rag/ingest", json=body) + assert response.status_code == 400, response.text + assert f"'{next(iter(field))}' must be a string" in response.json()["detail"]["error"] + + def test_route_validates_the_request_before_fetching_the_url(self, client_internal_user, monkeypatch): + monkeypatch.setattr("litellm.proxy.rag_endpoints.endpoints.fetch_file_url", _never_fetch) + vector_store = {"custom_llm_provider": "bedrock", "aws_sts_endpoint": "https://attacker.example/sts"} + body = {**FILE_URL_BODY, "ingest_options": {"vector_store": vector_store}} + response = client_internal_user.post("/v1/rag/ingest", json=body) + assert response.status_code == 400, response.text + assert "aws_sts_endpoint" in response.json()["detail"]["error"] + + def test_route_hands_fetched_bytes_to_ingestion_and_never_the_url(self, client_internal_user, monkeypatch): + fetch_url, seen = _fetcher_answering(200, b"benign document text\n") + monkeypatch.setattr("litellm.proxy.rag_endpoints.endpoints.fetch_file_url", fetch_url) + with ( + patch( # test-quality-ok: aingest is the endpoint's downstream boundary; the test asserts the forwarded bytes + "litellm.proxy.rag_endpoints.endpoints.litellm.aingest", + new_callable=AsyncMock, + return_value={"vector_store_id": "vs_new", "file_id": "file_123"}, + ) as mock_aingest + ): + response = client_internal_user.post("/v1/rag/ingest", json=FILE_URL_BODY) + assert response.status_code == 200, response.text + assert seen == [FILE_URL_BODY["file_url"]] + forwarded = mock_aingest.await_args.kwargs + assert forwarded["file_data"][1] == b"benign document text\n" + assert "file_url" not in forwarded + + def test_route_rejects_eicar_behind_a_url(self, client_internal_user, monkeypatch): + fetch_url, _seen = _fetcher_answering(200, EICAR.encode()) + monkeypatch.setattr("litellm.proxy.rag_endpoints.endpoints.fetch_file_url", fetch_url) + with ( + patch( # test-quality-ok: aingest is the endpoint's downstream boundary; the test asserts it is never reached + "litellm.proxy.rag_endpoints.endpoints.litellm.aingest", + new_callable=AsyncMock, + ) as mock_aingest + ): + response = client_internal_user.post("/v1/rag/ingest", json=FILE_URL_BODY) + assert response.status_code == 400, response.text + assert response.json()["detail"]["reason"] == "malware_detected" + mock_aingest.assert_not_awaited() + + class TestVectorStoreUploadControls: """End-to-end enforcement of pentest M4 upload controls on /v1/rag/ingest.""" @@ -1795,6 +1968,47 @@ class TestVectorStoreUploadControls: assert response.status_code == 400, response.text assert response.json()["detail"]["reason"] == "archive_not_allowed" + def test_route_runs_the_configured_scanner(self, client_internal_user, monkeypatch): + monkeypatch.setattr( + "litellm.proxy.proxy_server.proxy_config.rag_ingest_controls", + RagIngestControls(scanner=_FixedScanner(ScanVerdict.INFECTED)), + ) + with ( + patch( # test-quality-ok: aingest is the endpoint's downstream boundary; the test asserts it is never reached + "litellm.proxy.rag_endpoints.endpoints.litellm.aingest", + new_callable=AsyncMock, + ) as mock_aingest + ): + response = client_internal_user.post( + "/v1/rag/ingest", + files={"file": ("doc.txt", io.BytesIO(b"benign document text"), "text/plain")}, + data={"request": INGEST_REQUEST}, + ) + assert response.status_code == 400, response.text + assert response.json()["detail"]["reason"] == "malware_detected" + assert "Custom.Sig" in response.json()["detail"]["error"] + mock_aingest.assert_not_awaited() + + def test_route_lets_a_clean_verdict_from_the_configured_scanner_through(self, client_internal_user, monkeypatch): + monkeypatch.setattr( + "litellm.proxy.proxy_server.proxy_config.rag_ingest_controls", + RagIngestControls(scanner=_FixedScanner(ScanVerdict.CLEAN)), + ) + with ( + patch( # test-quality-ok: aingest is the endpoint's downstream boundary; the test asserts the forwarded bytes + "litellm.proxy.rag_endpoints.endpoints.litellm.aingest", + new_callable=AsyncMock, + return_value={"vector_store_id": "vs_new", "file_id": "file_123"}, + ) as mock_aingest + ): + response = client_internal_user.post( + "/v1/rag/ingest", + files={"file": ("clean_name.txt", io.BytesIO(EICAR.encode()), "text/plain")}, + data={"request": INGEST_REQUEST}, + ) + assert response.status_code == 200, response.text + assert mock_aingest.await_args.kwargs["file_data"][1] == EICAR.encode() + async def test_clean_text_upload_gets_server_generated_filename(self): from litellm.proxy.rag_endpoints.endpoints import parse_rag_ingest_request from litellm.proxy.rag_endpoints.upload_security import EicarTestMalwareScanner diff --git a/tests/unit/proxy/rag_endpoints/test_upload_security.py b/tests/unit/proxy/rag_endpoints/test_upload_security.py index 84375080904..e31387bce12 100644 --- a/tests/unit/proxy/rag_endpoints/test_upload_security.py +++ b/tests/unit/proxy/rag_endpoints/test_upload_security.py @@ -7,6 +7,7 @@ dependency-injected malware scanner validated with the EICAR test file. """ from dataclasses import dataclass +from types import SimpleNamespace import pytest @@ -14,6 +15,7 @@ from litellm.proxy.rag_endpoints.upload_security import ( EICAR_TEST_SIGNATURE, DetectedFormat, EicarTestMalwareScanner, + RagIngestConfigError, RejectedUpload, RejectionReason, ScanResult, @@ -21,9 +23,12 @@ from litellm.proxy.rag_endpoints.upload_security import ( SecuredUpload, generate_safe_filename, inspect_content, + parse_rag_ingest_settings, + resolve_malware_scanner, safe_download_headers, validate_upload, ) +from litellm.types.proxy.rag_ingest import RagIngestSettings @dataclass(frozen=True) @@ -174,3 +179,112 @@ def test_safe_download_headers_sanitize_injection(hostile): disposition = safe_download_headers(hostile)["Content-Disposition"] assert "\r" not in disposition and "\n" not in disposition assert disposition.count('"') == 2 + + +def _loader_returning(loaded: object): + calls: list[tuple[str, str | None]] = [] + + def load_instance(value: str, config_file_path: str | None) -> object: + calls.append((value, config_file_path)) + return loaded + + return load_instance, calls + + +def _loader_raising(error: Exception): + def load_instance(value: str, config_file_path: str | None) -> object: + raise error + + return load_instance + + +def test_parse_rag_ingest_settings_unset_block_runs_every_control_with_the_eicar_scanner(): + for rag_ingest in (None, {}, {"malware_scanner": None}): + assert parse_rag_ingest_settings(rag_ingest) == RagIngestSettings( + malware_scanner=None, files_api_controls=True + ), rag_ingest + + +def test_parse_rag_ingest_settings_reads_the_files_api_opt_out(): + parsed = parse_rag_ingest_settings({"malware_scanner": "custom_scanner.scanner", "files_api_controls": False}) + assert parsed == RagIngestSettings(malware_scanner="custom_scanner.scanner", files_api_controls=False) + + +@pytest.mark.parametrize( + "rag_ingest", + [ + {"malware_scaner": "custom_scanner.scanner"}, + {"malware_scanner": 42}, + {"files_api_controls": "sometimes"}, + "custom_scanner.scanner", + ["x"], + ], +) +def test_parse_rag_ingest_settings_rejects_an_invalid_block(rag_ingest): + parsed = parse_rag_ingest_settings(rag_ingest) + assert isinstance(parsed, RagIngestConfigError) + assert "general_settings.rag_ingest" in parsed.message + + +def test_resolve_malware_scanner_unset_runs_the_eicar_test_scanner(): + load_instance, calls = _loader_returning(_CLEAN_SCANNER) + resolved = resolve_malware_scanner( + RagIngestSettings(), config_file_path="/etc/litellm/config.yaml", load_instance=load_instance + ) + assert isinstance(resolved, EicarTestMalwareScanner) + assert calls == [] + + +def test_resolve_malware_scanner_loads_the_configured_instance_next_to_the_config(): + load_instance, calls = _loader_returning(_INFECTED_SCANNER) + resolved = resolve_malware_scanner( + RagIngestSettings(malware_scanner="custom_scanner.scanner"), + config_file_path="/etc/litellm/config.yaml", + load_instance=load_instance, + ) + assert resolved is _INFECTED_SCANNER + assert calls == [("custom_scanner.scanner", "/etc/litellm/config.yaml")] + assert validate_upload(content=_TEXT_BYTES, scanner=resolved).reason is RejectionReason.MALWARE_DETECTED + + +@pytest.mark.parametrize( + "loaded, expected_fragment", + [ + (object(), "no scan(content: bytes) -> ScanResult method"), + ("custom_scanner.scanner", "no scan(content: bytes) -> ScanResult method"), + (SimpleNamespace(scan="not callable"), "no scan(content: bytes) -> ScanResult method"), + (_StubScanner, "names the class _StubScanner"), + ], +) +def test_resolve_malware_scanner_rejects_an_object_that_is_not_a_scanner(loaded, expected_fragment): + load_instance, _calls = _loader_returning(loaded) + resolved = resolve_malware_scanner( + RagIngestSettings(malware_scanner="custom_scanner.scanner"), + config_file_path=None, + load_instance=load_instance, + ) + assert isinstance(resolved, RagIngestConfigError) + assert "general_settings.rag_ingest.malware_scanner" in resolved.message + assert "custom_scanner.scanner" in resolved.message + assert expected_fragment in resolved.message + + +@pytest.mark.parametrize( + "error", + [ + ImportError("Could not import scanner from custom_scanner"), + AttributeError("module has no attribute scanner"), + ValueError("Empty module name"), + RuntimeError("clamd socket is missing"), + ], +) +def test_resolve_malware_scanner_reports_a_failed_load_with_the_option_name(error): + resolved = resolve_malware_scanner( + RagIngestSettings(malware_scanner="custom_scanner.scanner"), + config_file_path=None, + load_instance=_loader_raising(error), + ) + assert isinstance(resolved, RagIngestConfigError) + assert "general_settings.rag_ingest.malware_scanner" in resolved.message + assert str(error) in resolved.message + diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index 8536cbaa399..11aa4d14777 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -3699,6 +3699,57 @@ async def test_load_config_role_permissions_usable_by_jwt_auth(tmp_path): assert get_role_based_models(rbac_role="team", general_settings=settings) is None +@pytest.mark.asyncio +async def test_load_config_wires_the_configured_malware_scanner_into_uploads(tmp_path): + from litellm.proxy.proxy_server import ProxyConfig + from litellm.proxy.rag_endpoints.upload_security import EicarTestMalwareScanner, ScanVerdict + + proxy_config: Final = ProxyConfig() + (tmp_path / "boot_scanner.py").write_text( + "from litellm.proxy.rag_endpoints.upload_security import ScanResult, ScanVerdict\n" + "class FlagEverything:\n" + " def scan(self, content):\n" + " return ScanResult(verdict=ScanVerdict.INFECTED, signature='Boot.Test')\n" + "scanner = FlagEverything()\n" + "not_a_scanner = object()\n" + ) + config_file: Final = tmp_path / "config.yaml" + + config_file.write_text( + yaml.dump({"model_list": [], "general_settings": {"rag_ingest": {"malware_scanner": "boot_scanner.scanner"}}}) + ) + await proxy_config.load_config(router=MagicMock(), config_file_path=str(config_file)) + assert proxy_config.rag_ingest_controls.scanner.scan(b"plain text").verdict is ScanVerdict.INFECTED + assert proxy_config.rag_ingest_controls.settings.files_api_controls is True + + config_file.write_text( + yaml.dump( + { + "model_list": [], + "general_settings": { + "rag_ingest": {"malware_scanner": "boot_scanner.scanner", "files_api_controls": False} + }, + } + ) + ) + await proxy_config.load_config(router=MagicMock(), config_file_path=str(config_file)) + assert proxy_config.rag_ingest_controls.settings.files_api_controls is False + assert proxy_config.rag_ingest_controls.scanner.scan(b"plain text").verdict is ScanVerdict.INFECTED + + config_file.write_text(yaml.dump({"model_list": []})) + await proxy_config.load_config(router=MagicMock(), config_file_path=str(config_file)) + assert isinstance(proxy_config.rag_ingest_controls.scanner, EicarTestMalwareScanner) + assert proxy_config.rag_ingest_controls.settings.files_api_controls is True + + config_file.write_text( + yaml.dump( + {"model_list": [], "general_settings": {"rag_ingest": {"malware_scanner": "boot_scanner.not_a_scanner"}}} + ) + ) + with pytest.raises(ValueError, match=re.escape("general_settings.rag_ingest.malware_scanner")): + await proxy_config.load_config(router=MagicMock(), config_file_path=str(config_file)) + + @pytest.mark.asyncio async def test_load_config_without_role_permissions_leaves_every_role_unrestricted(tmp_path): from litellm.proxy.auth.auth_checks import get_role_based_models, get_role_based_routes diff --git a/tests/unit/proxy/types_utils/test_db_overlay_remote_module_scrub.py b/tests/unit/proxy/types_utils/test_db_overlay_remote_module_scrub.py index 0072997a0d9..9f5ed39f232 100644 --- a/tests/unit/proxy/types_utils/test_db_overlay_remote_module_scrub.py +++ b/tests/unit/proxy/types_utils/test_db_overlay_remote_module_scrub.py @@ -182,6 +182,21 @@ def test_litellm_jwtauth_custom_validate_stripped(): assert cleaned["litellm_jwtauth"]["user_id_jwt_field"] == "sub" +@pytest.mark.parametrize("remote", ["s3://attacker/m.scanner", "gcs://attacker/m.scanner"]) +def test_rag_ingest_malware_scanner_stripped(remote): + overlay = {"rag_ingest": {"malware_scanner": remote}, "custom_auth": "my_auth.handler"} + cleaned = _scrub_db_overlay_remote_module_loads("general_settings", overlay) + assert cleaned["rag_ingest"]["malware_scanner"] is None + assert cleaned["custom_auth"] == "my_auth.handler" + assert overlay["rag_ingest"]["malware_scanner"] == remote + + +def test_rag_ingest_local_malware_scanner_preserved(): + overlay = {"rag_ingest": {"malware_scanner": "custom_scanner.scanner"}} + cleaned = _scrub_db_overlay_remote_module_loads("general_settings", overlay) + assert cleaned["rag_ingest"]["malware_scanner"] == "custom_scanner.scanner" + + def test_local_dotted_name_preserved(): # The scrub only targets s3:// / gcs:// scheme prefixes — legitimate # dotted module names (the documented operator flow) must pass diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index d14e490f606..7b81f1f2487 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -31262,6 +31262,8 @@ export interface components { * @default 30 */ proxy_config_reload_interval_seconds: number; + /** @description controls run on every upload that lands in a vector store (/v1/rag/ingest and /v1/files with purpose assistants or user_data); rag_ingest.malware_scanner names the scanner instance as . the way custom_auth does, and unset runs the EICAR test scanner; rag_ingest.files_api_controls: false lets the /v1/files uploads through unchecked - https://docs.litellm.ai/docs/rag_ingest */ + rag_ingest?: components["schemas"]["RagIngestSettings"] | null; /** * Reject Clientside Metadata Tags * @description When set to True, rejects requests that contain client-side 'metadata.tags' to prevent users from influencing budgets by sending different tags. Tags can only be inherited from the API key metadata. @@ -43274,6 +43276,20 @@ export interface components { }; } | null; }; + /** RagIngestSettings */ + RagIngestSettings: { + /** + * Files Api Controls + * @description Whether /v1/files uploads with purpose assistants or user_data run the upload controls (format, size and malware checks) before reaching the provider. false lets those uploads through unchecked, for a deployment that sends the provider a format the controls reject on that route, such as DOCX or images. /v1/rag/ingest runs the controls whatever this is set to + * @default true + */ + files_api_controls: boolean; + /** + * Malware Scanner + * @description The scanner every vector-store upload goes through, as . where the module sits next to config.yaml or is importable, resolved the way custom_auth is. The instance must expose scan(content: bytes) -> ScanResult and be safe to call from several threads at once. Unset runs the EICAR test scanner, which flags only the EICAR test file + */ + malware_scanner?: string | null; + }; /** * RankingOptions * @description Ranking options for search.