mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-13 14:15:52 +00:00
fix: better zip protection
This commit is contained in:
1 parent
8dd4c9a417
commit
663869868f
16 files changed
+1401
-146
No files matched your search
@@ -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,
|
||||
}
|
||||
),
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
+56
-77
@@ -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,
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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"]
|
||||
)
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
@@ -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):
|
||||
|
||||
@@ -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:
|
||||
|
||||
Reference in new issue
Block a user