From cc26342a3a820fd401de57b7528488f961cd97ae Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 25 Sep 2026 16:50:45 -0700 Subject: [PATCH 01/14] feat(rag): configurable malware scanner and upload controls on every vector-store upload route --- litellm/llms/custom_httpx/http_handler.py | 6 +- litellm/proxy/_types.py | 9 + .../openai_files_endpoints/files_endpoints.py | 24 ++ litellm/proxy/proxy_server.py | 45 ++- litellm/proxy/rag_endpoints/endpoints.py | 100 ++++-- .../proxy/rag_endpoints/upload_security.py | 53 ++- litellm/types/proxy/rag_ingest.py | 17 + .../test_files_endpoint.py | 128 ++++++- .../proxy/rag_endpoints/test_rag_endpoints.py | 339 ++++++++++++++---- .../rag_endpoints/test_upload_security.py | 96 +++++ tests/test_litellm/proxy/test_proxy_server.py | 35 ++ .../test_db_overlay_remote_module_scrub.py | 15 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 10 + 13 files changed, 777 insertions(+), 100 deletions(-) create mode 100644 litellm/types/proxy/rag_ingest.py diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 312f46bb021..f9581aa3841 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -553,6 +553,10 @@ class HTTPResponseLimitError(ValueError): pass +class HTTPResponseEncodingError(HTTPResponseLimitError): + pass + + 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 @@ -781,7 +785,7 @@ class AsyncHTTPHandler: if response.is_redirect or response.is_error: return httpx.Response(response.status_code, headers=response.headers, request=response.request) if response.headers.get("content-encoding", "identity").lower() != "identity": - raise HTTPResponseLimitError("Response size limits require an uncompressed response") + raise HTTPResponseEncodingError("Response size limits require an uncompressed response") if int(response.headers.get("content-length", "0")) > max_bytes: raise HTTPResponseLimitError("Response exceeds the configured size limit") with BytesIO() as body: diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 4597872d84e..bfc894798d3 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -52,6 +52,7 @@ from litellm.types.proxy.carried_budget_state import ( UserBudgetSnapshot, ) from litellm.types.proxy.control_plane_endpoints import WorkerRegistryEntry +from litellm.types.proxy.rag_ingest import RagIngestSettings from litellm.types.proxy.spend_capture_rate import SpendCaptureRateCheckSettings from litellm.types.router import RouterErrors, UpdateRouterConfig from litellm.types.router_weights import validate_router_settings_dict @@ -2761,6 +2762,14 @@ 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 - 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 f5e62da5962..4a6ad8adc47 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -94,6 +94,7 @@ 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 MalwareScanner, RejectedUpload, validate_upload from litellm.proxy.utils import PrismaClient, ProxyLogging, is_known_model from litellm.repositories.table_repositories import ManagedFileRepository from litellm.router import Router @@ -527,6 +528,26 @@ async def route_create_file( return response +_VECTOR_STORE_UPLOAD_PURPOSES: Final[frozenset[str]] = frozenset({"assistants", "user_data"}) + + +async def _reject_unsafe_vector_store_upload( + purpose: str, + file_source: bytes | BinaryIO, + scanner: MalwareScanner, +) -> None: + if purpose not in _VECTOR_STORE_UPLOAD_PURPOSES or not isinstance(file_source, bytes): + return + validation: Final = await asyncio.to_thread(validate_upload, content=file_source, scanner=scanner) + 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)], @@ -577,6 +598,7 @@ async def create_file( llm_router, proxy_config, proxy_logging_obj, + rag_upload_malware_scanner, version, ) @@ -655,6 +677,8 @@ async def create_file( if blocked_extension_failure is not None: raise_upload_validation_failure(blocked_extension_failure) + await _reject_unsafe_vector_store_upload(purpose, file_source, rag_upload_malware_scanner) + if passthrough: _validate_passthrough_upload( purpose=purpose, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 61304f0d919..22ba313ed0f 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -742,6 +742,12 @@ 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 ( + EicarTestMalwareScanner, + MalwareScanner, + MalwareScannerConfigError, + resolve_malware_scanner, +) 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 @@ -1035,6 +1041,7 @@ def cleanup_router_config_variables(): user_config_file_path, \ otel_logging, \ user_custom_auth, \ + rag_upload_malware_scanner, \ user_custom_auth_path, \ user_custom_key_generate, \ user_custom_key_update, \ @@ -1053,6 +1060,7 @@ def cleanup_router_config_variables(): user_config_file_path = None otel_logging = None user_custom_auth = None + rag_upload_malware_scanner = EicarTestMalwareScanner() user_custom_auth_path = None user_custom_key_generate = None user_custom_key_update = None @@ -2564,6 +2572,7 @@ polling_via_cache_enabled: Literal["all"] | list[str] | bool = False native_background_mode: list[str] = [] # Models that should use native provider background mode instead of polling polling_cache_ttl: int = 3600 # Default 1 hour TTL for polling cache user_custom_auth = None +rag_upload_malware_scanner: MalwareScanner = EicarTestMalwareScanner() user_custom_key_generate = None # Sentinel: prevents PKCE-no-Redis advisory from re-logging on config hot-reload. # Tests that need to reset it can patch 'litellm.proxy.proxy_server._pkce_no_redis_warning_emitted'. @@ -4511,6 +4520,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: @@ -4607,16 +4624,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"`` @@ -5931,6 +5949,7 @@ class ProxyConfig: user_config_file_path, \ otel_logging, \ user_custom_auth, \ + rag_upload_malware_scanner, \ user_custom_auth_path, \ user_custom_key_generate, \ user_custom_key_update, \ @@ -6347,6 +6366,15 @@ class ProxyConfig: TypeAdapter(MCPAdvertisedVersions).validate_python(general_settings["mcp_advertised_versions"]) + resolved_malware_scanner: Final = resolve_malware_scanner( + general_settings.get("rag_ingest"), + config_file_path=config_file_path, + load_instance=get_instance_fn, + ) + if isinstance(resolved_malware_scanner, MalwareScannerConfigError): + raise ValueError(resolved_malware_scanner.message) + rag_upload_malware_scanner = resolved_malware_scanner + 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: @@ -6463,7 +6491,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 974bff6338a..43a1cbfc1c7 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, status from fastapi.responses import ORJSONResponse, StreamingResponse @@ -23,6 +25,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 @@ -38,9 +46,9 @@ from litellm.proxy.common_utils.http_parsing_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 ( @@ -52,6 +60,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: @@ -388,17 +397,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( @@ -415,7 +467,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) @@ -443,15 +497,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") + file_id = _optional_str_field(request_data, "file_id") 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") + file_id = _optional_str_field(data, "file_id") # Handle base64-encoded file in JSON body file_obj = data.get("file") @@ -477,10 +531,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, @@ -522,6 +572,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 @@ -580,13 +632,14 @@ async def rag_ingest( llm_router, prisma_client, proxy_config, + rag_upload_malware_scanner, version, ) try: # Parse request ingest_options, file_data, file_url, file_id = await parse_rag_ingest_request( - request, scanner=EicarTestMalwareScanner() + request, scanner=rag_upload_malware_scanner ) # INTERNAL_USER_VIEW_ONLY can ingest to existing vector stores only @@ -634,6 +687,12 @@ async def rag_ingest( detail={"error": provider_error}, # mutable-ok: FastAPI serializes the detail as JSON ) + uploaded_file_data: Final = ( + await fetch_and_secure_file_url(file_url, scanner=rag_upload_malware_scanner, fetch_url=fetch_file_url) + if file_data is None and file_url is not None + else file_data + ) + # Add litellm data request_data: dict[str, Any] = {} request_data = await add_litellm_data_to_request( @@ -654,8 +713,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..e7f6a5e7fef 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 collections.abc import Callable, Mapping from dataclasses import dataclass 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,45 @@ class EicarTestMalwareScanner: return ScanResult(verdict=ScanVerdict.CLEAN) +@dataclass(frozen=True, slots=True) +class MalwareScannerConfigError: + message: str + + +InstanceLoader: TypeAlias = Callable[[str, str | None], object] + + +def resolve_malware_scanner( + rag_ingest: object, + *, + config_file_path: str | None, + load_instance: InstanceLoader, +) -> MalwareScanner | MalwareScannerConfigError: + try: + settings: Final = RagIngestSettings() if rag_ingest is None else RagIngestSettings.model_validate(rag_ingest) + except ValidationError as e: + return MalwareScannerConfigError(f"general_settings.rag_ingest is invalid: {e}") + if settings.malware_scanner is None: + return EicarTestMalwareScanner() + try: + loaded: Final = load_instance(settings.malware_scanner, config_file_path) + except Exception as e: + return MalwareScannerConfigError( + f"{MALWARE_SCANNER_OPTION}={settings.malware_scanner!r} could not be loaded: {e}" + ) + if isinstance(loaded, type): + return MalwareScannerConfigError( + 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 MalwareScannerConfigError( + 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 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..e091c6a0434 --- /dev/null +++ b/litellm/types/proxy/rag_ingest.py @@ -0,0 +1,17 @@ +"""``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( + 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" + ), + ) diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py index 84cf4ea7c32..23f4cc56371 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py @@ -4939,7 +4939,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"}, ) @@ -5907,3 +5907,129 @@ def test_create_file_passthrough_fails_closed_when_guardrails_would_scan_the_bat assert error["param"] == "passthrough" assert "guardrails" in error["message"] 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 ScanResult, ScanVerdict + + class _InfectedScanner: + def scan(self, content: bytes) -> ScanResult: + return ScanResult(verdict=ScanVerdict.INFECTED, signature="Custom.Sig") + + monkeypatch.setattr("litellm.proxy.proxy_server.rag_upload_malware_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 ScanResult, ScanVerdict + + class _InfectedScanner: + def scan(self, content: bytes) -> ScanResult: + return ScanResult(verdict=ScanVerdict.INFECTED, signature="Custom.Sig") + + monkeypatch.setattr("litellm.proxy.proxy_server.rag_upload_malware_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 diff --git a/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py b/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py index 2d654ea28ec..4907e5a9ba9 100644 --- a/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py +++ b/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py @@ -5,17 +5,28 @@ Covers: - internal_user_viewer restriction: can only ingest to existing vector stores (must provide vector_store_id) """ +import base64 import io import json +from dataclasses import dataclass from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest +from fastapi import HTTPException from fastapi.testclient import TestClient import litellm +from litellm.llms.custom_httpx.http_handler import HTTPResponseEncodingError, 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, + ScanResult, + ScanVerdict, +) @pytest.fixture @@ -58,19 +69,13 @@ def test_internal_user_viewer_rag_ingest_without_vector_store_id_rejected( response = client_internal_user_viewer.post( "/v1/rag/ingest", files={"file": ("sample.txt", io.BytesIO(b"test content"), "text/plain")}, - data={ - "request": '{"ingest_options":{"vector_store":{"custom_llm_provider":"openai"}}}' - }, + data={"request": '{"ingest_options":{"vector_store":{"custom_llm_provider":"openai"}}}'}, ) assert response.status_code == 403 detail = response.json() assert "detail" in detail - error_msg = ( - detail["detail"]["error"] - if isinstance(detail["detail"], dict) - else str(detail["detail"]) - ) + error_msg = detail["detail"]["error"] if isinstance(detail["detail"], dict) else str(detail["detail"]) assert "internal_user_viewer" in error_msg assert "vector_store_id" in error_msg @@ -97,8 +102,7 @@ def test_internal_user_viewer_rag_ingest_with_vector_store_id_passes_check( # Should not be 403 (role check passed) assert response.status_code != 403, ( - f"internal_user_viewer with vector_store_id should pass role check. " - f"Response: {response.json()}" + f"internal_user_viewer with vector_store_id should pass role check. Response: {response.json()}" ) @@ -114,15 +118,12 @@ def test_internal_user_rag_ingest_without_vector_store_id_allowed(client_interna response = client_internal_user.post( "/v1/rag/ingest", files={"file": ("sample.txt", io.BytesIO(b"test content"), "text/plain")}, - data={ - "request": '{"ingest_options":{"vector_store":{"custom_llm_provider":"openai"}}}' - }, + data={"request": '{"ingest_options":{"vector_store":{"custom_llm_provider":"openai"}}}'}, ) # Should not be 403 assert response.status_code != 403, ( - f"internal_user should be allowed to create new vector stores. " - f"Response: {response.json()}" + f"internal_user should be allowed to create new vector stores. Response: {response.json()}" ) @@ -170,13 +171,13 @@ def test_rag_ingest_blocks_clientside_credentials(client_internal_user, blocked_ }, }, ) - assert ( - response.status_code == 400 - ), f"Expected 400 when '{blocked_field}' is set clientside, got {response.status_code}: {response.json()}" + assert response.status_code == 400, ( + f"Expected 400 when '{blocked_field}' is set clientside, got {response.status_code}: {response.json()}" + ) body = response.json() - assert blocked_field in str( - body - ), f"Response should mention '{blocked_field}': {body}" + assert blocked_field in str(body), f"Response should mention '{blocked_field}': {body}" + + class TestRagIngestSSRFBlocked: """ aws_sts_endpoint and related credential-redirect fields must be rejected @@ -193,9 +194,7 @@ class TestRagIngestSSRFBlocked: ("aws_bedrock_runtime_endpoint", "https://attacker.example/bedrock"), ], ) - def test_ssrf_field_in_vector_store_config_rejected( - self, field, value, client_internal_user - ): + def test_ssrf_field_in_vector_store_config_rejected(self, field, value, client_internal_user): payload = { "file_url": "https://example.com/doc.pdf", "ingest_options": { @@ -215,12 +214,12 @@ class TestRagIngestSSRFBlocked: ) body = response.json() detail = body.get("detail", {}) - error_text = ( - detail.get("error", "") if isinstance(detail, dict) else str(detail) - ) + 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, @@ -230,14 +229,10 @@ class TestRagIngestSSRFBlocked: "/v1/rag/ingest", json={ "file_url": "https://example.com/doc.pdf", - "ingest_options": { - "vector_store": {"custom_llm_provider": "bedrock"} - }, + "ingest_options": {"vector_store": {"custom_llm_provider": "bedrock"}}, }, ) - assert response.status_code != 400, ( - f"Clean Bedrock ingest_options should not be rejected: {response.json()}" - ) + assert response.status_code != 400, f"Clean Bedrock ingest_options should not be rejected: {response.json()}" S3_REGISTRY_STORE = { @@ -703,12 +698,14 @@ def test_rag_query_returns_response_cost_header(client_internal_user): ) mock_response._hidden_params["response_cost"] = 3.45e-06 - with patch( - "litellm.proxy.rag_endpoints.endpoints.litellm.aquery", - new_callable=AsyncMock, - return_value=mock_response, - ), patch("litellm.vector_store_registry", None), patch( - "litellm.proxy.proxy_server.prisma_client", None + with ( + patch( + "litellm.proxy.rag_endpoints.endpoints.litellm.aquery", + new_callable=AsyncMock, + return_value=mock_response, + ), + patch("litellm.vector_store_registry", None), + patch("litellm.proxy.proxy_server.prisma_client", None), ): response = client_internal_user.post( "/v1/rag/query", @@ -742,7 +739,9 @@ def test_rag_query_surfaces_upstream_status_code(client_internal_user, upstream_ new=AsyncMock(side_effect=upstream_error), ), patch("litellm.vector_store_registry", None), # test-quality-ok: proxy module global, no injection seam - patch("litellm.proxy.proxy_server.prisma_client", None), # test-quality-ok: proxy module global, no injection seam + patch( + "litellm.proxy.proxy_server.prisma_client", None + ), # test-quality-ok: proxy module global, no injection seam ): response = client_internal_user.post( "/v1/rag/query", @@ -779,10 +778,14 @@ def test_rag_query_stream_returns_event_stream(client_internal_user): api_key="test-key", ) - with patch( - "litellm.proxy.rag_endpoints.endpoints.litellm.aquery", - new=AsyncMock(side_effect=fake_aquery), - ), patch("litellm.vector_store_registry", None), patch("litellm.proxy.proxy_server.prisma_client", None): + with ( + patch( + "litellm.proxy.rag_endpoints.endpoints.litellm.aquery", + new=AsyncMock(side_effect=fake_aquery), + ), + patch("litellm.vector_store_registry", None), + patch("litellm.proxy.proxy_server.prisma_client", None), + ): response = client_internal_user.post( "/v1/rag/query", json={ @@ -825,7 +828,9 @@ def test_rag_query_stream_pings_while_retrieval_is_still_running(client_internal new=AsyncMock(side_effect=slow_aquery), ), patch("litellm.vector_store_registry", None), # test-quality-ok: proxy module global, no injection seam - patch("litellm.proxy.proxy_server.prisma_client", None), # test-quality-ok: proxy module global, no injection seam + patch( + "litellm.proxy.proxy_server.prisma_client", None + ), # test-quality-ok: proxy module global, no injection seam ): response = client_internal_user.post( "/v1/rag/query", @@ -849,9 +854,7 @@ def test_rag_query_stream_pings_while_retrieval_is_still_running(client_internal assert response.text.endswith("data: [DONE]\n\n") -def test_rag_query_stream_keeps_response_headers_when_retrieval_beats_the_keepalive( - client_internal_user, monkeypatch -): +def test_rag_query_stream_keeps_response_headers_when_retrieval_beats_the_keepalive(client_internal_user, monkeypatch): import litellm as litellm_module monkeypatch.setattr(litellm_module, "sse_keepalive_ping_interval_seconds", 5) @@ -873,7 +876,9 @@ def test_rag_query_stream_keeps_response_headers_when_retrieval_beats_the_keepal new=AsyncMock(side_effect=fast_aquery), ), patch("litellm.vector_store_registry", None), # test-quality-ok: proxy module global, no injection seam - patch("litellm.proxy.proxy_server.prisma_client", None), # test-quality-ok: proxy module global, no injection seam + patch( + "litellm.proxy.proxy_server.prisma_client", None + ), # test-quality-ok: proxy module global, no injection seam ): response = client_internal_user.post( "/v1/rag/query", @@ -925,13 +930,17 @@ def test_rag_query_merges_managed_store_params(client_internal_user): model="gpt-4o-mini", ) - with patch( # test-quality-ok: aquery is the endpoint's downstream boundary; the forwarded config is what the test asserts - "litellm.proxy.rag_endpoints.endpoints.litellm.aquery", - new_callable=AsyncMock, - return_value=mock_response, - ) as mock_aquery, patch.object(litellm, "vector_store_registry", mock_registry), patch( # test-quality-ok: seeds the managed-store registry the merge under test reads and grants access so real store resolution runs - "litellm.proxy.vector_store_endpoints.utils.can_user_access_vector_store", - new=AsyncMock(return_value=True), + with ( + patch( # test-quality-ok: aquery is the endpoint's downstream boundary; the forwarded config is what the test asserts + "litellm.proxy.rag_endpoints.endpoints.litellm.aquery", + new_callable=AsyncMock, + return_value=mock_response, + ) as mock_aquery, + patch.object(litellm, "vector_store_registry", mock_registry), + patch( # test-quality-ok: seeds the managed-store registry the merge under test reads and grants access so real store resolution runs + "litellm.proxy.vector_store_endpoints.utils.can_user_access_vector_store", + new=AsyncMock(return_value=True), + ), ): response = client_internal_user.post( "/v1/rag/query", @@ -971,13 +980,17 @@ def test_rag_query_store_params_win_over_user_retrieval_config(client_internal_u model="gpt-4o-mini", ) - with patch( # test-quality-ok: aquery is the endpoint's downstream boundary; the forwarded config is what the test asserts - "litellm.proxy.rag_endpoints.endpoints.litellm.aquery", - new_callable=AsyncMock, - return_value=mock_response, - ) as mock_aquery, patch.object(litellm, "vector_store_registry", mock_registry), patch( # test-quality-ok: seeds the managed-store registry the merge under test reads and grants access so real store resolution runs - "litellm.proxy.vector_store_endpoints.utils.can_user_access_vector_store", - new=AsyncMock(return_value=True), + with ( + patch( # test-quality-ok: aquery is the endpoint's downstream boundary; the forwarded config is what the test asserts + "litellm.proxy.rag_endpoints.endpoints.litellm.aquery", + new_callable=AsyncMock, + return_value=mock_response, + ) as mock_aquery, + patch.object(litellm, "vector_store_registry", mock_registry), + patch( # test-quality-ok: seeds the managed-store registry the merge under test reads and grants access so real store resolution runs + "litellm.proxy.vector_store_endpoints.utils.can_user_access_vector_store", + new=AsyncMock(return_value=True), + ), ): response = client_internal_user.post( "/v1/rag/query", @@ -1124,6 +1137,163 @@ 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"] + + async def test_compressed_url_response_is_refused_without_calling_it_too_large(self): + async def fetch_url(url: str) -> httpx.Response: + raise HTTPResponseEncodingError("Response size limits require an uncompressed response") + + detail = await self._expect_rejection(fetch_url, None) + assert "uncompressed" 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.""" @@ -1155,6 +1325,43 @@ 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.rag_upload_malware_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.rag_upload_malware_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 @@ -1164,9 +1371,7 @@ class TestVectorStoreUploadControls: content=b"benign document text\n", content_type="text/plain", ) - _options, file_data, _url, _file_id = await parse_rag_ingest_request( - request, scanner=EicarTestMalwareScanner() - ) + _options, file_data, _url, _file_id = await parse_rag_ingest_request(request, scanner=EicarTestMalwareScanner()) assert file_data is not None server_filename, content_bytes, secured_content_type = file_data assert server_filename != "../../etc/passwd" diff --git a/tests/test_litellm/proxy/rag_endpoints/test_upload_security.py b/tests/test_litellm/proxy/rag_endpoints/test_upload_security.py index 84375080904..1347fea8c1f 100644 --- a/tests/test_litellm/proxy/rag_endpoints/test_upload_security.py +++ b/tests/test_litellm/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, + MalwareScannerConfigError, RejectedUpload, RejectionReason, ScanResult, @@ -21,6 +23,7 @@ from litellm.proxy.rag_endpoints.upload_security import ( SecuredUpload, generate_safe_filename, inspect_content, + resolve_malware_scanner, safe_download_headers, validate_upload, ) @@ -174,3 +177,96 @@ 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_resolve_malware_scanner_unset_runs_the_eicar_test_scanner(): + load_instance, calls = _loader_returning(_CLEAN_SCANNER) + for rag_ingest in (None, {}, {"malware_scanner": None}): + resolved = resolve_malware_scanner( + rag_ingest, config_file_path="/etc/litellm/config.yaml", load_instance=load_instance + ) + assert isinstance(resolved, EicarTestMalwareScanner), rag_ingest + 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( + {"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( + {"malware_scanner": "custom_scanner.scanner"}, + config_file_path=None, + load_instance=load_instance, + ) + assert isinstance(resolved, MalwareScannerConfigError) + 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( + {"malware_scanner": "custom_scanner.scanner"}, + config_file_path=None, + load_instance=_loader_raising(error), + ) + assert isinstance(resolved, MalwareScannerConfigError) + assert "general_settings.rag_ingest.malware_scanner" in resolved.message + assert str(error) in resolved.message + + +@pytest.mark.parametrize( + "rag_ingest", + [{"malware_scaner": "custom_scanner.scanner"}, {"malware_scanner": 42}, "custom_scanner.scanner", ["x"]], +) +def test_resolve_malware_scanner_rejects_an_invalid_rag_ingest_block(rag_ingest): + load_instance, calls = _loader_returning(_CLEAN_SCANNER) + resolved = resolve_malware_scanner(rag_ingest, config_file_path=None, load_instance=load_instance) + assert isinstance(resolved, MalwareScannerConfigError) + assert "general_settings.rag_ingest" in resolved.message + assert calls == [] diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 884a9c81500..ecf19e362e4 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -3518,6 +3518,41 @@ 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, monkeypatch): + from litellm.proxy.proxy_server import ProxyConfig + from litellm.proxy.rag_endpoints.upload_security import EicarTestMalwareScanner, ScanVerdict + + monkeypatch.setattr(proxy_server_module, "rag_upload_malware_scanner", EicarTestMalwareScanner()) + (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 ProxyConfig().load_config(router=MagicMock(), config_file_path=str(config_file)) + assert proxy_server_module.rag_upload_malware_scanner.scan(b"plain text").verdict is ScanVerdict.INFECTED + + config_file.write_text(yaml.dump({"model_list": []})) + await ProxyConfig().load_config(router=MagicMock(), config_file_path=str(config_file)) + assert isinstance(proxy_server_module.rag_upload_malware_scanner, EicarTestMalwareScanner) + + 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 ProxyConfig().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/test_litellm/proxy/types_utils/test_db_overlay_remote_module_scrub.py b/tests/test_litellm/proxy/types_utils/test_db_overlay_remote_module_scrub.py index 0072997a0d9..9f5ed39f232 100644 --- a/tests/test_litellm/proxy/types_utils/test_db_overlay_remote_module_scrub.py +++ b/tests/test_litellm/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 1e0aed46923..f63e4657a97 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -28451,6 +28451,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 - 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. @@ -39051,6 +39053,14 @@ export interface components { }; } | null; }; + /** RagIngestSettings */ + RagIngestSettings: { + /** + * 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. From 328fe2aec28d4e120e09a5f22060f219b2c4c93f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 25 Sep 2026 17:52:14 -0700 Subject: [PATCH 02/14] fix(files): opt-out flag, bounded read, and fail-closed source for the vector store upload controls --- litellm/proxy/_types.py | 3 +- .../openai_files_endpoints/files_endpoints.py | 60 ++++++++++++---- litellm/proxy/proxy_server.py | 16 ++++- .../proxy/rag_endpoints/upload_security.py | 27 ++++---- litellm/types/proxy/rag_ingest.py | 11 ++- .../test_files_endpoint.py | 69 +++++++++++++++++++ .../rag_endpoints/test_upload_security.py | 62 +++++++++++------ tests/test_litellm/proxy/test_proxy_server.py | 18 +++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 8 ++- 9 files changed, 219 insertions(+), 55 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index bfc894798d3..09c99969ba2 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2767,7 +2767,8 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): 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 - https://docs.litellm.ai/docs/rag_ingest" + "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( diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index 4a6ad8adc47..1dead8ac504 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -94,7 +94,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 MalwareScanner, RejectedUpload, validate_upload +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.table_repositories import ManagedFileRepository from litellm.router import Router @@ -105,6 +110,7 @@ from litellm.types.llms.openai import ( OpenAIFileObject, OpenAIFilesPurpose, ) +from litellm.types.proxy.rag_ingest import RagIngestSettings router: Final = APIRouter() @@ -531,14 +537,33 @@ async def route_create_file( _VECTOR_STORE_UPLOAD_PURPOSES: Final[frozenset[str]] = frozenset({"assistants", "user_data"}) -async def _reject_unsafe_vector_store_upload( - purpose: str, - file_source: bytes | BinaryIO, - scanner: MalwareScanner, -) -> None: - if purpose not in _VECTOR_STORE_UPLOAD_PURPOSES or not isinstance(file_source, bytes): - return - validation: Final = await asyncio.to_thread(validate_upload, content=file_source, scanner=scanner) +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}.", @@ -598,6 +623,7 @@ async def create_file( llm_router, proxy_config, proxy_logging_obj, + rag_ingest_settings, rag_upload_malware_scanner, version, ) @@ -613,18 +639,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_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 = ( @@ -677,7 +706,8 @@ async def create_file( if blocked_extension_failure is not None: raise_upload_validation_failure(blocked_extension_failure) - await _reject_unsafe_vector_store_upload(purpose, file_source, rag_upload_malware_scanner) + if controls_apply: + await _reject_unsafe_vector_store_upload(_in_memory_upload(file_source), rag_upload_malware_scanner) if passthrough: _validate_passthrough_upload( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 22ba313ed0f..58f7b436ba3 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -745,7 +745,8 @@ from litellm.proxy.rag_endpoints.endpoints import router as rag_router from litellm.proxy.rag_endpoints.upload_security import ( EicarTestMalwareScanner, MalwareScanner, - MalwareScannerConfigError, + RagIngestConfigError, + parse_rag_ingest_settings, resolve_malware_scanner, ) from litellm.proxy.rerank_endpoints.endpoints import router as rerank_router @@ -867,6 +868,7 @@ from litellm.types.proxy.model_deprecation import ( DEFAULT_DEPRECATION_WARN_DAYS, ModelDeprecationResponse, ) +from litellm.types.proxy.rag_ingest import RagIngestSettings from litellm.types.proxy.spend_capture_rate import SpendCaptureProvider, SpendCaptureRateCheckSettings from litellm.types.realtime import RealtimeQueryParams from litellm.types.router import ( @@ -1042,6 +1044,7 @@ def cleanup_router_config_variables(): otel_logging, \ user_custom_auth, \ rag_upload_malware_scanner, \ + rag_ingest_settings, \ user_custom_auth_path, \ user_custom_key_generate, \ user_custom_key_update, \ @@ -1061,6 +1064,7 @@ def cleanup_router_config_variables(): otel_logging = None user_custom_auth = None rag_upload_malware_scanner = EicarTestMalwareScanner() + rag_ingest_settings = RagIngestSettings() user_custom_auth_path = None user_custom_key_generate = None user_custom_key_update = None @@ -2573,6 +2577,7 @@ native_background_mode: list[str] = [] # Models that should use native provider polling_cache_ttl: int = 3600 # Default 1 hour TTL for polling cache user_custom_auth = None rag_upload_malware_scanner: MalwareScanner = EicarTestMalwareScanner() +rag_ingest_settings: RagIngestSettings = RagIngestSettings() user_custom_key_generate = None # Sentinel: prevents PKCE-no-Redis advisory from re-logging on config hot-reload. # Tests that need to reset it can patch 'litellm.proxy.proxy_server._pkce_no_redis_warning_emitted'. @@ -5950,6 +5955,7 @@ class ProxyConfig: otel_logging, \ user_custom_auth, \ rag_upload_malware_scanner, \ + rag_ingest_settings, \ user_custom_auth_path, \ user_custom_key_generate, \ user_custom_key_update, \ @@ -6366,13 +6372,17 @@ class ProxyConfig: TypeAdapter(MCPAdvertisedVersions).validate_python(general_settings["mcp_advertised_versions"]) + parsed_rag_ingest_settings: Final = parse_rag_ingest_settings(general_settings.get("rag_ingest")) + if isinstance(parsed_rag_ingest_settings, RagIngestConfigError): + raise ValueError(parsed_rag_ingest_settings.message) resolved_malware_scanner: Final = resolve_malware_scanner( - general_settings.get("rag_ingest"), + parsed_rag_ingest_settings, config_file_path=config_file_path, load_instance=get_instance_fn, ) - if isinstance(resolved_malware_scanner, MalwareScannerConfigError): + if isinstance(resolved_malware_scanner, RagIngestConfigError): raise ValueError(resolved_malware_scanner.message) + rag_ingest_settings = parsed_rag_ingest_settings rag_upload_malware_scanner = resolved_malware_scanner if os.getenv("NUM_WORKERS", "1") != "1" and redis_usage_cache is None: diff --git a/litellm/proxy/rag_endpoints/upload_security.py b/litellm/proxy/rag_endpoints/upload_security.py index e7f6a5e7fef..4579fd65cca 100644 --- a/litellm/proxy/rag_endpoints/upload_security.py +++ b/litellm/proxy/rag_endpoints/upload_security.py @@ -123,39 +123,42 @@ class EicarTestMalwareScanner: @dataclass(frozen=True, slots=True) -class MalwareScannerConfigError: +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( - rag_ingest: object, + settings: RagIngestSettings, *, config_file_path: str | None, load_instance: InstanceLoader, -) -> MalwareScanner | MalwareScannerConfigError: - try: - settings: Final = RagIngestSettings() if rag_ingest is None else RagIngestSettings.model_validate(rag_ingest) - except ValidationError as e: - return MalwareScannerConfigError(f"general_settings.rag_ingest is invalid: {e}") +) -> 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: - return MalwareScannerConfigError( - f"{MALWARE_SCANNER_OPTION}={settings.malware_scanner!r} could not be loaded: {e}" - ) + return RagIngestConfigError(f"{MALWARE_SCANNER_OPTION}={settings.malware_scanner!r} could not be loaded: {e}") if isinstance(loaded, type): - return MalwareScannerConfigError( + 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 MalwareScannerConfigError( + return RagIngestConfigError( f"{MALWARE_SCANNER_OPTION}={settings.malware_scanner!r} resolved to a {type(loaded).__qualname__}, " "which has no scan(content: bytes) -> ScanResult method" ) diff --git a/litellm/types/proxy/rag_ingest.py b/litellm/types/proxy/rag_ingest.py index e091c6a0434..01ebf070bf6 100644 --- a/litellm/types/proxy/rag_ingest.py +++ b/litellm/types/proxy/rag_ingest.py @@ -7,7 +7,7 @@ class RagIngestSettings(BaseModel): model_config = ConfigDict(frozen=True, extra="forbid") malware_scanner: str | None = Field( - None, + 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 " @@ -15,3 +15,12 @@ class RagIngestSettings(BaseModel): "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/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py index 23f4cc56371..b02de1346d9 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py @@ -18,7 +18,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, @@ -6033,3 +6035,70 @@ def test_create_file_other_purposes_skip_the_vector_store_upload_controls( ) 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 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.rag_upload_malware_scanner", _InfectedScanner()) + monkeypatch.setattr("litellm.proxy.proxy_server.rag_ingest_settings", RagIngestSettings(files_api_controls=False)) + 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" diff --git a/tests/test_litellm/proxy/rag_endpoints/test_upload_security.py b/tests/test_litellm/proxy/rag_endpoints/test_upload_security.py index 1347fea8c1f..e31387bce12 100644 --- a/tests/test_litellm/proxy/rag_endpoints/test_upload_security.py +++ b/tests/test_litellm/proxy/rag_endpoints/test_upload_security.py @@ -15,7 +15,7 @@ from litellm.proxy.rag_endpoints.upload_security import ( EICAR_TEST_SIGNATURE, DetectedFormat, EicarTestMalwareScanner, - MalwareScannerConfigError, + RagIngestConfigError, RejectedUpload, RejectionReason, ScanResult, @@ -23,10 +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) @@ -196,20 +198,47 @@ def _loader_raising(error: Exception): 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) - for rag_ingest in (None, {}, {"malware_scanner": None}): - resolved = resolve_malware_scanner( - rag_ingest, config_file_path="/etc/litellm/config.yaml", load_instance=load_instance - ) - assert isinstance(resolved, EicarTestMalwareScanner), rag_ingest + 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( - {"malware_scanner": "custom_scanner.scanner"}, + RagIngestSettings(malware_scanner="custom_scanner.scanner"), config_file_path="/etc/litellm/config.yaml", load_instance=load_instance, ) @@ -230,11 +259,11 @@ def test_resolve_malware_scanner_loads_the_configured_instance_next_to_the_confi 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( - {"malware_scanner": "custom_scanner.scanner"}, + RagIngestSettings(malware_scanner="custom_scanner.scanner"), config_file_path=None, load_instance=load_instance, ) - assert isinstance(resolved, MalwareScannerConfigError) + 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 @@ -251,22 +280,11 @@ def test_resolve_malware_scanner_rejects_an_object_that_is_not_a_scanner(loaded, ) def test_resolve_malware_scanner_reports_a_failed_load_with_the_option_name(error): resolved = resolve_malware_scanner( - {"malware_scanner": "custom_scanner.scanner"}, + RagIngestSettings(malware_scanner="custom_scanner.scanner"), config_file_path=None, load_instance=_loader_raising(error), ) - assert isinstance(resolved, MalwareScannerConfigError) + assert isinstance(resolved, RagIngestConfigError) assert "general_settings.rag_ingest.malware_scanner" in resolved.message assert str(error) in resolved.message - -@pytest.mark.parametrize( - "rag_ingest", - [{"malware_scaner": "custom_scanner.scanner"}, {"malware_scanner": 42}, "custom_scanner.scanner", ["x"]], -) -def test_resolve_malware_scanner_rejects_an_invalid_rag_ingest_block(rag_ingest): - load_instance, calls = _loader_returning(_CLEAN_SCANNER) - resolved = resolve_malware_scanner(rag_ingest, config_file_path=None, load_instance=load_instance) - assert isinstance(resolved, MalwareScannerConfigError) - assert "general_settings.rag_ingest" in resolved.message - assert calls == [] diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index ecf19e362e4..7e8647d69ab 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -3522,8 +3522,10 @@ async def test_load_config_role_permissions_usable_by_jwt_auth(tmp_path): async def test_load_config_wires_the_configured_malware_scanner_into_uploads(tmp_path, monkeypatch): from litellm.proxy.proxy_server import ProxyConfig from litellm.proxy.rag_endpoints.upload_security import EicarTestMalwareScanner, ScanVerdict + from litellm.types.proxy.rag_ingest import RagIngestSettings monkeypatch.setattr(proxy_server_module, "rag_upload_malware_scanner", EicarTestMalwareScanner()) + monkeypatch.setattr(proxy_server_module, "rag_ingest_settings", RagIngestSettings(files_api_controls=False)) (tmp_path / "boot_scanner.py").write_text( "from litellm.proxy.rag_endpoints.upload_security import ScanResult, ScanVerdict\n" "class FlagEverything:\n" @@ -3539,10 +3541,26 @@ async def test_load_config_wires_the_configured_malware_scanner_into_uploads(tmp ) await ProxyConfig().load_config(router=MagicMock(), config_file_path=str(config_file)) assert proxy_server_module.rag_upload_malware_scanner.scan(b"plain text").verdict is ScanVerdict.INFECTED + assert proxy_server_module.rag_ingest_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 ProxyConfig().load_config(router=MagicMock(), config_file_path=str(config_file)) + assert proxy_server_module.rag_ingest_settings.files_api_controls is False + assert proxy_server_module.rag_upload_malware_scanner.scan(b"plain text").verdict is ScanVerdict.INFECTED config_file.write_text(yaml.dump({"model_list": []})) await ProxyConfig().load_config(router=MagicMock(), config_file_path=str(config_file)) assert isinstance(proxy_server_module.rag_upload_malware_scanner, EicarTestMalwareScanner) + assert proxy_server_module.rag_ingest_settings.files_api_controls is True config_file.write_text( yaml.dump( diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index f63e4657a97..508e2f71248 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -28451,7 +28451,7 @@ 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 - https://docs.litellm.ai/docs/rag_ingest */ + /** @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 @@ -39055,6 +39055,12 @@ export interface components { }; /** 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 From a0866d7f0d5cb689a8edc5ec1b8376decce9c078 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 25 Sep 2026 18:49:08 -0700 Subject: [PATCH 03/14] fix(http): decode a compressed response under the bounded reader's byte cap instead of refusing it --- litellm/llms/custom_httpx/http_handler.py | 16 ++++++--- litellm/proxy/rag_endpoints/endpoints.py | 3 -- .../proxy/rag_endpoints/test_rag_endpoints.py | 9 +---- .../llms/custom_httpx/test_http_handler.py | 36 +++++++++++++++++++ 4 files changed, 48 insertions(+), 16 deletions(-) diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index f9581aa3841..216b32c9b50 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -553,8 +553,7 @@ class HTTPResponseLimitError(ValueError): pass -class HTTPResponseEncodingError(HTTPResponseLimitError): - pass +_WIRE_BODY_HEADERS: Final = frozenset({"content-encoding", "content-length"}) class MaskedHTTPStatusError(httpx.HTTPStatusError): @@ -609,6 +608,12 @@ 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) + ) + + class AsyncHTTPHandler: def __init__( self, @@ -784,8 +789,6 @@ class AsyncHTTPHandler: ) if response.is_redirect or response.is_error: return httpx.Response(response.status_code, headers=response.headers, request=response.request) - if response.headers.get("content-encoding", "identity").lower() != "identity": - raise HTTPResponseEncodingError("Response size limits require an uncompressed response") if int(response.headers.get("content-length", "0")) > max_bytes: raise HTTPResponseLimitError("Response exceeds the configured size limit") with BytesIO() as body: @@ -794,7 +797,10 @@ class AsyncHTTPHandler: raise HTTPResponseLimitError("Response exceeds the configured size limit") body.write(chunk) 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=body.getvalue(), + request=response.request, ) finally: await response.aclose() diff --git a/litellm/proxy/rag_endpoints/endpoints.py b/litellm/proxy/rag_endpoints/endpoints.py index 43a1cbfc1c7..908e702bb94 100644 --- a/litellm/proxy/rag_endpoints/endpoints.py +++ b/litellm/proxy/rag_endpoints/endpoints.py @@ -27,7 +27,6 @@ from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook i ) 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, ) @@ -420,8 +419,6 @@ async def fetch_file_url(url: str) -> httpx.Response: 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, diff --git a/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py b/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py index 4907e5a9ba9..22e4914912d 100644 --- a/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py +++ b/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py @@ -17,7 +17,7 @@ from fastapi import HTTPException from fastapi.testclient import TestClient import litellm -from litellm.llms.custom_httpx.http_handler import HTTPResponseEncodingError, HTTPResponseLimitError +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 @@ -1201,13 +1201,6 @@ class TestFileUrlUploadControls: detail = await self._expect_rejection(fetch_url, "file_too_large") assert str(MAX_UPLOAD_SIZE_BYTES) in detail["error"] - async def test_compressed_url_response_is_refused_without_calling_it_too_large(self): - async def fetch_url(url: str) -> httpx.Response: - raise HTTPResponseEncodingError("Response size limits require an uncompressed response") - - detail = await self._expect_rejection(fetch_url, None) - assert "uncompressed" 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") diff --git a/tests/unit/llms/custom_httpx/test_http_handler.py b/tests/unit/llms/custom_httpx/test_http_handler.py index 8358d15d30e..81e462bb46c 100644 --- a/tests/unit/llms/custom_httpx/test_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_http_handler.py @@ -1,5 +1,6 @@ import asyncio import gc +import gzip import io import os import pathlib @@ -21,6 +22,7 @@ from litellm.llms.custom_httpx.http_handler import ( _CLIENT_REFCOUNT_WHEN_HANDLER_IS_SOLE_REFERRER, AsyncHTTPHandler, HTTPHandler, + HTTPResponseLimitError, MaskedHTTPStatusError, _get_httpx_client, get_ssl_configuration, @@ -1773,6 +1775,40 @@ 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_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_closes_stream_on_cancellation(respx_mock, monkeypatch): monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") From 7d423b51a2aeccc4105374c9b74f426ce383123c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 25 Sep 2026 19:27:34 -0700 Subject: [PATCH 04/14] fix(http): apply the Content-Length pre-check only to identity-encoded responses --- litellm/llms/custom_httpx/http_handler.py | 8 ++++- .../test_openapi_to_mcp_generator.py | 35 ++++++++++++++++++- .../llms/custom_httpx/test_http_handler.py | 17 +++++++++ 3 files changed, 58 insertions(+), 2 deletions(-) diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 216b32c9b50..9138f68e930 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -614,6 +614,12 @@ def _headers_of_the_decoded_body(headers: httpx.Headers) -> httpx.Headers: ) +def _declared_decoded_length(headers: httpx.Headers) -> int: + if headers.get("content-encoding", "identity") != "identity": + return 0 + return int(headers.get("content-length", "0")) + + class AsyncHTTPHandler: def __init__( self, @@ -789,7 +795,7 @@ class AsyncHTTPHandler: ) if response.is_redirect or response.is_error: return httpx.Response(response.status_code, headers=response.headers, request=response.request) - if int(response.headers.get("content-length", "0")) > max_bytes: + if _declared_decoded_length(response.headers) > max_bytes: raise HTTPResponseLimitError("Response exceeds the configured size limit") with BytesIO() as body: async for chunk in response.aiter_bytes(chunk_size=65536): diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py index 6b0211c3866..262602a658f 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py @@ -1569,7 +1569,10 @@ 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"}], + ) async def test_unsafe_response_headers_reject_before_reading(self, respx_mock, monkeypatch, headers): import httpx from litellm.llms.custom_httpx.http_handler import HTTPResponseLimitError @@ -1593,6 +1596,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/llms/custom_httpx/test_http_handler.py b/tests/unit/llms/custom_httpx/test_http_handler.py index 81e462bb46c..c27f1536325 100644 --- a/tests/unit/llms/custom_httpx/test_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_http_handler.py @@ -1793,6 +1793,23 @@ async def test_bounded_get_decodes_a_compressed_body_and_reports_the_decoded_len 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") From 1e8672ded1b35a50b65dc137a8c335f82fa044c7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 25 Sep 2026 19:46:49 -0700 Subject: [PATCH 05/14] fix(http): bound the wire bytes of a compressed response under the size limit --- litellm/llms/custom_httpx/http_handler.py | 43 +++++++++++++------ .../test_openapi_to_mcp_generator.py | 6 ++- .../llms/custom_httpx/test_http_handler.py | 35 +++++++++++++++ 3 files changed, 71 insertions(+), 13 deletions(-) diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 9138f68e930..e6c5749efe6 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -8,7 +8,8 @@ import sys import threading import time import weakref -from collections.abc import AsyncIterable, Callable, Iterable, Mapping +from collections.abc import AsyncGenerator, AsyncIterable, Callable, Iterable, Mapping +from contextlib import aclosing from http.cookiejar import CookieJar, DefaultCookiePolicy from io import BytesIO from types import MappingProxyType @@ -554,6 +555,8 @@ class HTTPResponseLimitError(ValueError): _WIRE_BODY_HEADERS: Final = frozenset({"content-encoding", "content-length"}) +_ENCODED_OVERHEAD_DIVISOR: Final = 4096 +_ENCODED_OVERHEAD_BYTES: Final = 64 class MaskedHTTPStatusError(httpx.HTTPStatusError): @@ -614,10 +617,29 @@ def _headers_of_the_decoded_body(headers: httpx.Headers) -> httpx.Headers: ) -def _declared_decoded_length(headers: httpx.Headers) -> int: - if headers.get("content-encoding", "identity") != "identity": - return 0 - return int(headers.get("content-length", "0")) +def _wire_byte_limit(headers: httpx.Headers, max_bytes: int) -> int: + if headers.get("content-encoding", "identity").strip().lower() == "identity": + return max_bytes + return max_bytes + max_bytes // _ENCODED_OVERHEAD_DIVISOR + _ENCODED_OVERHEAD_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 + + +async def _decoded_within(response: httpx.Response, wire: AsyncGenerator[bytes, None], max_bytes: int) -> bytes: + decoding: Final = httpx.Response( + response.status_code, headers=response.headers, content=wire, request=response.request + ) + with BytesIO() as body: + async for chunk in decoding.aiter_bytes(chunk_size=65536): + if body.tell() + len(chunk) > max_bytes: + raise HTTPResponseLimitError("Response exceeds the configured size limit") + body.write(chunk) + return body.getvalue() class AsyncHTTPHandler: @@ -795,17 +817,14 @@ class AsyncHTTPHandler: ) if response.is_redirect or response.is_error: return httpx.Response(response.status_code, headers=response.headers, request=response.request) - if _declared_decoded_length(response.headers) > max_bytes: + wire_limit: Final = _wire_byte_limit(response.headers, max_bytes) + if int(response.headers.get("content-length", "0")) > 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=_headers_of_the_decoded_body(response.headers), - content=body.getvalue(), + content=await _decoded_within(response, wire, max_bytes), request=response.request, ) finally: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py index 262602a658f..dd6b7297e13 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py @@ -1571,7 +1571,11 @@ class TestBoundedOpenAPISpecLoading: @pytest.mark.asyncio @pytest.mark.parametrize( "headers", - [{"content-length": "1000000"}, {"content-length": "1000000", "content-encoding": "identity"}], + [ + {"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 diff --git a/tests/unit/llms/custom_httpx/test_http_handler.py b/tests/unit/llms/custom_httpx/test_http_handler.py index c27f1536325..c095ae3762f 100644 --- a/tests/unit/llms/custom_httpx/test_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_http_handler.py @@ -1826,6 +1826,41 @@ async def test_bounded_get_caps_the_decoded_size_of_a_compressed_body(respx_mock 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_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") From d6d2473668bf7ec0fa7daf3e0960554d7e37dc9d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 25 Sep 2026 20:33:35 -0700 Subject: [PATCH 06/14] fix(http): serve a body the transport already read through the bounded reader --- litellm/llms/custom_httpx/http_handler.py | 13 +++++++ .../llms/custom_httpx/test_http_handler.py | 39 +++++++++++++++++++ 2 files changed, 52 insertions(+) diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index e6c5749efe6..b75dcc050d1 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -630,6 +630,17 @@ async def _wire_bounded(response: httpx.Response, limit: int) -> AsyncGenerator[ 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, + ) + + async def _decoded_within(response: httpx.Response, wire: AsyncGenerator[bytes, None], max_bytes: int) -> bytes: decoding: Final = httpx.Response( response.status_code, headers=response.headers, content=wire, request=response.request @@ -817,6 +828,8 @@ class AsyncHTTPHandler: ) if response.is_redirect or response.is_error: return httpx.Response(response.status_code, headers=response.headers, request=response.request) + if response.is_stream_consumed: + return _already_read_within(response, max_bytes) wire_limit: Final = _wire_byte_limit(response.headers, max_bytes) if int(response.headers.get("content-length", "0")) > wire_limit: raise HTTPResponseLimitError("Response exceeds the configured size limit") diff --git a/tests/unit/llms/custom_httpx/test_http_handler.py b/tests/unit/llms/custom_httpx/test_http_handler.py index c095ae3762f..35e5728db42 100644 --- a/tests/unit/llms/custom_httpx/test_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_http_handler.py @@ -1839,6 +1839,45 @@ async def test_bounded_get_caps_the_wire_bytes_of_a_compressed_body_that_decodes 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 async def test_bounded_get_rejects_a_declared_compressed_length_over_the_wire_limit(respx_mock, monkeypatch): monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") From b234de051bace91d1f081a3a30c7bede55bf0c6c Mon Sep 17 00:00:00 2001 From: mateo Date: Sat, 26 Sep 2026 23:10:23 +0000 Subject: [PATCH 07/14] fix(http): decode a compressed response through output-bounded decoders so the size cap holds before inflation httpx.Response.aiter_bytes inflates each raw chunk in full before the bounded reader can measure it, so a small gzip, deflate, or brotli body could allocate gigabytes under a 1 MiB cap. The capped path now feeds the wire bytes through zlib.decompressobj(max_length) or brotli.Decompressor.process(output_buffer_limit) and refuses encodings it cannot bound (zstd, or br without the brotli package) Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/custom_httpx/http_handler.py | 81 +++++++++++++++++-- .../llms/custom_httpx/test_http_handler.py | 46 +++++++++++ 2 files changed, 120 insertions(+), 7 deletions(-) diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index b75dcc050d1..52cfdeac000 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -8,12 +8,13 @@ import sys import threading import time import weakref +import zlib from collections.abc import AsyncGenerator, AsyncIterable, Callable, Iterable, Mapping from contextlib import aclosing 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 @@ -61,6 +62,15 @@ try: except Exception: version = "0.0.0" +try: + from brotli import Decompressor as _BrotliInflater + from brotli import error as _BrotliError + + _BROTLI_ERRORS: Final[tuple[type[Exception], ...]] = (_BrotliError,) +except ImportError: + _BrotliInflater = None + _BROTLI_ERRORS = () + # aiohttp 3.10+ exposes a `socket_factory` kwarg on TCPConnector. Older # versions don't — detect once and skip the keep-alive wiring there. @@ -641,15 +651,72 @@ def _already_read_within(response: httpx.Response, max_bytes: int) -> httpx.Resp ) +class _BoundedDecoder(Protocol): + def decode(self, data: bytes, max_output: int) -> bytes: ... + + +class _IdentityDecoder: + def decode(self, data: bytes, max_output: int) -> bytes: + return data + + +class _ZlibDecoder: + def __init__(self, wbits: int) -> None: + self._inflate = zlib.decompressobj(wbits) + self._raw_fallback = wbits == zlib.MAX_WBITS + + def decode(self, data: bytes, max_output: int) -> bytes: + try: + head: Final = self._inflate.decompress(data, max_output) + except zlib.error as exc: + if not self._raw_fallback: + raise httpx.DecodingError(str(exc)) from exc + self._inflate = zlib.decompressobj(-zlib.MAX_WBITS) + self._raw_fallback = False + return self.decode(data, max_output) + self._raw_fallback = False + if len(head) >= max_output or not self._inflate.unconsumed_tail: + return head + return head + self.decode(self._inflate.unconsumed_tail, max_output - len(head)) + + +class _BrotliDecoder: + def __init__(self, inflate: "_BrotliInflater") -> None: + self._inflate: Final = inflate + + def decode(self, data: bytes, max_output: int) -> bytes: + try: + head: bytes = self._inflate.process(data, output_buffer_limit=max_output) + except _BROTLI_ERRORS as exc: + raise httpx.DecodingError(str(exc)) from exc + if len(head) >= max_output or self._inflate.can_accept_more_data(): + return head + return head + self.decode(b"", max_output - len(head)) + + +def _bounded_decoder(headers: httpx.Headers) -> _BoundedDecoder: + match headers.get("content-encoding", "identity").strip().lower(): + case "identity": + return _IdentityDecoder() + case "gzip" | "x-gzip": + return _ZlibDecoder(zlib.MAX_WBITS | 16) + case "deflate": + return _ZlibDecoder(zlib.MAX_WBITS) + case "br" if _BrotliInflater is not None: + return _BrotliDecoder(_BrotliInflater()) + case _: + raise HTTPResponseLimitError( + "Response size limits require an identity, gzip, deflate, or br encoded response" + ) + + async def _decoded_within(response: httpx.Response, wire: AsyncGenerator[bytes, None], max_bytes: int) -> bytes: - decoding: Final = httpx.Response( - response.status_code, headers=response.headers, content=wire, request=response.request - ) + decoder: Final = _bounded_decoder(response.headers) with BytesIO() as body: - async for chunk in decoding.aiter_bytes(chunk_size=65536): - if body.tell() + len(chunk) > max_bytes: + 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") - body.write(chunk) return body.getvalue() diff --git a/tests/unit/llms/custom_httpx/test_http_handler.py b/tests/unit/llms/custom_httpx/test_http_handler.py index 35e5728db42..918d3b6b947 100644 --- a/tests/unit/llms/custom_httpx/test_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_http_handler.py @@ -6,11 +6,14 @@ import os import pathlib import ssl import threading +import tracemalloc import weakref +import zlib from collections.abc import Callable, Mapping from typing import Final from unittest.mock import MagicMock, patch +import brotli import certifi import httpx import pytest @@ -1878,6 +1881,49 @@ async def test_bounded_get_bounds_an_already_read_body_on_its_decoded_length(mon await handler.close() +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("encoding", "compress"), + [("gzip", gzip.compress), ("deflate", zlib.compress), ("br", brotli.compress)], +) +async def test_bounded_get_never_inflates_a_compressed_body_past_the_cap( + respx_mock, monkeypatch, encoding: str, compress: Callable[[bytes], bytes] +): + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + inflated: Final = 64 * 1024 * 1024 + cap: Final = 1024 * 1024 + bomb: Final = 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 +async def test_bounded_get_refuses_an_encoding_it_cannot_decode_under_the_cap(respx_mock, monkeypatch): + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + respx_mock.get("https://cdn.example/zst").respond(200, content=b"\x28\xb5\x2f\xfd", 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") From 3c63bd87a8472b85813cb24e3f5892a94a11392a Mon Sep 17 00:00:00 2001 From: mateo Date: Sat, 26 Sep 2026 23:31:07 +0000 Subject: [PATCH 08/14] fix(http): fail truncated compressed downloads, drain brotli iteratively, and keep brotli optional in tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/custom_httpx/http_handler.py | 71 ++++++++++++------- .../llms/custom_httpx/test_http_handler.py | 30 ++++++-- 2 files changed, 73 insertions(+), 28 deletions(-) diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 52cfdeac000..5ba96f87c53 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -9,7 +9,7 @@ import threading import time import weakref import zlib -from collections.abc import AsyncGenerator, AsyncIterable, Callable, Iterable, Mapping +from collections.abc import AsyncGenerator, AsyncIterable, Callable, Iterable, Iterator, Mapping from contextlib import aclosing from http.cookiejar import CookieJar, DefaultCookiePolicy from io import BytesIO @@ -65,11 +65,11 @@ except Exception: try: from brotli import Decompressor as _BrotliInflater from brotli import error as _BrotliError - - _BROTLI_ERRORS: Final[tuple[type[Exception], ...]] = (_BrotliError,) except ImportError: _BrotliInflater = None - _BROTLI_ERRORS = () + + class _BrotliError(Exception): + """Stand-in so the except clause below type-checks; unreachable without brotli.""" # aiohttp 3.10+ exposes a `socket_factory` kwarg on TCPConnector. Older @@ -654,11 +654,17 @@ def _already_read_within(response: httpx.Response, max_bytes: int) -> httpx.Resp 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: @@ -679,35 +685,51 @@ class _ZlibDecoder: return head return head + self.decode(self._inflate.unconsumed_tail, max_output - len(head)) + def finish(self) -> None: + if not self._inflate.eof: + raise httpx.DecodingError("Compressed response ended before the end of the stream") + class _BrotliDecoder: def __init__(self, inflate: "_BrotliInflater") -> None: self._inflate: Final = inflate def decode(self, data: bytes, max_output: int) -> bytes: - try: - head: bytes = self._inflate.process(data, output_buffer_limit=max_output) - except _BROTLI_ERRORS as exc: - raise httpx.DecodingError(str(exc)) from exc - if len(head) >= max_output or self._inflate.can_accept_more_data(): - return head - return head + self.decode(b"", max_output - len(head)) + 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 _BrotliError as exc: + raise httpx.DecodingError(str(exc)) from exc + produced += len(piece) # rebind-ok: see above + pending = b"" # rebind-ok: see above + yield piece + if self._inflate.can_accept_more_data(): + 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") def _bounded_decoder(headers: httpx.Headers) -> _BoundedDecoder: - match headers.get("content-encoding", "identity").strip().lower(): - case "identity": - return _IdentityDecoder() - case "gzip" | "x-gzip": - return _ZlibDecoder(zlib.MAX_WBITS | 16) - case "deflate": - return _ZlibDecoder(zlib.MAX_WBITS) - case "br" if _BrotliInflater is not None: - return _BrotliDecoder(_BrotliInflater()) - case _: - raise HTTPResponseLimitError( - "Response size limits require an identity, gzip, deflate, or br encoded response" - ) + 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 _BrotliInflater is None else {"br": lambda: _BrotliDecoder(_BrotliInflater())}), + } + ) + factory: Final = factories.get(headers.get("content-encoding", "identity").strip().lower()) + if factory is None: + raise HTTPResponseLimitError("Response size limits require an identity, gzip, deflate, or br encoded response") + return factory() async def _decoded_within(response: httpx.Response, wire: AsyncGenerator[bytes, None], max_bytes: int) -> bytes: @@ -717,6 +739,7 @@ async def _decoded_within(response: httpx.Response, wire: AsyncGenerator[bytes, 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() diff --git a/tests/unit/llms/custom_httpx/test_http_handler.py b/tests/unit/llms/custom_httpx/test_http_handler.py index 918d3b6b947..d7c105bec74 100644 --- a/tests/unit/llms/custom_httpx/test_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_http_handler.py @@ -13,7 +13,6 @@ from collections.abc import Callable, Mapping from typing import Final from unittest.mock import MagicMock, patch -import brotli import certifi import httpx import pytest @@ -1884,15 +1883,15 @@ async def test_bounded_get_bounds_an_already_read_body_on_its_decoded_length(mon @pytest.mark.asyncio @pytest.mark.parametrize( ("encoding", "compress"), - [("gzip", gzip.compress), ("deflate", zlib.compress), ("br", brotli.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] + 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(b"\0" * inflated) + bomb: Final = (compress or pytest.importorskip("brotli").compress)(b"\0" * inflated) assert len(bomb) < cap class WireStream(httpx.AsyncByteStream): @@ -1912,6 +1911,29 @@ async def test_bounded_get_never_inflates_a_compressed_body_past_the_cap( 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_refuses_an_encoding_it_cannot_decode_under_the_cap(respx_mock, monkeypatch): monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") From 01c9dc5995f06fcaed25e3ba99a18d53ac180296 Mon Sep 17 00:00:00 2001 From: mateo Date: Sat, 26 Sep 2026 23:37:58 +0000 Subject: [PATCH 09/14] fix(http): drop unused rebind suppressions and narrow the optional brotli import for the type gates Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/custom_httpx/http_handler.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 5ba96f87c53..9cdfd912155 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -68,8 +68,7 @@ try: except ImportError: _BrotliInflater = None - class _BrotliError(Exception): - """Stand-in so the except clause below type-checks; unreachable without brotli.""" + class _BrotliError(Exception): ... # aiohttp 3.10+ exposes a `socket_factory` kwarg on TCPConnector. Older @@ -703,8 +702,8 @@ class _BrotliDecoder: piece: bytes = self._inflate.process(pending, output_buffer_limit=max_output - produced) except _BrotliError as exc: raise httpx.DecodingError(str(exc)) from exc - produced += len(piece) # rebind-ok: see above - pending = b"" # rebind-ok: see above + produced += len(piece) + pending = b"" yield piece if self._inflate.can_accept_more_data(): return @@ -717,13 +716,14 @@ class _BrotliDecoder: def _bounded_decoder(headers: httpx.Headers) -> _BoundedDecoder: + brotli_inflater: Final = _BrotliInflater 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 _BrotliInflater is None else {"br": lambda: _BrotliDecoder(_BrotliInflater())}), + **({} if brotli_inflater is None else {"br": lambda: _BrotliDecoder(brotli_inflater())}), } ) factory: Final = factories.get(headers.get("content-encoding", "identity").strip().lower()) From 2b06c56dfbf902cacceb6758ab16e4f1a4b97c4a Mon Sep 17 00:00:00 2001 From: mateo Date: Sat, 26 Sep 2026 23:46:43 +0000 Subject: [PATCH 10/14] fix(http): end the brotli drain loop on empty output or a finished stream Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/custom_httpx/http_handler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 9cdfd912155..9cc2ed15540 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -705,7 +705,7 @@ class _BrotliDecoder: produced += len(piece) pending = b"" yield piece - if self._inflate.can_accept_more_data(): + if not piece or self._inflate.is_finished(): return return b"".join(drained()) From 14328aa37c4c6db3f8ba1c5799448c52cb1f0dab Mon Sep 17 00:00:00 2001 From: mateo Date: Sun, 27 Sep 2026 01:10:21 +0000 Subject: [PATCH 11/14] fix(http): drain zlib output iteratively and stream the zstd refusal fixture Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/custom_httpx/http_handler.py | 33 ++++++++++++------- .../llms/custom_httpx/test_http_handler.py | 7 +++- 2 files changed, 28 insertions(+), 12 deletions(-) diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 9cc2ed15540..c2624407132 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -670,19 +670,30 @@ class _ZlibDecoder: self._inflate = zlib.decompressobj(wbits) self._raw_fallback = wbits == zlib.MAX_WBITS - def decode(self, data: bytes, max_output: int) -> bytes: + def _inflate_once(self, data: bytes, max_output: int) -> bytes: try: - head: Final = self._inflate.decompress(data, max_output) - except zlib.error as exc: + return self._inflate.decompress(data, max_output) + except zlib.error as zlib_wrapped: if not self._raw_fallback: - raise httpx.DecodingError(str(exc)) from exc - self._inflate = zlib.decompressobj(-zlib.MAX_WBITS) - self._raw_fallback = False - return self.decode(data, max_output) - self._raw_fallback = False - if len(head) >= max_output or not self._inflate.unconsumed_tail: - return head - return head + self.decode(self._inflate.unconsumed_tail, max_output - len(head)) + 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: diff --git a/tests/unit/llms/custom_httpx/test_http_handler.py b/tests/unit/llms/custom_httpx/test_http_handler.py index d7c105bec74..7fa0ae80f2d 100644 --- a/tests/unit/llms/custom_httpx/test_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_http_handler.py @@ -1937,7 +1937,12 @@ async def test_bounded_get_decodes_a_compressed_body_under_the_cap_and_rejects_a @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") - respx_mock.get("https://cdn.example/zst").respond(200, content=b"\x28\xb5\x2f\xfd", headers={"content-encoding": "zstd"}) + + 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"): From c5252b4bea1809f042679925f3837c03f7c297bc Mon Sep 17 00:00:00 2001 From: mateo Date: Sun, 27 Sep 2026 02:14:18 +0000 Subject: [PATCH 12/14] fix(http_handler): undo stacked Content-Encoding values on the size-capped path Comma-separated encodings are decoded last applied first with every intermediate stage bounded, so gzip, gzip downloads work again while the cap still holds. Unsupported encodings raise HTTPResponseEncodingError, which /v1/rag/ingest reports as a fetch error instead of file_too_large Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/custom_httpx/http_handler.py | 54 ++++++++++++++++--- litellm/proxy/rag_endpoints/endpoints.py | 3 ++ .../llms/custom_httpx/test_http_handler.py | 33 ++++++++++++ 3 files changed, 82 insertions(+), 8 deletions(-) diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index c2624407132..af466706aad 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -9,7 +9,7 @@ import threading import time import weakref import zlib -from collections.abc import AsyncGenerator, AsyncIterable, Callable, Iterable, Iterator, Mapping +from collections.abc import AsyncGenerator, AsyncIterable, Callable, Iterable, Iterator, Mapping, Sequence from contextlib import aclosing from http.cookiejar import CookieJar, DefaultCookiePolicy from io import BytesIO @@ -563,6 +563,10 @@ 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 @@ -626,12 +630,22 @@ def _headers_of_the_decoded_body(headers: httpx.Headers) -> httpx.Headers: ) -def _wire_byte_limit(headers: httpx.Headers, max_bytes: int) -> int: - if headers.get("content-encoding", "identity").strip().lower() == "identity": - return max_bytes +def _encodings_of(headers: httpx.Headers) -> tuple[str, ...]: + declared: Final[str] = headers.get("content-encoding", "identity") + names: Final = tuple(name.strip().lower() for name in 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: @@ -726,6 +740,27 @@ class _BrotliDecoder: raise httpx.DecodingError("Compressed response ended before the end of the stream") +class _ChainedDecoder: + """Undo stacked Content-Encoding values, last applied first, bounding every intermediate stage.""" + + def __init__(self, stages: Sequence[_BoundedDecoder]) -> None: + self._stages: Final = tuple(stages) + + def decode(self, data: bytes, max_output: int) -> bytes: + intermediate_cap: Final = _encoded_byte_limit(max_output) + *outer, last = self._stages + passed = data # rebind-ok: each stage's output feeds the next stage + for stage in outer: + passed = stage.decode(passed, intermediate_cap + 1) + if len(passed) > intermediate_cap: + raise HTTPResponseLimitError("Response exceeds the configured size limit") + return last.decode(passed, max_output) + + def finish(self) -> None: + for stage in self._stages: + stage.finish() + + def _bounded_decoder(headers: httpx.Headers) -> _BoundedDecoder: brotli_inflater: Final = _BrotliInflater factories: Final[Mapping[str, Callable[[], _BoundedDecoder]]] = MappingProxyType( @@ -737,10 +772,13 @@ def _bounded_decoder(headers: httpx.Headers) -> _BoundedDecoder: **({} if brotli_inflater is None else {"br": lambda: _BrotliDecoder(brotli_inflater())}), } ) - factory: Final = factories.get(headers.get("content-encoding", "identity").strip().lower()) - if factory is None: - raise HTTPResponseLimitError("Response size limits require an identity, gzip, deflate, or br encoded response") - return factory() + 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) async def _decoded_within(response: httpx.Response, wire: AsyncGenerator[bytes, None], max_bytes: int) -> bytes: diff --git a/litellm/proxy/rag_endpoints/endpoints.py b/litellm/proxy/rag_endpoints/endpoints.py index 908e702bb94..43a1cbfc1c7 100644 --- a/litellm/proxy/rag_endpoints/endpoints.py +++ b/litellm/proxy/rag_endpoints/endpoints.py @@ -27,6 +27,7 @@ from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook i ) 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, ) @@ -419,6 +420,8 @@ async def fetch_file_url(url: str) -> httpx.Response: 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, diff --git a/tests/unit/llms/custom_httpx/test_http_handler.py b/tests/unit/llms/custom_httpx/test_http_handler.py index 7fa0ae80f2d..93380548c1d 100644 --- a/tests/unit/llms/custom_httpx/test_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_http_handler.py @@ -1934,6 +1934,39 @@ async def test_bounded_get_decodes_a_compressed_body_under_the_cap_and_rejects_a 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_refuses_an_encoding_it_cannot_decode_under_the_cap(respx_mock, monkeypatch): monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") From 3e188f3309c0bc4107e5fc2926d6a2cba6508e15 Mon Sep 17 00:00:00 2001 From: mateo Date: Sun, 27 Sep 2026 02:53:41 +0000 Subject: [PATCH 13/14] fix(http_handler): cap each stacked decoding stage across the whole response, not per wire chunk Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/custom_httpx/http_handler.py | 46 +++++++++++++------ .../llms/custom_httpx/test_http_handler.py | 43 ++++++++++++++--- 2 files changed, 70 insertions(+), 19 deletions(-) diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index af466706aad..a93c53594e3 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -740,28 +740,48 @@ class _BrotliDecoder: 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]) -> None: - self._stages: Final = tuple(stages) + 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: - intermediate_cap: Final = _encoded_byte_limit(max_output) - *outer, last = self._stages passed = data # rebind-ok: each stage's output feeds the next stage - for stage in outer: - passed = stage.decode(passed, intermediate_cap + 1) - if len(passed) > intermediate_cap: - raise HTTPResponseLimitError("Response exceeds the configured size limit") - return last.decode(passed, max_output) + for stage in self._outer: + passed = stage.decode(passed) + return self._last.decode(passed, max_output) def finish(self) -> None: - for stage in self._stages: + for stage in self._outer: stage.finish() + self._last.finish() -def _bounded_decoder(headers: httpx.Headers) -> _BoundedDecoder: +def _bounded_decoder(headers: httpx.Headers, max_bytes: int) -> _BoundedDecoder: brotli_inflater: Final = _BrotliInflater factories: Final[Mapping[str, Callable[[], _BoundedDecoder]]] = MappingProxyType( { @@ -778,11 +798,11 @@ def _bounded_decoder(headers: httpx.Headers) -> _BoundedDecoder: "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) + 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) + 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())) diff --git a/tests/unit/llms/custom_httpx/test_http_handler.py b/tests/unit/llms/custom_httpx/test_http_handler.py index 93380548c1d..18406a0343e 100644 --- a/tests/unit/llms/custom_httpx/test_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_http_handler.py @@ -1443,7 +1443,6 @@ async def test_connection_error_retry_forwards_content(method: str): await handler.close() - @pytest.fixture def forward_proxy_server(): """Plain HTTP forward proxy that records the absolute URIs it is asked to fetch.""" @@ -1574,9 +1573,7 @@ def private_ca_tls_upstream(tmp_path: pathlib.Path): ca_pem.write_bytes(cert.public_bytes(serialization.Encoding.PEM)) key_pem = tmp_path / "key.pem" key_pem.write_bytes( - key.private_bytes( - serialization.Encoding.PEM, serialization.PrivateFormat.PKCS8, serialization.NoEncryption() - ) + key.private_bytes(serialization.Encoding.PEM, serialization.PrivateFormat.PKCS8, serialization.NoEncryption()) ) class OkTlsHandler(BaseHTTPRequestHandler): @@ -1751,8 +1748,11 @@ async def test_bounded_get_preserves_sdk_redirect_auth_and_query_handling(respx_ handler = AsyncHTTPHandler() try: response = await handler.get( - "https://example.com/spec.json?original=1", max_response_bytes=100, follow_redirects=True, - headers={"Authorization": "Bearer sentinel", "Accept-Encoding": "gzip"}, timeout=2.0, + "https://example.com/spec.json?original=1", + max_response_bytes=100, + follow_redirects=True, + headers={"Authorization": "Bearer sentinel", "Accept-Encoding": "gzip"}, + timeout=2.0, ) finally: await handler.close() @@ -1967,6 +1967,37 @@ async def test_bounded_get_undoes_stacked_encodings_and_still_caps_the_decoded_s 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") From c6206d3f5947031ad4beef434db2afa885a77eb5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 3 Nov 2026 11:25:10 -0800 Subject: [PATCH 14/14] test(rag): stub the file_url fetcher in JSON ingest tests that reach it after the main merge --- .../unit/proxy/rag_endpoints/test_rag_endpoints.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py b/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py index a42ef269bb8..d2e096668ed 100644 --- a/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py +++ b/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py @@ -499,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"}, @@ -525,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"},