mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
Merge pull request #38135 from BerriAI/litellm_lit_5704_vectorstore_upload_controls
feat(proxy): enforce vector-store upload security controls on /v1/rag/ingest
This commit is contained in:
Mateo Wang
GitHub
commit
1e9cecd05f
5 changed files with 593 additions and 3 deletions
|
|
@ -27,6 +27,13 @@ from litellm.proxy.common_utils.http_parsing_utils import (
|
|||
_safe_get_request_headers,
|
||||
get_form_data,
|
||||
)
|
||||
from litellm.proxy.rag_endpoints.upload_security import (
|
||||
MAX_UPLOAD_SIZE_BYTES,
|
||||
EicarTestMalwareScanner,
|
||||
MalwareScanner,
|
||||
RejectedUpload,
|
||||
validate_upload,
|
||||
)
|
||||
from litellm.proxy.vector_store_endpoints.utils import (
|
||||
assert_user_can_access_vector_store_id,
|
||||
)
|
||||
|
|
@ -287,8 +294,22 @@ 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)
|
||||
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
|
||||
|
||||
|
||||
async def parse_rag_ingest_request(
|
||||
request: Request,
|
||||
scanner: MalwareScanner,
|
||||
) -> tuple[dict[str, Any], tuple[str, bytes, str] | None, str | None, str | None]:
|
||||
"""
|
||||
Parse RAG ingest request.
|
||||
|
|
@ -297,6 +318,11 @@ async def parse_rag_ingest_request(
|
|||
- Form: file + request JSON in form field
|
||||
- JSON body for URL-based ingestion
|
||||
|
||||
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.
|
||||
|
||||
Returns:
|
||||
Tuple of (ingest_options, file_data, file_url, file_id)
|
||||
"""
|
||||
|
|
@ -315,7 +341,7 @@ async def parse_rag_ingest_request(
|
|||
# Get file
|
||||
file_obj = form_data.get("file")
|
||||
if file_obj is not None and hasattr(file_obj, "read"):
|
||||
file_content = await file_obj.read()
|
||||
file_content = await file_obj.read(MAX_UPLOAD_SIZE_BYTES + 1)
|
||||
file_data = (file_obj.filename, file_content, file_obj.content_type)
|
||||
|
||||
# Parse JSON from 'request' form field (contains full request body as JSON)
|
||||
|
|
@ -357,6 +383,10 @@ 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,
|
||||
|
|
@ -398,7 +428,7 @@ async def parse_rag_ingest_request(
|
|||
},
|
||||
)
|
||||
|
||||
return ingest_options, file_data, file_url, file_id
|
||||
return ingest_options, secured_file_data, file_url, file_id
|
||||
|
||||
|
||||
@router.post(
|
||||
|
|
@ -461,7 +491,9 @@ async def rag_ingest(
|
|||
|
||||
try:
|
||||
# Parse request
|
||||
ingest_options, file_data, file_url, file_id = await parse_rag_ingest_request(request)
|
||||
ingest_options, file_data, file_url, file_id = await parse_rag_ingest_request(
|
||||
request, scanner=EicarTestMalwareScanner()
|
||||
)
|
||||
|
||||
# INTERNAL_USER_VIEW_ONLY can ingest to existing vector stores only
|
||||
if user_api_key_dict.user_role == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value and not ingest_options.get(
|
||||
|
|
|
|||
289
litellm/proxy/rag_endpoints/upload_security.py
Normal file
289
litellm/proxy/rag_endpoints/upload_security.py
Normal file
|
|
@ -0,0 +1,289 @@
|
|||
"""Security controls for vector-store file uploads.
|
||||
|
||||
Content is classified by inspecting its actual bytes (magic signatures and a
|
||||
strict UTF-8 decode), never by trusting the client-supplied filename or
|
||||
content-type. Uploads are restricted to an allowlist of non-executable formats,
|
||||
capped in size, screened for archives, and passed through a dependency-injected
|
||||
malware scanner before they are accepted. Accepted uploads are given a
|
||||
server-generated filename so the client-controlled name never reaches storage.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import uuid
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
from types import MappingProxyType
|
||||
from typing import Final, Protocol, TypeAlias, runtime_checkable
|
||||
|
||||
from typing_extensions import assert_never
|
||||
|
||||
MAX_UPLOAD_SIZE_BYTES: Final = 512 * 1024 * 1024
|
||||
|
||||
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, ...]] = (
|
||||
b"PK\x03\x04",
|
||||
b"PK\x05\x06",
|
||||
b"PK\x07\x08",
|
||||
b"\x1f\x8b",
|
||||
b"\xfd7zXZ\x00",
|
||||
b"7z\xbc\xaf\x27\x1c",
|
||||
b"Rar!\x1a\x07\x00",
|
||||
b"Rar!\x1a\x07\x01\x00",
|
||||
b"\x04\x22\x4d\x18",
|
||||
b"\x28\xb5\x2f\xfd",
|
||||
)
|
||||
|
||||
_ARCHIVE_MAGIC_PREFIXES_ASCII_AMBIGUOUS: Final[tuple[bytes, ...]] = (b"BZh",)
|
||||
|
||||
_EXECUTABLE_MAGIC_PREFIXES: Final[tuple[bytes, ...]] = (
|
||||
b"\x7fELF",
|
||||
b"\xca\xfe\xba\xbe",
|
||||
b"\xfe\xed\xfa\xce",
|
||||
b"\xfe\xed\xfa\xcf",
|
||||
b"\xce\xfa\xed\xfe",
|
||||
b"\xcf\xfa\xed\xfe",
|
||||
b"\x00asm",
|
||||
)
|
||||
|
||||
_EXECUTABLE_MAGIC_PREFIXES_ASCII_AMBIGUOUS: Final[tuple[bytes, ...]] = (b"MZ", b"dex\n")
|
||||
|
||||
_TAR_USTAR_MAGIC: Final = b"ustar"
|
||||
_TAR_USTAR_OFFSET: Final = 257
|
||||
|
||||
_UTF8_BOM: Final = b"\xef\xbb\xbf"
|
||||
|
||||
|
||||
class DetectedFormat(str, Enum):
|
||||
PDF = "pdf"
|
||||
TEXT = "text"
|
||||
|
||||
|
||||
class DisallowedKind(str, Enum):
|
||||
ARCHIVE = "archive"
|
||||
EXECUTABLE = "executable"
|
||||
UNKNOWN_BINARY = "unknown_binary"
|
||||
|
||||
|
||||
class RejectionReason(str, Enum):
|
||||
EMPTY_FILE = "empty_file"
|
||||
FILE_TOO_LARGE = "file_too_large"
|
||||
ARCHIVE_NOT_ALLOWED = "archive_not_allowed"
|
||||
EXECUTABLE_NOT_ALLOWED = "executable_not_allowed"
|
||||
UNSUPPORTED_FORMAT = "unsupported_format"
|
||||
MALWARE_DETECTED = "malware_detected"
|
||||
MALWARE_SCAN_ERROR = "malware_scan_error"
|
||||
|
||||
|
||||
class ScanVerdict(str, Enum):
|
||||
CLEAN = "clean"
|
||||
INFECTED = "infected"
|
||||
ERROR = "error"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ScanResult:
|
||||
verdict: ScanVerdict
|
||||
signature: str | None = None
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class MalwareScanner(Protocol):
|
||||
def scan(self, content: bytes) -> ScanResult: ...
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class EicarTestMalwareScanner:
|
||||
"""Placeholder scanner that only flags the EICAR anti-malware test file.
|
||||
|
||||
It exists to prove the scan hook is wired end to end and to satisfy the
|
||||
EICAR retest; it provides no real protection. Inject a scanner backed by a
|
||||
real engine through the ``scanner`` parameter of :func:`validate_upload` to
|
||||
screen production uploads.
|
||||
"""
|
||||
|
||||
def scan(self, content: bytes) -> ScanResult:
|
||||
if EICAR_TEST_SIGNATURE in content:
|
||||
return ScanResult(verdict=ScanVerdict.INFECTED, signature="EICAR-STANDARD-ANTIVIRUS-TEST-FILE")
|
||||
return ScanResult(verdict=ScanVerdict.CLEAN)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class AllowedContent:
|
||||
format: DetectedFormat
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class DisallowedContent:
|
||||
kind: DisallowedKind
|
||||
|
||||
|
||||
ContentInspection: TypeAlias = AllowedContent | DisallowedContent
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class SecuredUpload:
|
||||
safe_filename: str
|
||||
content_type: str
|
||||
detected_format: DetectedFormat
|
||||
size_bytes: int
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class RejectedUpload:
|
||||
reason: RejectionReason
|
||||
message: str
|
||||
|
||||
|
||||
UploadValidation: TypeAlias = SecuredUpload | RejectedUpload
|
||||
|
||||
_SAFE_EXTENSION: Final[Mapping[DetectedFormat, str]] = MappingProxyType(
|
||||
{
|
||||
DetectedFormat.PDF: "pdf",
|
||||
DetectedFormat.TEXT: "txt",
|
||||
}
|
||||
)
|
||||
|
||||
_SAFE_CONTENT_TYPE: Final[Mapping[DetectedFormat, str]] = MappingProxyType(
|
||||
{
|
||||
DetectedFormat.PDF: "application/pdf",
|
||||
DetectedFormat.TEXT: "text/plain",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _starts_with_any(content: bytes, prefixes: tuple[bytes, ...]) -> bool:
|
||||
return any(content.startswith(prefix) for prefix in prefixes)
|
||||
|
||||
|
||||
def _is_archive(content: bytes) -> bool:
|
||||
if _starts_with_any(content, _ARCHIVE_MAGIC_PREFIXES):
|
||||
return True
|
||||
tar_magic_end: Final = _TAR_USTAR_OFFSET + len(_TAR_USTAR_MAGIC)
|
||||
if len(content) >= tar_magic_end and content[_TAR_USTAR_OFFSET:tar_magic_end] == _TAR_USTAR_MAGIC:
|
||||
return True
|
||||
return _starts_with_any(content, _ARCHIVE_MAGIC_PREFIXES_ASCII_AMBIGUOUS) and not _is_utf8_text(content)
|
||||
|
||||
|
||||
def _is_utf8_text(content: bytes) -> bool:
|
||||
if b"\x00" in content:
|
||||
return False
|
||||
try:
|
||||
content.decode("utf-8")
|
||||
except UnicodeDecodeError:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def _is_executable_binary(content: bytes) -> bool:
|
||||
if _starts_with_any(content, _EXECUTABLE_MAGIC_PREFIXES):
|
||||
return True
|
||||
return _starts_with_any(content, _EXECUTABLE_MAGIC_PREFIXES_ASCII_AMBIGUOUS) and not _is_utf8_text(content)
|
||||
|
||||
|
||||
def _looks_like_shebang(content: bytes) -> bool:
|
||||
body: Final = content.removeprefix(_UTF8_BOM).lstrip()
|
||||
return body.startswith(b"#!")
|
||||
|
||||
|
||||
def inspect_content(content: bytes) -> ContentInspection:
|
||||
if _looks_like_shebang(content):
|
||||
return DisallowedContent(DisallowedKind.EXECUTABLE)
|
||||
if content.startswith(b"%PDF-"):
|
||||
return AllowedContent(DetectedFormat.PDF)
|
||||
if _is_archive(content):
|
||||
return DisallowedContent(DisallowedKind.ARCHIVE)
|
||||
if _is_executable_binary(content):
|
||||
return DisallowedContent(DisallowedKind.EXECUTABLE)
|
||||
if _is_utf8_text(content):
|
||||
return AllowedContent(DetectedFormat.TEXT)
|
||||
return DisallowedContent(DisallowedKind.UNKNOWN_BINARY)
|
||||
|
||||
|
||||
def generate_safe_filename(detected_format: DetectedFormat) -> str:
|
||||
return f"{uuid.uuid4().hex}.{_SAFE_EXTENSION[detected_format]}"
|
||||
|
||||
|
||||
def _reject_disallowed(kind: DisallowedKind) -> RejectedUpload:
|
||||
match kind:
|
||||
case DisallowedKind.ARCHIVE:
|
||||
return RejectedUpload(
|
||||
RejectionReason.ARCHIVE_NOT_ALLOWED,
|
||||
"Archive uploads are not allowed.",
|
||||
)
|
||||
case DisallowedKind.EXECUTABLE:
|
||||
return RejectedUpload(
|
||||
RejectionReason.EXECUTABLE_NOT_ALLOWED,
|
||||
"Executable uploads are not allowed.",
|
||||
)
|
||||
case DisallowedKind.UNKNOWN_BINARY:
|
||||
return RejectedUpload(
|
||||
RejectionReason.UNSUPPORTED_FORMAT,
|
||||
"Only PDF and UTF-8 text documents are accepted.",
|
||||
)
|
||||
assert_never(kind)
|
||||
|
||||
|
||||
def _scan_rejection(content: bytes, scanner: MalwareScanner) -> RejectedUpload | None:
|
||||
result: Final = scanner.scan(content)
|
||||
match result.verdict:
|
||||
case ScanVerdict.CLEAN:
|
||||
return None
|
||||
case ScanVerdict.INFECTED:
|
||||
return RejectedUpload(
|
||||
RejectionReason.MALWARE_DETECTED,
|
||||
f"Uploaded file was flagged by malware scanning ({result.signature or 'unknown signature'}).",
|
||||
)
|
||||
case ScanVerdict.ERROR:
|
||||
return RejectedUpload(
|
||||
RejectionReason.MALWARE_SCAN_ERROR,
|
||||
"Malware scanning could not complete; upload rejected.",
|
||||
)
|
||||
assert_never(result.verdict)
|
||||
|
||||
|
||||
def validate_upload(
|
||||
*,
|
||||
content: bytes,
|
||||
scanner: MalwareScanner,
|
||||
max_size_bytes: int = MAX_UPLOAD_SIZE_BYTES,
|
||||
) -> UploadValidation:
|
||||
size: Final = len(content)
|
||||
if size == 0:
|
||||
return RejectedUpload(RejectionReason.EMPTY_FILE, "Uploaded file is empty.")
|
||||
if size > max_size_bytes:
|
||||
return RejectedUpload(
|
||||
RejectionReason.FILE_TOO_LARGE,
|
||||
f"Uploaded file is {size} bytes, exceeding the {max_size_bytes}-byte limit.",
|
||||
)
|
||||
|
||||
inspection: Final = inspect_content(content)
|
||||
if isinstance(inspection, DisallowedContent):
|
||||
return _reject_disallowed(inspection.kind)
|
||||
|
||||
scan_rejection: Final = _scan_rejection(content, scanner)
|
||||
if scan_rejection is not None:
|
||||
return scan_rejection
|
||||
|
||||
return SecuredUpload(
|
||||
safe_filename=generate_safe_filename(inspection.format),
|
||||
content_type=_SAFE_CONTENT_TYPE[inspection.format],
|
||||
detected_format=inspection.format,
|
||||
size_bytes=size,
|
||||
)
|
||||
|
||||
|
||||
def _sanitize_header_filename(filename: str) -> str:
|
||||
stripped: Final = "".join(char for char in filename if char not in '"\\\r\n').strip()
|
||||
return stripped or "download"
|
||||
|
||||
|
||||
def safe_download_headers(filename: str) -> Mapping[str, str]:
|
||||
return MappingProxyType(
|
||||
{
|
||||
"Content-Disposition": f'attachment; filename="{_sanitize_header_filename(filename)}"',
|
||||
"X-Content-Type-Options": "nosniff",
|
||||
}
|
||||
)
|
||||
|
|
@ -17,6 +17,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import (
|
|||
handle_model_based_routing,
|
||||
prepare_data_with_credentials,
|
||||
)
|
||||
from litellm.proxy.rag_endpoints.upload_security import safe_download_headers
|
||||
from litellm.proxy.vector_store_endpoints.utils import (
|
||||
assert_user_can_access_vector_store_id,
|
||||
is_allowed_to_call_vector_store_files_endpoint,
|
||||
|
|
@ -885,6 +886,9 @@ async def vector_store_file_content(
|
|||
if original_managed_file_id:
|
||||
response = _replace_file_id_in_response(response, original_managed_file_id)
|
||||
|
||||
for header_name, header_value in safe_download_headers(file_id).items():
|
||||
fastapi_response.headers[header_name] = header_value
|
||||
|
||||
return response
|
||||
except Exception as e: # noqa: BLE001
|
||||
raise await processor._handle_llm_api_exception(
|
||||
|
|
|
|||
|
|
@ -322,3 +322,92 @@ def test_rag_query_stream_returns_event_stream(client_internal_user):
|
|||
assert response.headers.get("content-type", "").startswith("text/event-stream")
|
||||
assert '"object":"chat.completion.chunk"' in response.text
|
||||
assert "data: [DONE]" in response.text
|
||||
|
||||
|
||||
EICAR = r"X5O!P%@AP[4\PZX54(P^)7CC)7}$EICAR-STANDARD-ANTIVIRUS-TEST-FILE!$H+H*"
|
||||
INGEST_REQUEST = '{"ingest_options":{"vector_store":{"custom_llm_provider":"openai"}}}'
|
||||
|
||||
|
||||
def _multipart_ingest_request(*, filename: str, content: bytes, content_type: str):
|
||||
from starlette.requests import Request
|
||||
|
||||
boundary = "litellmuploadtestboundary"
|
||||
head = (
|
||||
f"--{boundary}\r\n"
|
||||
f'Content-Disposition: form-data; name="file"; filename="{filename}"\r\n'
|
||||
f"Content-Type: {content_type}\r\n\r\n"
|
||||
).encode()
|
||||
tail = (
|
||||
f"\r\n--{boundary}\r\n"
|
||||
f'Content-Disposition: form-data; name="request"\r\n\r\n'
|
||||
f"{INGEST_REQUEST}\r\n"
|
||||
f"--{boundary}--\r\n"
|
||||
).encode()
|
||||
body = head + content + tail
|
||||
scope = {
|
||||
"type": "http",
|
||||
"method": "POST",
|
||||
"path": "/v1/rag/ingest",
|
||||
"headers": [
|
||||
(b"content-type", f"multipart/form-data; boundary={boundary}".encode()),
|
||||
(b"content-length", str(len(body)).encode()),
|
||||
],
|
||||
"state": {},
|
||||
}
|
||||
|
||||
async def receive():
|
||||
return {"type": "http.request", "body": body, "more_body": False}
|
||||
|
||||
return Request(scope, receive)
|
||||
|
||||
|
||||
class TestVectorStoreUploadControls:
|
||||
"""End-to-end enforcement of pentest M4 upload controls on /v1/rag/ingest."""
|
||||
|
||||
def test_eicar_upload_blocked_by_malware_scanner(self, client_internal_user):
|
||||
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 == 400, response.text
|
||||
assert response.json()["detail"]["reason"] == "malware_detected"
|
||||
|
||||
def test_executable_upload_rejected(self, client_internal_user):
|
||||
elf = b"\x7fELF\x02\x01\x01\x00" + b"\x00" * 40
|
||||
response = client_internal_user.post(
|
||||
"/v1/rag/ingest",
|
||||
files={"file": ("doc.txt", io.BytesIO(elf), "text/plain")},
|
||||
data={"request": INGEST_REQUEST},
|
||||
)
|
||||
assert response.status_code == 400, response.text
|
||||
assert response.json()["detail"]["reason"] == "executable_not_allowed"
|
||||
|
||||
def test_zip_archive_upload_rejected(self, client_internal_user):
|
||||
response = client_internal_user.post(
|
||||
"/v1/rag/ingest",
|
||||
files={"file": ("doc.pdf", io.BytesIO(b"PK\x03\x04\x14\x00\x00\x00payload"), "application/pdf")},
|
||||
data={"request": INGEST_REQUEST},
|
||||
)
|
||||
assert response.status_code == 400, response.text
|
||||
assert response.json()["detail"]["reason"] == "archive_not_allowed"
|
||||
|
||||
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
|
||||
|
||||
request = _multipart_ingest_request(
|
||||
filename="../../etc/passwd",
|
||||
content=b"benign document text\n",
|
||||
content_type="text/plain",
|
||||
)
|
||||
_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"
|
||||
assert "/" not in server_filename and "\\" not in server_filename
|
||||
assert server_filename.endswith(".txt")
|
||||
assert secured_content_type == "text/plain"
|
||||
assert content_bytes == b"benign document text\n"
|
||||
|
|
|
|||
176
tests/test_litellm/proxy/rag_endpoints/test_upload_security.py
Normal file
176
tests/test_litellm/proxy/rag_endpoints/test_upload_security.py
Normal file
|
|
@ -0,0 +1,176 @@
|
|||
"""Unit tests for vector-store upload security controls.
|
||||
|
||||
These pin the pentest M4 remediation: an allowlist enforced by real content
|
||||
inspection (not extension/mime trust), a size cap, archive and executable
|
||||
rejection, server-generated filenames, safe download headers, and a
|
||||
dependency-injected malware scanner validated with the EICAR test file.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.proxy.rag_endpoints.upload_security import (
|
||||
EICAR_TEST_SIGNATURE,
|
||||
DetectedFormat,
|
||||
EicarTestMalwareScanner,
|
||||
RejectedUpload,
|
||||
RejectionReason,
|
||||
ScanResult,
|
||||
ScanVerdict,
|
||||
SecuredUpload,
|
||||
generate_safe_filename,
|
||||
inspect_content,
|
||||
safe_download_headers,
|
||||
validate_upload,
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _StubScanner:
|
||||
result: ScanResult
|
||||
|
||||
def scan(self, content: bytes) -> ScanResult:
|
||||
return self.result
|
||||
|
||||
|
||||
_CLEAN_SCANNER = _StubScanner(ScanResult(ScanVerdict.CLEAN))
|
||||
_INFECTED_SCANNER = _StubScanner(ScanResult(ScanVerdict.INFECTED, signature="Test.Sig"))
|
||||
_ERROR_SCANNER = _StubScanner(ScanResult(ScanVerdict.ERROR))
|
||||
|
||||
_PDF_BYTES = b"%PDF-1.7\n1 0 obj<<>>endobj\n"
|
||||
_TEXT_BYTES = "the quick brown fox\n".encode("utf-8")
|
||||
_ELF_BYTES = b"\x7fELF\x02\x01\x01\x00" + b"\x00" * 32
|
||||
_PE_BYTES = b"MZ\x90\x00\x03\x00\x00\x00\x04\x00\x00\x00"
|
||||
_ZIP_BYTES = b"PK\x03\x04\x14\x00\x00\x00"
|
||||
_GZIP_BYTES = b"\x1f\x8b\x08\x00\x00\x00\x00\x00"
|
||||
_SHEBANG_BYTES = b"#!/bin/bash\nrm -rf /\n"
|
||||
|
||||
|
||||
def _tar_bytes() -> bytes:
|
||||
header = bytearray(512)
|
||||
header[257:262] = b"ustar"
|
||||
return bytes(header)
|
||||
|
||||
|
||||
def _expect_rejected(content: bytes, reason: RejectionReason, *, max_size_bytes: int = 512 * 1024 * 1024) -> None:
|
||||
result = validate_upload(content=content, scanner=_CLEAN_SCANNER, max_size_bytes=max_size_bytes)
|
||||
assert isinstance(result, RejectedUpload), f"expected rejection, got {result!r}"
|
||||
assert result.reason is reason, f"expected {reason}, got {result.reason}"
|
||||
|
||||
|
||||
def test_empty_file_rejected():
|
||||
_expect_rejected(b"", RejectionReason.EMPTY_FILE)
|
||||
|
||||
|
||||
def test_oversized_file_rejected():
|
||||
_expect_rejected(b"%PDF-" + b"a" * 100, RejectionReason.FILE_TOO_LARGE, max_size_bytes=10)
|
||||
|
||||
|
||||
def test_zip_archive_rejected():
|
||||
_expect_rejected(_ZIP_BYTES, RejectionReason.ARCHIVE_NOT_ALLOWED)
|
||||
|
||||
|
||||
def test_gzip_archive_rejected():
|
||||
_expect_rejected(_GZIP_BYTES, RejectionReason.ARCHIVE_NOT_ALLOWED)
|
||||
|
||||
|
||||
def test_tar_archive_rejected():
|
||||
_expect_rejected(_tar_bytes(), RejectionReason.ARCHIVE_NOT_ALLOWED)
|
||||
|
||||
|
||||
def test_elf_executable_rejected():
|
||||
_expect_rejected(_ELF_BYTES, RejectionReason.EXECUTABLE_NOT_ALLOWED)
|
||||
|
||||
|
||||
def test_windows_pe_executable_rejected():
|
||||
_expect_rejected(_PE_BYTES, RejectionReason.EXECUTABLE_NOT_ALLOWED)
|
||||
|
||||
|
||||
def test_shebang_script_rejected():
|
||||
_expect_rejected(_SHEBANG_BYTES, RejectionReason.EXECUTABLE_NOT_ALLOWED)
|
||||
|
||||
|
||||
def test_unknown_binary_rejected():
|
||||
_expect_rejected(b"\x89\x01\x02\x00\xff\xfe garbage", RejectionReason.UNSUPPORTED_FORMAT)
|
||||
|
||||
|
||||
def test_pdf_accepted_with_server_filename_and_content_type():
|
||||
result = validate_upload(content=_PDF_BYTES, scanner=_CLEAN_SCANNER)
|
||||
assert isinstance(result, SecuredUpload)
|
||||
assert result.detected_format is DetectedFormat.PDF
|
||||
assert result.content_type == "application/pdf"
|
||||
assert result.safe_filename.endswith(".pdf")
|
||||
assert result.size_bytes == len(_PDF_BYTES)
|
||||
|
||||
|
||||
def test_utf8_text_accepted():
|
||||
result = validate_upload(content=_TEXT_BYTES, scanner=_CLEAN_SCANNER)
|
||||
assert isinstance(result, SecuredUpload)
|
||||
assert result.detected_format is DetectedFormat.TEXT
|
||||
assert result.content_type == "text/plain"
|
||||
assert result.safe_filename.endswith(".txt")
|
||||
|
||||
|
||||
def test_inspect_content_classifies_directly():
|
||||
from litellm.proxy.rag_endpoints.upload_security import AllowedContent, DisallowedContent, DisallowedKind
|
||||
|
||||
assert inspect_content(_PDF_BYTES) == AllowedContent(DetectedFormat.PDF)
|
||||
assert inspect_content(_TEXT_BYTES) == AllowedContent(DetectedFormat.TEXT)
|
||||
assert inspect_content(_ZIP_BYTES) == DisallowedContent(DisallowedKind.ARCHIVE)
|
||||
assert inspect_content(_ELF_BYTES) == DisallowedContent(DisallowedKind.EXECUTABLE)
|
||||
|
||||
|
||||
def test_server_generated_filenames_are_unique_and_ignore_client_name():
|
||||
first = generate_safe_filename(DetectedFormat.PDF)
|
||||
second = generate_safe_filename(DetectedFormat.PDF)
|
||||
assert first != second
|
||||
assert first.endswith(".pdf")
|
||||
assert "/" not in first and "\\" not in first
|
||||
|
||||
|
||||
def test_malware_hook_blocks_infected_clean_format():
|
||||
result = validate_upload(content=_TEXT_BYTES, scanner=_INFECTED_SCANNER)
|
||||
assert isinstance(result, RejectedUpload)
|
||||
assert result.reason is RejectionReason.MALWARE_DETECTED
|
||||
assert "Test.Sig" in result.message
|
||||
|
||||
|
||||
def test_malware_scan_error_fails_closed():
|
||||
result = validate_upload(content=_TEXT_BYTES, scanner=_ERROR_SCANNER)
|
||||
assert isinstance(result, RejectedUpload)
|
||||
assert result.reason is RejectionReason.MALWARE_SCAN_ERROR
|
||||
|
||||
|
||||
def test_injected_clean_scanner_allows_valid_file():
|
||||
result = validate_upload(content=_TEXT_BYTES, scanner=_CLEAN_SCANNER)
|
||||
assert isinstance(result, SecuredUpload)
|
||||
|
||||
|
||||
def test_eicar_default_scanner_flags_only_eicar():
|
||||
scanner = EicarTestMalwareScanner()
|
||||
assert scanner.scan(EICAR_TEST_SIGNATURE).verdict is ScanVerdict.INFECTED
|
||||
assert scanner.scan(b"totally benign text").verdict is ScanVerdict.CLEAN
|
||||
|
||||
|
||||
def test_eicar_upload_passes_format_but_blocked_by_scanner():
|
||||
"""EICAR is valid ASCII text, so only the malware hook can stop it."""
|
||||
format_only = validate_upload(content=EICAR_TEST_SIGNATURE, scanner=_CLEAN_SCANNER)
|
||||
assert isinstance(format_only, SecuredUpload)
|
||||
|
||||
scanned = validate_upload(content=EICAR_TEST_SIGNATURE, scanner=EicarTestMalwareScanner())
|
||||
assert isinstance(scanned, RejectedUpload)
|
||||
assert scanned.reason is RejectionReason.MALWARE_DETECTED
|
||||
|
||||
|
||||
def test_safe_download_headers_force_attachment_and_nosniff():
|
||||
headers = safe_download_headers("file_abc123")
|
||||
assert headers["Content-Disposition"] == 'attachment; filename="file_abc123"'
|
||||
assert headers["X-Content-Type-Options"] == "nosniff"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("hostile", ['a"; drop', "a\r\nSet-Cookie: x=1", "../../etc/passwd", ""])
|
||||
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
|
||||
Loading…
Add table
Reference in a new issue