From 663869868fedd934452d816dc5d1db7a8b5da24c Mon Sep 17 00:00:00 2001 From: Alex Date: Wed, 12 Aug 2026 16:55:33 +0100 Subject: [PATCH] fix: better zip protection --- application/api/user/attachments/routes.py | 33 +- application/api/user/sources/upload.py | 201 +++++++--- application/api/user/tools/routes.py | 37 +- application/app.py | 51 +++ application/core/settings.py | 10 + application/security/zip_archive.py | 403 ++++++++++++++++++++ application/upload_limits.py | 146 +++++++ application/worker.py | 133 +++---- docs/content/Deploying/DocsGPT-Settings.mdx | 24 ++ tests/api/user/attachments/test_routes.py | 23 ++ tests/api/user/sources/test_upload.py | 192 ++++++++++ tests/api/user/test_tools_routes.py | 42 +- tests/test_app_routes.py | 52 +++ tests/test_upload_limits.py | 39 ++ tests/test_worker_utils.py | 23 ++ tests/test_zip_extraction_security.py | 138 ++++++- 16 files changed, 1401 insertions(+), 146 deletions(-) create mode 100644 application/security/zip_archive.py create mode 100644 application/upload_limits.py create mode 100644 tests/test_upload_limits.py diff --git a/application/api/user/attachments/routes.py b/application/api/user/attachments/routes.py index d384cb9f..b7d51a9c 100644 --- a/application/api/user/attachments/routes.py +++ b/application/api/user/attachments/routes.py @@ -8,6 +8,7 @@ import uuid from flask import Response, current_app, jsonify, make_response, redirect, request from flask_restx import fields, Namespace, Resource +from werkzeug.datastructures import FileStorage from application.api import api from application.cache import get_redis_instance @@ -36,6 +37,11 @@ from application.stt.live_session import ( ) from application.stt.stt_creator import STTCreator from application.tts.tts_creator import TTSCreator +from application.upload_limits import ( + copy_upload_to_path, + upload_limit_message, + UploadTooLargeError, +) from application.utils import ( get_agent_image_content_type, is_external_image_url, @@ -122,6 +128,8 @@ def _enforce_uploaded_audio_size_limit(file, filename: str) -> None: def _get_store_attachment_user_error(exc: Exception) -> str: if isinstance(exc, AudioFileTooLargeError): return build_stt_file_size_limit_message() + if isinstance(exc, UploadTooLargeError): + return upload_limit_message() return "Failed to process file" @@ -200,7 +208,21 @@ class StoreAttachment(Resource): f"{str(attachment_id)}/{original_filename}" ) - metadata = storage.save_file(file, relative_path) + # Validate a bounded local spool before storage reads it. + # This protects local and S3 backends even without a + # trustworthy Content-Length header. + with tempfile.TemporaryDirectory() as temp_dir: + staged_path = os.path.join(temp_dir, original_filename) + copy_upload_to_path(file, staged_path) + with open(staged_path, "rb") as staged_stream: + staged_upload = FileStorage( + stream=staged_stream, + filename=original_filename, + content_type=file.mimetype, + ) + metadata = storage.save_file( + staged_upload, relative_path + ) file_info = { "filename": original_filename, "attachment_id": str(attachment_id), @@ -224,15 +246,18 @@ class StoreAttachment(Resource): }) if not tasks: + size_limit_messages = { + build_stt_file_size_limit_message(), + upload_limit_message(), + } if errors and all( - error.get("error") == build_stt_file_size_limit_message() - for error in errors + error.get("error") in size_limit_messages for error in errors ): return make_response( jsonify( { "success": False, - "message": build_stt_file_size_limit_message(), + "message": errors[0]["error"], "errors": errors, } ), diff --git a/application/api/user/sources/upload.py b/application/api/user/sources/upload.py index 1cd41fc5..80901880 100644 --- a/application/api/user/sources/upload.py +++ b/application/api/user/sources/upload.py @@ -17,6 +17,13 @@ from application.core.settings import settings from application.storage.db.source_ids import derive_source_id as _derive_source_id from application.parser.connectors.connector_creator import ConnectorCreator from application.parser.file.constants import SUPPORTED_SOURCE_EXTENSIONS +from application.security.zip_archive import ( + extract_zip_safely, + safe_zip_error_message, + ZipExtractionBudget, + ZipExtractionError, + ZipExtractionLimits, +) from application.storage.db.repositories.idempotency import IdempotencyRepository from application.storage.db.repositories.sources import SourcesRepository from application.storage.db.source_config import SourceConfig @@ -28,6 +35,11 @@ from application.stt.upload_limits import ( enforce_audio_file_size_limit, is_audio_filename, ) +from application.upload_limits import ( + copy_upload_to_path, + upload_limit_message, + UploadTooLargeError, +) from application.utils import check_required_fields, safe_filename @@ -145,15 +157,23 @@ def _release_claim(key): "Failed to release task_dedup claim for key=%s", key, ) - - - def _enforce_audio_path_size_limit(file_path: str, filename: str) -> None: if not is_audio_filename(filename): return enforce_audio_file_size_limit(os.path.getsize(file_path)) +def _source_archive_limits() -> ZipExtractionLimits: + """Return configured limits shared by all archives in one upload.""" + return ZipExtractionLimits( + max_uncompressed_bytes=settings.UPLOAD_MAX_ARCHIVE_BYTES, + max_files=settings.UPLOAD_MAX_ARCHIVE_FILES, + max_compression_ratio=settings.UPLOAD_MAX_ARCHIVE_RATIO, + max_member_bytes=settings.UPLOAD_MAX_FILE_BYTES, + max_depth=settings.UPLOAD_MAX_ARCHIVE_DEPTH, + ) + + @sources_upload_ns.route("/upload") class UploadFile(Resource): @api.expect( @@ -222,67 +242,66 @@ class UploadFile(Resource): base_path = f"{settings.UPLOAD_FOLDER}/{safe_user}/{dir_name}" file_name_map = {} + active_upload_name = "uploaded archive" try: storage = StorageCreator.get_storage() - for file in files: - original_filename = os.path.basename(file.filename) - safe_file = safe_filename(original_filename) - if original_filename: - file_name_map[safe_file] = original_filename + # Stage and validate every file before the first storage write. A + # later oversized file/archive cannot leave a partially accepted + # source behind. All ZIPs in the request share one expansion budget. + with tempfile.TemporaryDirectory() as temp_dir: + pending_files: list[tuple[str, str]] = [] + archive_budget = ZipExtractionBudget() + archive_limits = _source_archive_limits() - with tempfile.TemporaryDirectory() as temp_dir: - temp_file_path = os.path.join(temp_dir, safe_file) - file.save(temp_file_path) + for index, file in enumerate(files): + original_filename = os.path.basename(file.filename) + safe_file = safe_filename(original_filename) + active_upload_name = safe_zip_error_message( + original_filename or safe_file, max_chars=200 + ) + if original_filename: + file_name_map[safe_file] = original_filename + + upload_dir = os.path.join(temp_dir, str(index)) + os.makedirs(upload_dir, exist_ok=True) + temp_file_path = os.path.join(upload_dir, safe_file) + copy_upload_to_path(file, temp_file_path) _enforce_audio_path_size_limit(temp_file_path, safe_file) - # Only extract actual .zip files, not Office formats (.docx, .xlsx, .pptx) - # which are technically zip archives but should be processed as-is + # Office/e-book containers are ZIP-based formats but must + # be parsed as documents, not expanded as user archives. is_office_format = safe_file.lower().endswith( (".docx", ".xlsx", ".pptx", ".odt", ".ods", ".odp", ".epub") ) if zipfile.is_zipfile(temp_file_path) and not is_office_format: - try: - with zipfile.ZipFile(temp_file_path, "r") as zip_ref: - zip_ref.extractall(path=temp_dir) - - # Walk through extracted files and upload them - - for root, _, files in os.walk(temp_dir): - for extracted_file in files: - if ( - os.path.join(root, extracted_file) - == temp_file_path - ): - continue - rel_path = os.path.relpath( - os.path.join(root, extracted_file), temp_dir - ) - storage_path = f"{base_path}/{rel_path}" - _enforce_audio_path_size_limit( - os.path.join(root, extracted_file), - extracted_file, - ) - - with open( - os.path.join(root, extracted_file), "rb" - ) as f: - storage.save_file(f, storage_path) - except Exception as e: - current_app.logger.error( - f"Error extracting zip: {e}", exc_info=True - ) - # If zip extraction fails, save the original zip file - - file_path = f"{base_path}/{safe_file}" - with open(temp_file_path, "rb") as f: - storage.save_file(f, file_path) + extract_dir = os.path.join(upload_dir, "extracted") + extract_zip_safely( + temp_file_path, + extract_dir, + archive_limits, + archive_budget, + ) + for root, _, extracted_files in os.walk(extract_dir): + for extracted_file in extracted_files: + local_path = os.path.join(root, extracted_file) + _enforce_audio_path_size_limit( + local_path, extracted_file + ) + rel_path = os.path.relpath( + local_path, extract_dir + ).replace(os.sep, "/") + pending_files.append( + (local_path, f"{base_path}/{rel_path}") + ) else: - # For non-zip files, save directly + pending_files.append( + (temp_file_path, f"{base_path}/{safe_file}") + ) - file_path = f"{base_path}/{safe_file}" - with open(temp_file_path, "rb") as f: - storage.save_file(f, file_path) + for local_path, storage_path in pending_files: + with open(local_path, "rb") as staged_file: + storage.save_file(staged_file, storage_path) # Mint the source UUID up here so the HTTP response and the # worker's SSE envelopes share one id. With an idempotency # key we reuse the deterministic uuid5 (retried task lands on @@ -324,6 +343,33 @@ class UploadFile(Resource): ), 413, ) + except (UploadTooLargeError, ZipExtractionError) as err: + zip_error_detail = safe_zip_error_message(err) + if isinstance(err, ZipExtractionError): + current_app.logger.warning( + "Rejected unsafe ZIP upload %s: %s", + active_upload_name, + zip_error_detail, + ) + else: + current_app.logger.warning( + "Rejected oversized upload %s", active_upload_name + ) + if scoped_key: + _release_claim(scoped_key) + return make_response( + jsonify( + { + "success": False, + "message": ( + f"{active_upload_name}: {zip_error_detail}" + if isinstance(err, ZipExtractionError) + else upload_limit_message() + ), + } + ), + 413, + ) except Exception as err: current_app.logger.error(f"Error uploading file: {err}", exc_info=True) if scoped_key: @@ -678,21 +724,43 @@ class ManageSourceFiles(Resource): cached["source_id"] = resolved_source_id return make_response(jsonify(cached), 200) - added_files = [] - map_updated = False - target_dir = source_file_path if parent_dir: target_dir = f"{source_file_path}/{parent_dir}" - for file in files: - if file.filename: + added_files = [] + map_updated = False + with tempfile.TemporaryDirectory() as temp_dir: + staged_files: list[tuple[str, str, str, str]] = [] + for index, file in enumerate(files): + if not file.filename: + continue original_filename = os.path.basename(file.filename) safe_filename_str = safe_filename(original_filename) - file_path = f"{target_dir}/{safe_filename_str}" + staged_path = os.path.join( + temp_dir, f"{index}-{safe_filename_str}" + ) + copy_upload_to_path(file, staged_path) + _enforce_audio_path_size_limit( + staged_path, safe_filename_str + ) + staged_files.append( + ( + staged_path, + f"{target_dir}/{safe_filename_str}", + safe_filename_str, + original_filename, + ) + ) - # Save file to storage - - storage.save_file(file, file_path) + # All files are bounded before the first storage mutation. + for ( + staged_path, + file_path, + safe_filename_str, + original_filename, + ) in staged_files: + with open(staged_path, "rb") as staged_file: + storage.save_file(staged_file, file_path) added_files.append(safe_filename_str) if original_filename: relative_key = ( @@ -977,6 +1045,17 @@ class ManageSourceFiles(Resource): ), 200, ) + except (UploadTooLargeError, AudioFileTooLargeError) as err: + if scoped_key and not claim_transferred: + _release_claim(scoped_key) + message = ( + build_stt_file_size_limit_message() + if isinstance(err, AudioFileTooLargeError) + else upload_limit_message() + ) + return make_response( + jsonify({"success": False, "message": message}), 413 + ) except Exception as err: # Release the dedup claim only if it wasn't transferred to # a worker. Without this, a same-key retry within the 24h diff --git a/application/api/user/tools/routes.py b/application/api/user/tools/routes.py index 4c1f3f92..6e4b5a99 100644 --- a/application/api/user/tools/routes.py +++ b/application/api/user/tools/routes.py @@ -18,6 +18,7 @@ from application.agents.tools.tool_manager import ToolManager from application.api import api from application.api.user.artifacts.authz import Principal, authorize_artifact from application.api.user.team_sharing import effective_write_owner, visible_with_access +from application.core.settings import settings from application.core.url_validation import SSRFError, validate_url from application.security.encryption import decrypt_credentials, encrypt_credentials from application.storage.db.base_repository import looks_like_uuid @@ -27,6 +28,11 @@ from application.storage.db.repositories.todos import TodosRepository from application.storage.db.repositories.user_tools import UserToolsRepository from application.storage.db.repositories.users import UsersRepository from application.storage.db.session import db_readonly, db_session +from application.upload_limits import ( + read_text_upload_limited, + upload_limit_message, + UploadTooLargeError, +) from application.utils import check_required_fields, validate_function_name tool_config = {} @@ -821,7 +827,21 @@ class ParseSpec(Resource): jsonify({"success": False, "message": "No file selected"}), 400 ) try: - spec_content = file.read().decode("utf-8") + spec_content = read_text_upload_limited( + file, max_bytes=settings.PARSE_SPEC_MAX_BYTES + ) + except UploadTooLargeError: + return make_response( + jsonify( + { + "success": False, + "message": upload_limit_message( + settings.PARSE_SPEC_MAX_BYTES + ), + } + ), + 413, + ) except UnicodeDecodeError: return make_response( jsonify({"success": False, "message": "Invalid file encoding"}), 400 @@ -833,6 +853,21 @@ class ParseSpec(Resource): return make_response( jsonify({"success": False, "message": "No spec provided"}), 400 ) + if ( + not isinstance(spec_content, str) + or len(spec_content.encode("utf-8")) > settings.PARSE_SPEC_MAX_BYTES + ): + return make_response( + jsonify( + { + "success": False, + "message": upload_limit_message( + settings.PARSE_SPEC_MAX_BYTES + ), + } + ), + 413, + ) if not spec_content or not spec_content.strip(): return make_response( jsonify({"success": False, "message": "Empty spec content"}), 400 diff --git a/application/app.py b/application/app.py index b717ba9c..ecfa03cf 100644 --- a/application/app.py +++ b/application/app.py @@ -6,6 +6,7 @@ import uuid import dotenv from flask import Flask, Response, jsonify, redirect, request from jose import jwt +from werkzeug.exceptions import RequestEntityTooLarge from application.auth import handle_auth @@ -35,6 +36,10 @@ from application.stt.upload_limits import ( # noqa: E402 build_stt_file_size_limit_message, should_reject_stt_request, ) +from application.upload_limits import ( # noqa: E402 + is_document_upload_path, + upload_request_limit_message, +) if platform.system() == "Windows": @@ -85,6 +90,52 @@ app.config.update( celery.config_from_object("application.celeryconfig") api.init_app(app) + +def _upload_limit_error_payload() -> dict[str, bool | str]: + """Build a consistent 413 payload using the active route-specific limit.""" + active_limit = request.max_content_length + if active_limit is None: + active_limit = settings.UPLOAD_MAX_REQUEST_BYTES + return { + "success": False, + "message": upload_request_limit_message(active_limit), + } + + +@app.errorhandler(RequestEntityTooLarge) +def handle_request_entity_too_large(_error): + """Return API-shaped JSON instead of Werkzeug's HTML 413 page.""" + return jsonify(_upload_limit_error_payload()), 413 + + +@api.errorhandler(RequestEntityTooLarge) +def handle_restx_request_entity_too_large(_error): + """Keep Flask-RESTX from replacing the configured upload-limit response.""" + return _upload_limit_error_payload(), 413 + + +@app.before_request +def enforce_document_upload_request_size_limit(): + """Bound multipart parsing for public file-upload routes only. + + Internal worker index uploads are intentionally excluded; they are trusted + service traffic and can legitimately exceed the end-user document limit. + """ + if request.method == "OPTIONS" or not is_document_upload_path(request.path): + return None + request_limit = ( + settings.PARSE_SPEC_MAX_BYTES + if request.path == "/api/parse_spec" and request.is_json + else settings.UPLOAD_MAX_REQUEST_BYTES + ) + request.max_content_length = request_limit + if ( + request.content_length is not None + and request.content_length > request_limit + ): + raise RequestEntityTooLarge() + return None + # The same stable secret also signs opaque agent-image capabilities, including # in no-auth mode. Production replicas must receive one shared configured key; # only local development may use the atomic filesystem fallback. diff --git a/application/core/settings.py b/application/core/settings.py index b170fb22..2892da80 100644 --- a/application/core/settings.py +++ b/application/core/settings.py @@ -88,6 +88,16 @@ class Settings(BaseSettings): "request_limit": 500, } UPLOAD_FOLDER: str = "inputs" + # Public upload request cap is applied by Flask before multipart parsing. + # The per-file cap is also enforced while copying each controlled stream. + UPLOAD_MAX_REQUEST_BYTES: int = Field(default=256 * 1024 * 1024, gt=0) + UPLOAD_MAX_FILE_BYTES: int = Field(default=100 * 1024 * 1024, gt=0) + PARSE_SPEC_MAX_BYTES: int = Field(default=10 * 1024 * 1024, gt=0) + # ZIP limits apply cumulatively across nested archives in one extraction. + UPLOAD_MAX_ARCHIVE_BYTES: int = Field(default=250 * 1024 * 1024, gt=0) + UPLOAD_MAX_ARCHIVE_FILES: int = Field(default=10_000, gt=0) + UPLOAD_MAX_ARCHIVE_RATIO: int = Field(default=1000, gt=0) + UPLOAD_MAX_ARCHIVE_DEPTH: int = Field(default=3, ge=0) PARSE_PDF_AS_IMAGE: bool = False PARSE_IMAGE_REMOTE: bool = False DOCLING_OCR_ENABLED: bool = False # Enable OCR for docling parsers (PDF, images) diff --git a/application/security/zip_archive.py b/application/security/zip_archive.py new file mode 100644 index 00000000..04ceb8a3 --- /dev/null +++ b/application/security/zip_archive.py @@ -0,0 +1,403 @@ +"""Safe, bounded extraction for untrusted ZIP archives.""" + +from __future__ import annotations + +import os +import shutil +import stat +import tempfile +import zipfile +from dataclasses import dataclass +from pathlib import PurePosixPath + + +_COPY_CHUNK_BYTES = 64 * 1024 +_MAX_MEMBER_NAME_BYTES = 4096 +_MAX_MEMBER_DEPTH = 64 + + +class ZipExtractionError(ValueError): + """Raised when a ZIP is invalid or exceeds an extraction safety limit.""" + + +@dataclass(frozen=True) +class ZipExtractionLimits: + """Resource ceilings for one extraction tree.""" + + max_uncompressed_bytes: int + max_files: int + max_compression_ratio: int + max_member_bytes: int | None = None + max_depth: int = 3 + + +@dataclass +class ZipExtractionBudget: + """Cumulative resources reserved across nested archives.""" + + uncompressed_bytes: int = 0 + files: int = 0 + + +@dataclass(frozen=True) +class _ValidatedMember: + info: zipfile.ZipInfo + relative_path: str + + +def _collision_safe_member_path( + relative_path: str, + assignments: dict[tuple[tuple[str, ...], str], str], + occupied_names: dict[tuple[str, ...], set[str]], +) -> str: + """Map case-colliding member components to stable, portable names.""" + original_parent: tuple[str, ...] = () + assigned_parent: tuple[str, ...] = () + assigned_parts: list[str] = [] + for component in relative_path.split("/"): + assignment_key = (original_parent, component) + assigned_component = assignments.get(assignment_key) + if assigned_component is None: + used = occupied_names.setdefault(assigned_parent, set()) + assigned_component = component + suffix = 2 + stem, extension = os.path.splitext(component) + while assigned_component.casefold() in used: + assigned_component = f"{stem} ({suffix}){extension}" + suffix += 1 + assignments[assignment_key] = assigned_component + used.add(assigned_component.casefold()) + assigned_parts.append(assigned_component) + original_parent += (component,) + assigned_parent += (assigned_component,) + return "/".join(assigned_parts) + + +def _safe_member_path(name: str) -> str: + """Normalize a ZIP member name and reject paths escaping extraction.""" + if not name or "\x00" in name: + raise ZipExtractionError("ZIP contains an invalid empty member name") + if len(name.encode("utf-8", errors="surrogatepass")) > _MAX_MEMBER_NAME_BYTES: + raise ZipExtractionError("ZIP member name exceeds the safety limit") + + normalized = name.replace("\\", "/") + path = PurePosixPath(normalized) + parts = path.parts + if ( + path.is_absolute() + or normalized.startswith("//") + or any(part in {"", ".", ".."} for part in parts) + or (parts and parts[0].endswith(":")) + ): + raise ZipExtractionError(f"ZIP contains a path traversal attempt: {name}") + if len(parts) > _MAX_MEMBER_DEPTH: + raise ZipExtractionError("ZIP member path is nested too deeply") + return "/".join(parts) + + +def _reject_special_member(info: zipfile.ZipInfo) -> None: + """Reject symlinks and device-like Unix entries.""" + mode = info.external_attr >> 16 + file_type = stat.S_IFMT(mode) + if file_type not in {0, stat.S_IFREG, stat.S_IFDIR}: + raise ZipExtractionError(f"ZIP contains a special file: {info.filename}") + if info.flag_bits & 0x1: + raise ZipExtractionError("Encrypted ZIP members are not supported") + + +def _validate_members( + archive: zipfile.ZipFile, + limits: ZipExtractionLimits, + budget: ZipExtractionBudget, +) -> tuple[list[_ValidatedMember], int, int]: + members: list[_ValidatedMember] = [] + seen_paths: dict[str, bool] = {} + component_assignments: dict[tuple[tuple[str, ...], str], str] = {} + occupied_names: dict[tuple[str, ...], set[str]] = {} + added_bytes = 0 + added_files = 0 + added_compressed_bytes = 0 + + for info in archive.infolist(): + relative_path = _safe_member_path(info.filename) + is_directory = info.is_dir() + prior_is_directory = seen_paths.get(relative_path) + if prior_is_directory is not None and not ( + prior_is_directory and is_directory + ): + raise ZipExtractionError(f"ZIP contains a duplicate path: {info.filename}") + seen_paths[relative_path] = is_directory + _reject_special_member(info) + extraction_path = _collision_safe_member_path( + relative_path, component_assignments, occupied_names + ) + + if is_directory: + members.append(_ValidatedMember(info=info, relative_path=extraction_path)) + continue + + added_files += 1 + added_bytes += int(info.file_size) + added_compressed_bytes += int(info.compress_size) + if limits.max_member_bytes and info.file_size > limits.max_member_bytes: + raise ZipExtractionError( + "ZIP member exceeds the " + f"{limits.max_member_bytes}-byte per-file limit" + ) + if budget.files + added_files > limits.max_files: + raise ZipExtractionError( + f"ZIP contains too many files (>{limits.max_files})" + ) + if budget.uncompressed_bytes + added_bytes > limits.max_uncompressed_bytes: + raise ZipExtractionError( + "ZIP extraction exceeds the " + f"{limits.max_uncompressed_bytes}-byte expansion limit" + ) + + members.append(_ValidatedMember(info=info, relative_path=extraction_path)) + + if added_bytes: + if added_compressed_bytes <= 0: + raise ZipExtractionError("ZIP has an invalid compression ratio") + ratio = added_bytes / added_compressed_bytes + if ratio > limits.max_compression_ratio: + raise ZipExtractionError( + "ZIP exceeds the " + f"{limits.max_compression_ratio}:1 compression ratio limit" + ) + + return members, added_bytes, added_files + + +def validate_zip_archive( + zip_path: str | os.PathLike[str], + limits: ZipExtractionLimits, + budget: ZipExtractionBudget | None = None, +) -> None: + """Validate central-directory metadata without extracting any content.""" + active_budget = budget or ZipExtractionBudget() + try: + with zipfile.ZipFile(zip_path, "r") as archive: + _validate_members(archive, limits, active_budget) + except (OSError, zipfile.BadZipFile) as exc: + raise ZipExtractionError(f"Invalid or corrupted ZIP file: {exc}") from exc + + +def _extract_archive_into( + zip_path: str | os.PathLike[str], + destination_root: str, + limits: ZipExtractionLimits, + budget: ZipExtractionBudget, +) -> list[str]: + """Extract one validated archive into a private staging directory.""" + extracted_files: list[str] = [] + with zipfile.ZipFile(zip_path, "r") as archive: + members, added_bytes, added_files = _validate_members( + archive, limits, budget + ) + + # Reserve declared resources before decompression so nested callers + # cannot reset the allowance between layers. + budget.uncompressed_bytes += added_bytes + budget.files += added_files + + for member in members: + target = os.path.realpath( + os.path.join(destination_root, *member.relative_path.split("/")) + ) + try: + inside_destination = ( + os.path.commonpath((destination_root, target)) + == destination_root + ) + except ValueError as exc: + raise ZipExtractionError( + f"ZIP contains a path traversal attempt: {member.info.filename}" + ) from exc + if not inside_destination: + raise ZipExtractionError( + "ZIP contains a path traversal attempt: " + f"{member.info.filename}" + ) + if member.info.is_dir(): + os.makedirs(target, exist_ok=True) + continue + if os.path.lexists(target): + raise ZipExtractionError( + f"ZIP member would overwrite an existing path: {member.info.filename}" + ) + + os.makedirs(os.path.dirname(target), exist_ok=True) + written = 0 + with archive.open(member.info, "r") as source, open(target, "xb") as output: + while True: + chunk = source.read(_COPY_CHUNK_BYTES) + if not chunk: + break + written += len(chunk) + if written > member.info.file_size: + raise ZipExtractionError( + "ZIP member expanded past its declared size: " + f"{member.info.filename}" + ) + output.write(chunk) + if written != member.info.file_size: + raise ZipExtractionError( + f"ZIP member size did not match metadata: {member.info.filename}" + ) + extracted_files.append(target) + return extracted_files + + +def _extract_tree_into( + zip_path: str | os.PathLike[str], + destination_root: str, + limits: ZipExtractionLimits, + budget: ZipExtractionBudget, + depth: int, +) -> None: + """Extract an archive and recursively expand nested ``.zip`` members.""" + extracted_files = _extract_archive_into( + zip_path, destination_root, limits, budget + ) + for extracted_path in extracted_files: + if not extracted_path.lower().endswith(".zip"): + continue + if not zipfile.is_zipfile(extracted_path): + continue + if depth >= limits.max_depth: + raise ZipExtractionError( + f"ZIP nesting exceeds the {limits.max_depth}-level limit" + ) + parent = os.path.dirname(extracted_path) + stem = os.path.splitext(os.path.basename(extracted_path))[0] or "archive" + occupied_names = {name.casefold() for name in os.listdir(parent)} + nested_name = stem + suffix = 2 + while nested_name.casefold() in occupied_names: + nested_name = f"{stem} ({suffix})" + suffix += 1 + nested_root = os.path.join(parent, nested_name) + os.makedirs(nested_root) + _extract_tree_into( + extracted_path, + nested_root, + limits, + budget, + depth + 1, + ) + os.remove(extracted_path) + + +def _commit_staging(staging_root: str, destination_root: str) -> None: + """Commit a complete staging tree, rolling back an interrupted merge.""" + if not os.path.lexists(destination_root): + os.replace(staging_root, destination_root) + return + if not os.path.isdir(destination_root): + raise ZipExtractionError("ZIP extraction destination is not a directory") + + staged_directories: list[tuple[str, str]] = [] + staged_files: list[tuple[str, str]] = [] + for root, directories, files in os.walk(staging_root): + relative_root = os.path.relpath(root, staging_root) + destination_directory = ( + destination_root + if relative_root == "." + else os.path.join(destination_root, relative_root) + ) + for directory in directories: + source = os.path.join(root, directory) + target = os.path.join(destination_directory, directory) + if os.path.lexists(target) and not os.path.isdir(target): + raise ZipExtractionError( + "ZIP member would overwrite an existing path" + ) + staged_directories.append((source, target)) + for filename in files: + source = os.path.join(root, filename) + target = os.path.join(destination_directory, filename) + if os.path.lexists(target): + raise ZipExtractionError( + "ZIP member would overwrite an existing path" + ) + staged_files.append((source, target)) + + created_directories: list[str] = [] + moved_files: list[str] = [] + try: + for _, target in staged_directories: + if not os.path.isdir(target): + os.makedirs(target) + created_directories.append(target) + for source, target in staged_files: + os.makedirs(os.path.dirname(target), exist_ok=True) + os.replace(source, target) + moved_files.append(target) + except Exception: + for target in reversed(moved_files): + try: + os.remove(target) + except OSError: + pass + for target in reversed(created_directories): + try: + os.rmdir(target) + except OSError: + pass + raise + + +def extract_zip_safely( + zip_path: str | os.PathLike[str], + extract_to: str | os.PathLike[str], + limits: ZipExtractionLimits, + budget: ZipExtractionBudget | None = None, +) -> ZipExtractionBudget: + """Transactionally extract a ZIP tree under cumulative resource limits. + + Nested ``.zip`` members are expanded inside a private staging directory. + The destination is changed only after every layer has passed validation + and extraction, so a corrupt member cannot leave partial content behind. + """ + active_budget = budget or ZipExtractionBudget() + destination_root = os.path.realpath(extract_to) + destination_parent = os.path.dirname(destination_root) + staging_root: str | None = None + + try: + os.makedirs(destination_parent, exist_ok=True) + staging_root = tempfile.mkdtemp( + prefix=".zip-extract-", dir=destination_parent + ) + _extract_tree_into(zip_path, staging_root, limits, active_budget, depth=0) + _commit_staging(staging_root, destination_root) + except ZipExtractionError: + raise + except ( + OSError, + RuntimeError, + ValueError, + zipfile.BadZipFile, + zipfile.LargeZipFile, + ) as exc: + raise ZipExtractionError(f"Unable to extract ZIP safely: {exc}") from exc + finally: + if staging_root is not None: + shutil.rmtree(staging_root, ignore_errors=True) + + return active_budget + + +def safe_zip_error_message(error: BaseException | str, max_chars: int = 512) -> str: + """Return a bounded, single-line representation safe for logs and JSON.""" + raw = str(error) + escaped = "".join( + f"\\x{ord(character):02x}" + if ord(character) < 32 or 127 <= ord(character) <= 159 + else character + for character in raw + ) + if len(escaped) > max_chars: + return f"{escaped[: max_chars - 3]}..." + return escaped diff --git a/application/upload_limits.py b/application/upload_limits.py new file mode 100644 index 00000000..dc8ff379 --- /dev/null +++ b/application/upload_limits.py @@ -0,0 +1,146 @@ +"""Bound user-controlled upload streams before storage or parsing.""" + +from __future__ import annotations + +import io +import os +from typing import BinaryIO + +from application.core.settings import settings + + +_COPY_CHUNK_BYTES = 64 * 1024 +_DOCUMENT_UPLOAD_PATHS = frozenset( + { + "/api/upload", + "/api/manage_source_files", + "/api/store_attachment", + "/api/parse_spec", + "/api/create_agent", + } +) + + +class UploadTooLargeError(ValueError): + """Raised when one uploaded file exceeds the configured byte cap.""" + + +class _LimitedRawReader(io.RawIOBase): + """Expose a binary stream while enforcing an exact cumulative read cap.""" + + def __init__(self, stream: BinaryIO, max_bytes: int) -> None: + super().__init__() + self._stream = stream + self._max_bytes = max_bytes + self._bytes_read = 0 + + def readable(self) -> bool: + """Return whether this wrapper supports reads.""" + return True + + def readinto(self, buffer: bytearray) -> int: + """Read one bounded chunk into ``buffer``.""" + remaining_with_probe = self._max_bytes - self._bytes_read + 1 + chunk = self._stream.read(min(len(buffer), remaining_with_probe)) + if not chunk: + return 0 + self._bytes_read += len(chunk) + if self._bytes_read > self._max_bytes: + raise UploadTooLargeError(upload_limit_message(self._max_bytes)) + buffer[: len(chunk)] = chunk + return len(chunk) + + +def is_document_upload_path(path: str) -> bool: + """Return whether a route accepts a user-controlled document/image file.""" + return path in _DOCUMENT_UPLOAD_PATHS or path.startswith("/api/update_agent/") + + +def upload_limit_message(max_bytes: int | None = None) -> str: + """Return the stable client-facing file-size rejection message.""" + limit = int(max_bytes or settings.UPLOAD_MAX_FILE_BYTES) + return f"File exceeds the {limit}-byte upload limit" + + +def upload_request_limit_message(max_bytes: int | None = None) -> str: + """Return the stable client-facing request-size rejection message.""" + limit = int(max_bytes or settings.UPLOAD_MAX_REQUEST_BYTES) + return f"Request exceeds the {limit}-byte upload limit" + + +def copy_upload_to_path( + upload: BinaryIO, + destination: str | os.PathLike[str], + max_bytes: int | None = None, +) -> int: + """Copy an upload to disk while enforcing a hard streaming byte limit. + + The limit is checked during reads, so it does not rely on a trustworthy + ``Content-Length`` header or on the input stream being seekable. + """ + limit = int(max_bytes or settings.UPLOAD_MAX_FILE_BYTES) + stream = getattr(upload, "stream", upload) + try: + stream.seek(0) + except (AttributeError, OSError): + pass + + total = 0 + with open(destination, "wb") as target: + while True: + chunk = stream.read(min(_COPY_CHUNK_BYTES, limit - total + 1)) + if not chunk: + break + total += len(chunk) + if total > limit: + raise UploadTooLargeError(upload_limit_message(limit)) + target.write(chunk) + return total + + +def read_upload_limited(upload: BinaryIO, max_bytes: int | None = None) -> bytes: + """Read at most one configured file into memory, rejecting overflow.""" + limit = int(max_bytes or settings.UPLOAD_MAX_FILE_BYTES) + stream = getattr(upload, "stream", upload) + try: + stream.seek(0) + except (AttributeError, OSError): + pass + data = bytearray() + while True: + chunk = stream.read(min(_COPY_CHUNK_BYTES, limit - len(data) + 1)) + if not chunk: + break + data.extend(chunk) + if len(data) > limit: + raise UploadTooLargeError(upload_limit_message(limit)) + return bytes(data) + + +def read_text_upload_limited( + upload: BinaryIO, + max_bytes: int | None = None, + encoding: str = "utf-8", +) -> str: + """Decode a bounded upload without first duplicating it into ``bytes``.""" + limit = int(max_bytes or settings.UPLOAD_MAX_FILE_BYTES) + stream = getattr(upload, "stream", upload) + try: + stream.seek(0) + except (AttributeError, OSError): + pass + + raw_reader = _LimitedRawReader(stream, limit) + buffered_reader = io.BufferedReader(raw_reader, buffer_size=_COPY_CHUNK_BYTES) + text_reader = io.TextIOWrapper(buffered_reader, encoding=encoding) + try: + return text_reader.read() + finally: + try: + text_reader.detach() + except (ValueError, OSError): + pass + try: + buffered_reader.detach() + except (ValueError, OSError): + pass diff --git a/application/worker.py b/application/worker.py index 19b298c5..f22368f2 100755 --- a/application/worker.py +++ b/application/worker.py @@ -8,7 +8,6 @@ import string import tempfile import threading from typing import Any, Dict -import zipfile import uuid from collections import Counter @@ -31,6 +30,14 @@ from application.parser.remote.remote_creator import ( normalize_remote_data, ) from application.parser.schema.base import Document +from application.security.zip_archive import ( + extract_zip_safely, + safe_zip_error_message, + validate_zip_archive, + ZipExtractionBudget, + ZipExtractionError, + ZipExtractionLimits, +) from application.storage.db.base_repository import looks_like_uuid from application.storage.db.repositories.agents import AgentsRepository @@ -293,15 +300,11 @@ current_dir = os.path.dirname( os.path.dirname(os.path.dirname(os.path.abspath(__file__))) ) -# Zip extraction security limits -MAX_UNCOMPRESSED_SIZE = 500 * 1024 * 1024 # 500 MB max uncompressed size -MAX_FILE_COUNT = 10000 # Maximum number of files to extract -MAX_COMPRESSION_RATIO = 100 # Maximum compression ratio (to detect zip bombs) - - -class ZipExtractionError(Exception): - """Raised when zip extraction fails due to security constraints.""" - pass +# Zip extraction security limits. Kept as module constants for backward +# compatibility with worker callers/tests; values come from operator settings. +MAX_UNCOMPRESSED_SIZE = settings.UPLOAD_MAX_ARCHIVE_BYTES +MAX_FILE_COUNT = settings.UPLOAD_MAX_ARCHIVE_FILES +MAX_COMPRESSION_RATIO = settings.UPLOAD_MAX_ARCHIVE_RATIO def _is_path_safe(base_path: str, target_path: str) -> bool: @@ -337,58 +340,26 @@ def _validate_zip_safety(zip_path: str, extract_to: str) -> None: Raises: ZipExtractionError: If the zip file fails security validation. """ - try: - with zipfile.ZipFile(zip_path, "r") as zip_ref: - # Get compressed size - compressed_size = os.path.getsize(zip_path) - - # Calculate total uncompressed size and file count - total_uncompressed = 0 - file_count = 0 - - for info in zip_ref.infolist(): - file_count += 1 - - # Check file count limit - if file_count > MAX_FILE_COUNT: - raise ZipExtractionError( - f"Zip file contains too many files (>{MAX_FILE_COUNT}). " - "This may be a zip bomb attack." - ) - - # Accumulate uncompressed size - total_uncompressed += info.file_size - - # Check total uncompressed size - if total_uncompressed > MAX_UNCOMPRESSED_SIZE: - raise ZipExtractionError( - f"Zip file uncompressed size exceeds limit " - f"({total_uncompressed / (1024*1024):.1f} MB > " - f"{MAX_UNCOMPRESSED_SIZE / (1024*1024):.1f} MB). " - "This may be a zip bomb attack." - ) - - # Check for path traversal (zip slip) - target_path = os.path.join(extract_to, info.filename) - if not _is_path_safe(extract_to, target_path): - raise ZipExtractionError( - f"Zip file contains path traversal attempt: {info.filename}" - ) - - # Check compression ratio (only if compressed size is meaningful) - if compressed_size > 0 and total_uncompressed > 0: - compression_ratio = total_uncompressed / compressed_size - if compression_ratio > MAX_COMPRESSION_RATIO: - raise ZipExtractionError( - f"Zip file has suspicious compression ratio ({compression_ratio:.1f}:1 > " - f"{MAX_COMPRESSION_RATIO}:1). This may be a zip bomb attack." - ) - - except zipfile.BadZipFile as e: - raise ZipExtractionError(f"Invalid or corrupted zip file: {e}") + del extract_to # Path checks are platform-independent inside the validator. + validate_zip_archive( + zip_path, + ZipExtractionLimits( + max_uncompressed_bytes=MAX_UNCOMPRESSED_SIZE, + max_files=MAX_FILE_COUNT, + max_compression_ratio=MAX_COMPRESSION_RATIO, + max_member_bytes=settings.UPLOAD_MAX_FILE_BYTES, + max_depth=RECURSION_DEPTH, + ), + ) -def extract_zip_recursive(zip_path, extract_to, current_depth=0, max_depth=5): +def extract_zip_recursive( + zip_path: str, + extract_to: str, + current_depth: int = 0, + max_depth: int = 5, + _budget: ZipExtractionBudget | None = None, +) -> None: """ Recursively extract zip files with security protections. @@ -410,16 +381,27 @@ def extract_zip_recursive(zip_path, extract_to, current_depth=0, max_depth=5): return try: - # Validate zip file safety before extraction - _validate_zip_safety(zip_path, extract_to) - - # Safe to extract - with zipfile.ZipFile(zip_path, "r") as zip_ref: - zip_ref.extractall(extract_to) + budget = _budget or ZipExtractionBudget() + extract_zip_safely( + zip_path, + extract_to, + ZipExtractionLimits( + max_uncompressed_bytes=MAX_UNCOMPRESSED_SIZE, + max_files=MAX_FILE_COUNT, + max_compression_ratio=MAX_COMPRESSION_RATIO, + max_member_bytes=settings.UPLOAD_MAX_FILE_BYTES, + max_depth=max_depth - current_depth, + ), + budget, + ) os.remove(zip_path) # Remove the zip file after extracting except ZipExtractionError as e: - logging.error(f"Zip security validation failed for {zip_path}: {e}") + logging.error( + "Zip security validation failed for %s: %s", + safe_zip_error_message(zip_path), + safe_zip_error_message(e), + ) # Remove the potentially malicious zip file try: os.remove(zip_path) @@ -427,17 +409,14 @@ def extract_zip_recursive(zip_path, extract_to, current_depth=0, max_depth=5): pass return except Exception as e: - logging.error(f"Error extracting zip file {zip_path}: {e}", exc_info=True) + logging.error( + "Error extracting zip file %s: %s", + safe_zip_error_message(zip_path), + safe_zip_error_message(e), + exc_info=True, + ) return - # Check for nested zip files and extract them - for root, dirs, files in os.walk(extract_to): - for file in files: - if file.endswith(".zip"): - # If a nested zip file is found, extract it recursively - file_path = os.path.join(root, file) - extract_zip_recursive(file_path, root, current_depth + 1, max_depth) - def download_file(url, params, dest_path): try: @@ -623,7 +602,7 @@ def ingest_worker( f.write(file_data.read()) # Handle zip files - if temp_filename.endswith(".zip"): + if temp_filename.lower().endswith(".zip"): logging.info(f"Extracting zip file: {temp_filename}") extract_zip_recursive( temp_file_path, diff --git a/docs/content/Deploying/DocsGPT-Settings.mdx b/docs/content/Deploying/DocsGPT-Settings.mdx index 3e25ee7c..b71c93d6 100644 --- a/docs/content/Deploying/DocsGPT-Settings.mdx +++ b/docs/content/Deploying/DocsGPT-Settings.mdx @@ -174,6 +174,30 @@ the built-in catalog — it does **not** fail to start. If a YAML declares an unknown provider name or has a schema error, the app **does** fail to start, with the offending file path in the message. +## Document Upload Limits + +DocsGPT bounds public document and attachment uploads before storage or parsing. +ZIP limits are cumulative across nested archives. Unsafe archive paths, +duplicate file entries, encrypted members, symlinks, and other special files +are rejected. Repeated directory records are ignored, and case-colliding names +are disambiguated without overwriting either file. Nested archives expand into +a directory named after the archive; files that merely end in `.zip` remain +ordinary files. + +| Setting | Default | Description | +| --- | --- | --- | +| `UPLOAD_MAX_REQUEST_BYTES` | `268435456` | Maximum request body for a public file-upload route, including multipart overhead. | +| `UPLOAD_MAX_FILE_BYTES` | `104857600` | Maximum encoded size of one uploaded file or one extracted ZIP member. | +| `PARSE_SPEC_MAX_BYTES` | `10485760` | Maximum UTF-8 size of an uploaded or JSON API specification. | +| `UPLOAD_MAX_ARCHIVE_BYTES` | `262144000` | Maximum total uncompressed bytes across all ZIP layers in one extraction. | +| `UPLOAD_MAX_ARCHIVE_FILES` | `10000` | Maximum total extracted files across all nested layers; directory records do not count. | +| `UPLOAD_MAX_ARCHIVE_RATIO` | `1000` | Maximum uncompressed-to-compressed expansion ratio for each archive layer. | +| `UPLOAD_MAX_ARCHIVE_DEPTH` | `3` | Maximum number of nested ZIP layers expanded after the outer archive. | + +These byte/count limits and the non-negative nesting depth can be overridden in `.env`. +The request limit intentionally applies only to public upload routes; internal +worker index transfers use their own trusted service boundary. + ## Speech-to-Text Settings DocsGPT can transcribe audio in two places: diff --git a/tests/api/user/attachments/test_routes.py b/tests/api/user/attachments/test_routes.py index 02949aa4..2c2b33e2 100644 --- a/tests/api/user/attachments/test_routes.py +++ b/tests/api/user/attachments/test_routes.py @@ -68,6 +68,29 @@ class FakeRedis: class TestStoreAttachmentEndpoint: + @patch("application.api.user.tasks.store_attachment.delay") + def test_store_attachment_rejects_oversized_non_audio_file( + self, mock_store_attachment, flask_app + ): + from application.api.user.attachments.routes import StoreAttachment + + app = Flask(__name__) + mock_storage = MagicMock() + with patch("application.api.user.base.storage", mock_storage), patch( + "application.upload_limits.settings.UPLOAD_MAX_FILE_BYTES", 4 + ), app.test_request_context( + "/api/store_attachment", + method="POST", + data={"file": (io.BytesIO(b"12345"), "large.txt")}, + content_type="multipart/form-data", + ): + request.decoded_token = {"sub": "test_user"} + response = StoreAttachment().post() + + assert _get_response_status(response) == 413 + mock_storage.save_file.assert_not_called() + mock_store_attachment.assert_not_called() + @patch("application.api.user.tasks.store_attachment.delay") def test_store_attachment_preserves_upload_indexes_for_partial_failures( self, mock_store_attachment, flask_app diff --git a/tests/api/user/sources/test_upload.py b/tests/api/user/sources/test_upload.py index fbe84a31..cabb92be 100644 --- a/tests/api/user/sources/test_upload.py +++ b/tests/api/user/sources/test_upload.py @@ -124,6 +124,34 @@ class TestUploadFile: assert response.json["success"] is True assert response.json["task_id"] == "task-1" + def test_rejects_file_over_configured_upload_limit(self, app): + from application.api.user.sources.upload import UploadFile + + fake_storage = MagicMock() + with patch( + "application.api.user.sources.upload.settings.UPLOAD_MAX_FILE_BYTES", + 4, + ), patch( + "application.api.user.sources.upload.StorageCreator.get_storage", + return_value=fake_storage, + ), patch( + "application.api.user.sources.upload.ingest.apply_async", + ) as apply_async, app.test_request_context( + "/api/upload", method="POST", + data={ + "user": "alice", "name": "job", + "file": (io.BytesIO(b"12345"), "large.txt"), + }, + content_type="multipart/form-data", + ): + from flask import request + request.decoded_token = {"sub": "alice"} + response = UploadFile().post() + + assert response.status_code == 413 + fake_storage.save_file.assert_not_called() + apply_async.assert_not_called() + def test_storage_error_returns_400(self, app): from application.api.user.sources.upload import UploadFile @@ -181,6 +209,140 @@ class TestUploadFile: # save_file called at least twice (for 2 files in zip) assert fake_storage.save_file.call_count >= 2 + def test_rejects_zip_with_excessive_expansion_ratio(self, app): + from application.api.user.sources.upload import UploadFile + import zipfile + + zip_buffer = io.BytesIO() + with zipfile.ZipFile(zip_buffer, "w", zipfile.ZIP_DEFLATED) as zf: + zf.writestr("bomb.txt", b"A" * (1024 * 1024)) + zip_buffer.seek(0) + + fake_storage = MagicMock() + with patch( + "application.api.user.sources.upload.StorageCreator.get_storage", + return_value=fake_storage, + ), patch( + "application.api.user.sources.upload.ingest.apply_async", + ) as apply_async, app.test_request_context( + "/api/upload", method="POST", + data={ + "user": "alice", "name": "job", + "file": (zip_buffer, "bomb.zip"), + }, + content_type="multipart/form-data", + ): + from flask import request + request.decoded_token = {"sub": "alice"} + response = UploadFile().post() + + assert response.status_code == 413 + fake_storage.save_file.assert_not_called() + apply_async.assert_not_called() + + def test_accepts_highly_compressible_csv(self, app): + from application.api.user.sources.upload import UploadFile + import zipfile + + zip_buffer = io.BytesIO() + with zipfile.ZipFile(zip_buffer, "w", zipfile.ZIP_DEFLATED) as archive: + archive.writestr("export.csv", b"a,b,c\n" * 900_000) + zip_buffer.seek(0) + + fake_storage = MagicMock() + with patch( + "application.api.user.sources.upload.StorageCreator.get_storage", + return_value=fake_storage, + ), patch( + "application.api.user.sources.upload.ingest.apply_async", + return_value=MagicMock(id="csv-task"), + ), app.test_request_context( + "/api/upload", + method="POST", + data={ + "user": "alice", + "name": "job", + "file": (zip_buffer, "export.zip"), + }, + content_type="multipart/form-data", + ): + from flask import request + request.decoded_token = {"sub": "alice"} + response = UploadFile().post() + + assert response.status_code == 200 + assert fake_storage.save_file.call_count == 1 + + def test_archive_error_names_file_and_escapes_log_controls(self, app): + from application.api.user.sources.upload import UploadFile + import zipfile + + zip_buffer = io.BytesIO() + with zipfile.ZipFile(zip_buffer, "w") as archive: + archive.writestr("../bad\nFORGED", b"bad") + zip_buffer.seek(0) + + with patch( + "application.api.user.sources.upload.StorageCreator.get_storage", + return_value=MagicMock(), + ), patch.object(app.logger, "warning") as warning, app.test_request_context( + "/api/upload", + method="POST", + data={ + "user": "alice", + "name": "job", + "file": (zip_buffer, "documents.zip"), + }, + content_type="multipart/form-data", + ): + from flask import request + request.decoded_token = {"sub": "alice"} + response = UploadFile().post() + + assert response.status_code == 413 + assert "documents.zip" in response.json["message"] + assert "\n" not in response.json["message"] + logged_values = " ".join(str(value) for value in warning.call_args.args) + assert "\n" not in logged_values + + def test_expands_nested_zip_before_storage(self, app): + from application.api.user.sources.upload import UploadFile + import zipfile + + inner = io.BytesIO() + with zipfile.ZipFile(inner, "w") as archive: + archive.writestr("inside.txt", b"inside") + outer = io.BytesIO() + with zipfile.ZipFile(outer, "w") as archive: + archive.writestr("inner.zip", inner.getvalue()) + outer.seek(0) + + fake_storage = MagicMock() + with patch( + "application.api.user.sources.upload.StorageCreator.get_storage", + return_value=fake_storage, + ), patch( + "application.api.user.sources.upload.ingest.apply_async", + return_value=MagicMock(id="nested-task"), + ), app.test_request_context( + "/api/upload", + method="POST", + data={ + "user": "alice", + "name": "job", + "file": (outer, "outer.zip"), + }, + content_type="multipart/form-data", + ): + from flask import request + request.decoded_token = {"sub": "alice"} + response = UploadFile().post() + + assert response.status_code == 200 + saved_paths = [call.args[1] for call in fake_storage.save_file.call_args_list] + assert any(path.endswith("/inner/inside.txt") for path in saved_paths) + assert not any(path.endswith(".zip") for path in saved_paths) + def test_office_format_zip_saved_as_is(self, app): from application.api.user.sources.upload import UploadFile import zipfile @@ -448,6 +610,36 @@ class TestManageSourceFiles: response = ManageSourceFiles().post() assert response.status_code == 400 + def test_add_rejects_file_over_upload_limit(self, app, pg_conn): + from application.api.user.sources.upload import ManageSourceFiles + + user = "u-add-large" + src = _seed_source(pg_conn, user=user, file_path="/data/src") + fake_storage = MagicMock() + with _patch_db(pg_conn), patch( + "application.api.user.sources.upload.StorageCreator.get_storage", + return_value=fake_storage, + ), patch( + "application.upload_limits.settings.UPLOAD_MAX_FILE_BYTES", 4 + ), patch( + "application.api.user.tasks.reingest_source_task.apply_async" + ) as apply_async, app.test_request_context( + "/api/manage_source_files", method="POST", + data={ + "source_id": str(src["id"]), + "operation": "add", + "file": (io.BytesIO(b"12345"), "large.txt"), + }, + content_type="multipart/form-data", + ): + from flask import request + request.decoded_token = {"sub": user} + response = ManageSourceFiles().post() + + assert response.status_code == 413 + fake_storage.save_file.assert_not_called() + apply_async.assert_not_called() + def test_add_files_success(self, app, pg_conn): from application.api.user.sources.upload import ManageSourceFiles diff --git a/tests/api/user/test_tools_routes.py b/tests/api/user/test_tools_routes.py index 936b1f6d..d3164591 100644 --- a/tests/api/user/test_tools_routes.py +++ b/tests/api/user/test_tools_routes.py @@ -933,6 +933,46 @@ class TestParseSpec: assert response.status_code == 200 assert response.json["success"] is True + def test_rejects_oversized_spec_upload(self, app): + from application.api.user.tools.routes import ParseSpec + from io import BytesIO + + with patch( + "application.api.user.tools.routes.settings.PARSE_SPEC_MAX_BYTES", 4 + ), patch( + "application.api.user.tools.routes.parse_spec" + ) as parse_mock, app.test_request_context( + "/api/parse_spec", + method="POST", + content_type="multipart/form-data", + data={"file": (BytesIO(b"12345"), "spec.yaml")}, + ): + from flask import request + request.decoded_token = {"sub": "user1"} + response = ParseSpec().post() + + assert response.status_code == 413 + parse_mock.assert_not_called() + + def test_rejects_oversized_json_spec(self, app): + from application.api.user.tools.routes import ParseSpec + + with patch( + "application.api.user.tools.routes.settings.PARSE_SPEC_MAX_BYTES", 4 + ), patch( + "application.api.user.tools.routes.parse_spec" + ) as parse_mock, app.test_request_context( + "/api/parse_spec", + method="POST", + json={"spec_content": "12345"}, + ): + from flask import request + request.decoded_token = {"sub": "user1"} + response = ParseSpec().post() + + assert response.status_code == 413 + parse_mock.assert_not_called() + def test_returns_400_file_no_filename(self, app): from application.api.user.tools.routes import ParseSpec from io import BytesIO @@ -1689,5 +1729,3 @@ class TestDualRegisteredToggle: assert "scheduler" in ( user_doc["tool_preferences"]["disabled_default_tools"] ) - - diff --git a/tests/test_app_routes.py b/tests/test_app_routes.py index d5d5d2af..d5ef5be6 100644 --- a/tests/test_app_routes.py +++ b/tests/test_app_routes.py @@ -147,6 +147,58 @@ class TestSttRequestSizeLimits: assert response.status_code == 413 +class TestDocumentUploadRequestSizeLimits: + + @pytest.mark.unit + def test_oversized_upload_rejected_before_multipart_parsing(self, client): + with patch( + "application.app.settings.UPLOAD_MAX_REQUEST_BYTES", + 32, + ): + response = client.post( + "/api/upload", + data=b"x" * 100, + content_type="multipart/form-data", + ) + + assert response.status_code == 413 + assert response.get_json() == { + "success": False, + "message": "Request exceeds the 32-byte upload limit", + } + + @pytest.mark.unit + def test_internal_worker_upload_is_not_subject_to_user_limit(self, client): + with patch( + "application.app.settings.UPLOAD_MAX_REQUEST_BYTES", + 32, + ): + response = client.post( + "/api/upload_index", + data=b"x" * 100, + content_type="application/octet-stream", + ) + + assert response.status_code != 413 + + @pytest.mark.unit + def test_json_spec_uses_dedicated_request_limit(self, client): + with patch( + "application.app.settings.PARSE_SPEC_MAX_BYTES", + 32, + ): + response = client.post( + "/api/parse_spec", + json={"spec_content": "x" * 100}, + ) + + assert response.status_code == 413 + assert response.get_json() == { + "success": False, + "message": "Request exceeds the 32-byte upload limit", + } + + class TestAuthenticateRequest: @pytest.mark.unit diff --git a/tests/test_upload_limits.py b/tests/test_upload_limits.py new file mode 100644 index 00000000..46818421 --- /dev/null +++ b/tests/test_upload_limits.py @@ -0,0 +1,39 @@ +"""Tests for bounded user-upload stream helpers.""" + +import io + +import pytest + +from application.upload_limits import ( + copy_upload_to_path, + read_upload_limited, + read_text_upload_limited, + UploadTooLargeError, +) + + +class _ShortReadStream(io.BytesIO): + """Return small chunks even when the caller asks for more.""" + + def read(self, size=-1): + return super().read(min(size, 2) if size >= 0 else 2) + + +def test_limited_read_accumulates_short_reads_before_rejecting(): + with pytest.raises(UploadTooLargeError): + read_upload_limited(_ShortReadStream(b"12345"), max_bytes=4) + + +def test_limited_copy_rejects_before_writing_overflow_byte(tmp_path): + target = tmp_path / "upload.bin" + with pytest.raises(UploadTooLargeError): + copy_upload_to_path(io.BytesIO(b"12345"), target, max_bytes=4) + + assert target.stat().st_size <= 4 + + +def test_limited_text_read_decodes_incrementally_and_rejects_overflow(): + assert read_text_upload_limited(io.BytesIO("café".encode()), max_bytes=5) == "café" + + with pytest.raises(UploadTooLargeError): + read_text_upload_limited(_ShortReadStream(b"12345"), max_bytes=4) diff --git a/tests/test_worker_utils.py b/tests/test_worker_utils.py index e8c11007..3b8f033d 100644 --- a/tests/test_worker_utils.py +++ b/tests/test_worker_utils.py @@ -203,6 +203,29 @@ class TestExtractZipRecursive: found = list(extract_to.rglob("inside.txt")) assert found, "expected nested zip to be extracted" + def test_nested_archives_share_one_expansion_budget(self, tmp_path): + from application.worker import extract_zip_recursive + + inner = tmp_path / "inner.zip" + with zipfile.ZipFile(inner, "w", zipfile.ZIP_STORED) as zf: + zf.writestr("inside.txt", b"x" * 60) + + outer = tmp_path / "outer.zip" + with zipfile.ZipFile(outer, "w", zipfile.ZIP_STORED) as zf: + zf.write(inner, "inner.zip") + + # The outer archive fits, as does the inner archive in isolation, but + # extracting both exceeds one cumulative request budget. + limit = inner.stat().st_size + 30 + extract_to = tmp_path / "out" + extract_to.mkdir() + with patch("application.worker.MAX_UNCOMPRESSED_SIZE", limit), patch( + "application.worker.MAX_COMPRESSION_RATIO", 10_000, + ): + extract_zip_recursive(str(outer), str(extract_to)) + + assert not (extract_to / "inside.txt").exists() + class TestDownloadFile: def test_writes_file_on_success(self, tmp_path): diff --git a/tests/test_zip_extraction_security.py b/tests/test_zip_extraction_security.py index c53452f7..65214aa7 100644 --- a/tests/test_zip_extraction_security.py +++ b/tests/test_zip_extraction_security.py @@ -1,6 +1,7 @@ """Tests for zip extraction security measures.""" import os +import stat import tempfile import zipfile @@ -13,6 +14,124 @@ from application.worker import ( extract_zip_recursive, MAX_FILE_COUNT, ) +from application.security.zip_archive import ( + extract_zip_safely, + ZipExtractionLimits, +) + + +def _test_limits(**overrides): + values = { + "max_uncompressed_bytes": 20 * 1024 * 1024, + "max_files": 100, + "max_compression_ratio": 2_000, + "max_member_bytes": 10 * 1024 * 1024, + "max_depth": 3, + } + values.update(overrides) + return ZipExtractionLimits(**values) + + +class TestTransactionalZipExtraction: + def test_default_file_budget_accepts_large_source_archives(self, tmp_path): + from application.api.user.sources.upload import _source_archive_limits + from application.security.zip_archive import validate_zip_archive + + zip_path = tmp_path / "repository.zip" + with zipfile.ZipFile(zip_path, "w") as archive: + for index in range(1_845): + archive.writestr(f"src/file-{index}.txt", b"x") + + validate_zip_archive(zip_path, _source_archive_limits()) + + def test_directory_records_and_case_distinct_files_are_allowed(self, tmp_path): + zip_path = tmp_path / "portable.zip" + with zipfile.ZipFile(zip_path, "w") as archive: + archive.writestr("docs/", b"") + with pytest.warns(UserWarning, match="Duplicate name"): + archive.writestr("docs/", b"") + archive.writestr("docs/README.txt", b"upper") + archive.writestr("docs/readme.txt", b"lower") + + destination = tmp_path / "out" + extract_zip_safely(zip_path, destination, _test_limits(max_files=2)) + + extracted_contents = { + path.read_bytes() for path in (destination / "docs").iterdir() + } + assert extracted_contents == {b"upper", b"lower"} + + def test_failure_leaves_destination_unchanged(self, tmp_path): + zip_path = tmp_path / "corrupt.zip" + with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_STORED) as archive: + archive.writestr("first.txt", b"FIRST") + archive.writestr("second.txt", b"SECOND") + + payload = bytearray(zip_path.read_bytes()) + second_offset = payload.index(b"SECOND") + payload[second_offset] ^= 0xFF + zip_path.write_bytes(payload) + + destination = tmp_path / "out" + destination.mkdir() + marker = destination / "existing.txt" + marker.write_text("keep", encoding="utf-8") + + with pytest.raises(ZipExtractionError): + extract_zip_safely(zip_path, destination, _test_limits()) + + assert marker.read_text(encoding="utf-8") == "keep" + assert not (destination / "first.txt").exists() + assert not (destination / "second.txt").exists() + + def test_nested_archives_are_expanded_with_one_budget(self, tmp_path): + inner_path = tmp_path / "inner.zip" + with zipfile.ZipFile(inner_path, "w") as archive: + archive.writestr("inside.txt", b"nested") + + outer_path = tmp_path / "outer.zip" + with zipfile.ZipFile(outer_path, "w") as archive: + archive.writestr("outer.txt", b"outer") + archive.writestr("inner.zip", inner_path.read_bytes()) + + destination = tmp_path / "out" + budget = extract_zip_safely(outer_path, destination, _test_limits()) + + assert (destination / "outer.txt").read_bytes() == b"outer" + assert (destination / "inner" / "inside.txt").read_bytes() == b"nested" + assert not (destination / "inner.zip").exists() + assert budget.files == 3 + + def test_nested_archive_uses_private_directory_for_colliding_names( + self, tmp_path + ): + inner_path = tmp_path / "n.zip" + with zipfile.ZipFile(inner_path, "w") as archive: + archive.writestr("a.txt", b"inner") + + outer_path = tmp_path / "outer.zip" + with zipfile.ZipFile(outer_path, "w") as archive: + archive.writestr("a.txt", b"outer") + archive.writestr("n.zip", inner_path.read_bytes()) + + destination = tmp_path / "out" + extract_zip_safely(outer_path, destination, _test_limits()) + + assert (destination / "a.txt").read_bytes() == b"outer" + assert (destination / "n" / "a.txt").read_bytes() == b"inner" + + @pytest.mark.parametrize("payload", [b"plain text", b"PK\x03\x04truncated"]) + def test_non_zip_member_with_zip_suffix_is_preserved(self, tmp_path, payload): + outer_path = tmp_path / "outer.zip" + with zipfile.ZipFile(outer_path, "w") as archive: + archive.writestr("readme.txt", b"content") + archive.writestr("notes.zip", payload) + + destination = tmp_path / "out" + extract_zip_safely(outer_path, destination, _test_limits()) + + assert (destination / "readme.txt").read_bytes() == b"content" + assert (destination / "notes.zip").read_bytes() == payload class TestIsPathSafe: @@ -161,7 +280,7 @@ class TestExtractZipRecursive: assert os.path.exists(os.path.join(extract_to, "outer.txt")) # Check inner zip was extracted - assert os.path.exists(os.path.join(extract_to, "inner.txt")) + assert os.path.exists(os.path.join(extract_to, "inner", "inner.txt")) # Check both zips were removed assert not os.path.exists(zip_path) @@ -216,6 +335,23 @@ class TestExtractZipRecursive: # Malicious file should NOT exist outside extract_to assert not os.path.exists("/tmp/malicious.txt") + def test_rejects_symlink_members(self): + """Archive-controlled links must never be materialized.""" + with tempfile.TemporaryDirectory() as temp_dir: + zip_path = os.path.join(temp_dir, "links.zip") + extract_to = os.path.join(temp_dir, "extract") + os.makedirs(extract_to) + + link = zipfile.ZipInfo("escape-link") + link.create_system = 3 + link.external_attr = (stat.S_IFLNK | 0o777) << 16 + with zipfile.ZipFile(zip_path, "w") as zf: + zf.writestr(link, "../../outside") + + extract_zip_recursive(zip_path, extract_to) + + assert not os.path.lexists(os.path.join(extract_to, "escape-link")) + def test_handles_corrupted_zip_gracefully(self): """Corrupted zip should be handled gracefully without crashing.""" with tempfile.TemporaryDirectory() as temp_dir: