diff --git a/.env-template b/.env-template index afff3835..81ba3d42 100644 --- a/.env-template +++ b/.env-template @@ -40,7 +40,8 @@ MICROSOFT_AUTHORITY=https://{tenantId}.ciamlogin.com/{tenantId} # Authentication (optional - default is no auth; see docs: Deploying -> App Configuration) # AUTH_TYPE=None|simple_jwt|session_jwt|oidc -# JWT_SECRET_KEY= +# JWT_SECRET_KEY= +# Local development auto-generates .jwt_secret_key when this is unset. # OIDC SSO (only when AUTH_TYPE=oidc; works with Authentik, Keycloak, Okta, ...) # OIDC_ISSUER= diff --git a/application/api/user/agents/routes.py b/application/api/user/agents/routes.py index 62cc6ee7..fc8f3753 100644 --- a/application/api/user/agents/routes.py +++ b/application/api/user/agents/routes.py @@ -11,6 +11,7 @@ from pydantic import ValidationError as PydanticValidationError from application.api import api from application.guardrails.config import AgentConfig from application.api.user.base import ( + copy_agent_image_for_user, handle_image_upload, resolve_prompt_name, resolve_source_details, @@ -244,7 +245,11 @@ def _format_agent_output( "slug": agent.get("slug", "") or "", "description": agent.get("description", "") or "", "image": ( - generate_image_url(agent["image"]) if agent.get("image") else "" + generate_image_url( + agent["image"], agent["id"], agent.get("user_id") + ) + if agent.get("image") + else "" ), "source": source_value, "sources": sources_list, @@ -308,7 +313,7 @@ def _build_create_kwargs(data: dict, *, image_url: str, agent_type: str) -> dict allowed_fields = set(schema["fields"]) for key in ( - "description", "agent_type", "key", "image", "retriever", + "description", "agent_type", "key", "retriever", "default_model_id", ): if key in allowed_fields and data.get(key) not in (None, ""): @@ -787,8 +792,8 @@ class UpdateAgent(Resource): "description": fields.String( required=True, description="New description of the agent" ), - "image": fields.String( - required=False, description="New image URL or identifier" + "image": fields.Raw( + required=False, description="Image file upload", type="file" ), "source": fields.String( required=False, description="Source ID (legacy single source)" @@ -913,8 +918,12 @@ class UpdateAgent(Resource): 404, ) pg_agent_id = str(existing_agent["id"]) + existing_image = existing_agent.get("image", "") or "" image_url, image_error = handle_image_upload( - request, existing_agent.get("image", "") or "", user, storage, + request, + existing_image, + existing_agent.get("user_id") or user, + storage, ) if image_error: return image_error @@ -923,7 +932,6 @@ class UpdateAgent(Resource): allowed_fields = [ "name", "description", - "image", "source", "sources", "chunks", @@ -1133,7 +1141,7 @@ class UpdateAgent(Resource): f"Field '{field}' cannot be empty", user, field ) update_fields[field] = value - if image_url: + if image_url and image_url != existing_image: update_fields["image"] = image_url if not update_fields: return _reject("No valid update data provided", user) @@ -1543,7 +1551,11 @@ class PinnedAgents(Resource): "name": agent.get("name", ""), "description": agent.get("description", ""), "image": ( - generate_image_url(agent["image"]) if agent.get("image") else "" + generate_image_url( + agent["image"], agent["id"], agent.get("user_id") + ) + if agent.get("image") + else "" ), "source": str(source_id) if source_id else "", "chunks": str(agent["chunks"]) if agent.get("chunks") is not None else "", @@ -1591,7 +1603,13 @@ class GetTemplateAgents(Resource): "id": str(agent["id"]), "name": agent.get("name"), "description": agent.get("description") or "", - "image": agent.get("image") or "", + "image": ( + generate_image_url( + agent["image"], agent["id"], agent.get("user_id") + ) + if agent.get("image") + else "" + ), } for agent in template_rows ] @@ -1647,7 +1665,7 @@ class AdoptAgent(Resource): # and must get its own values (slug stays NULL until exported). create_kwargs: dict = {} for col in ( - "description", "agent_type", "image", "retriever", + "description", "agent_type", "retriever", "default_model_id", "source_id", "prompt_id", "folder_id", "workflow_id", "extra_source_ids", @@ -1655,6 +1673,16 @@ class AdoptAgent(Resource): val = template.get(col) if val not in (None, ""): create_kwargs[col] = val + + # The avatar must live under the adopter's own upload + # directory: image paths are validated against their owner, + # so a copied reference to the template owner's blob would + # fail closed and render blank. + adopted_image = copy_agent_image_for_user( + template.get("image"), user, storage + ) + if adopted_image: + create_kwargs["image"] = adopted_image for col in ("tools", "json_schema", "models", "shared_metadata", "config"): if template.get(col) is not None: create_kwargs[col] = template[col] diff --git a/application/api/user/agents/sharing.py b/application/api/user/agents/sharing.py index 2d1758b4..df797733 100644 --- a/application/api/user/agents/sharing.py +++ b/application/api/user/agents/sharing.py @@ -29,7 +29,11 @@ def _serialize_agent_basic(agent: dict) -> dict: "user": agent.get("user_id", ""), "name": agent.get("name", ""), "image": ( - generate_image_url(agent["image"]) if agent.get("image") else "" + generate_image_url( + agent["image"], agent["id"], agent.get("user_id") + ) + if agent.get("image") + else "" ), "description": agent.get("description", ""), "source": str(source_id) if source_id else "", @@ -158,7 +162,11 @@ class SharedAgents(Resource): "name": agent.get("name", ""), "description": agent.get("description", ""), "image": ( - generate_image_url(agent["image"]) if agent.get("image") else "" + generate_image_url( + agent["image"], agent["id"], agent.get("user_id") + ) + if agent.get("image") + else "" ), "tools": agent.get("tools", []) or [], "tool_details": resolve_tool_details( diff --git a/application/api/user/attachments/routes.py b/application/api/user/attachments/routes.py index 32e876e7..f3e7a8b5 100644 --- a/application/api/user/attachments/routes.py +++ b/application/api/user/attachments/routes.py @@ -6,12 +6,14 @@ from pathlib import Path import uuid -from flask import current_app, jsonify, make_response, request +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 from application.core.settings import settings +from application.storage.db.base_repository import looks_like_uuid from application.storage.db.repositories.agents import AgentsRepository from application.storage.db.session import db_readonly from application.stt.constants import ( @@ -35,13 +37,42 @@ from application.stt.live_session import ( ) from application.stt.stt_creator import STTCreator from application.tts.tts_creator import TTSCreator -from application.utils import safe_filename +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, + is_safe_agent_image_path, + safe_filename, + verify_agent_image_capability, +) attachments_ns = Namespace( "attachments", description="File attachments and media operations", path="/api" ) +_AGENT_IMAGE_CHUNK_BYTES = 64 * 1024 +_AGENT_IMAGE_PRESIGNED_TTL_SECONDS = 300 +_AGENT_IMAGE_REDIRECT_CACHE_SECONDS = _AGENT_IMAGE_PRESIGNED_TTL_SECONDS - 60 + + +def _stream_file(file_obj, size_bytes: int): + """Yield up to the size-checked file range and always close the handle.""" + remaining = size_bytes + try: + while remaining > 0: + chunk = file_obj.read(min(_AGENT_IMAGE_CHUNK_BYTES, remaining)) + if not chunk: + break + remaining -= len(chunk) + yield chunk + finally: + file_obj.close() + def _resolve_authenticated_user(): decoded_token = getattr(request, "decoded_token", None) @@ -97,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" @@ -175,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), @@ -199,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, } ), @@ -630,25 +680,70 @@ class LiveSpeechToTextFinish(Resource): ) -@attachments_ns.route("/images/") +@attachments_ns.route("/images//") class ServeImage(Resource): - @api.doc(description="Serve an image from storage") - def get(self, image_path): - if ".." in image_path or image_path.startswith("/") or "\x00" in image_path: - return make_response( - jsonify({"success": False, "message": "Invalid image path"}), 400 - ) + @api.doc(description="Serve an agent image using an opaque capability URL") + def get(self, agent_id, capability): try: + if not looks_like_uuid(agent_id): + return make_response( + jsonify({"success": False, "message": "Image not found"}), 404 + ) + with db_readonly() as conn: + agent = AgentsRepository(conn).find_image_record(agent_id) + if not agent: + return make_response( + jsonify({"success": False, "message": "Image not found"}), 404 + ) + + image_path = agent.get("image") + user_id = agent.get("user_id") + content_type = get_agent_image_content_type(image_path) + if ( + is_external_image_url(image_path) + or not content_type + or not is_safe_agent_image_path(image_path, user_id) + or not verify_agent_image_capability( + capability, agent_id, image_path, user_id + ) + ): + return make_response( + jsonify({"success": False, "message": "Image not found"}), 404 + ) + from application.api.user.base import storage + size_bytes = storage.get_file_size(image_path) + if size_bytes < 0 or size_bytes > settings.AGENT_IMAGE_MAX_BYTES: + current_app.logger.warning( + "Rejected oversized legacy agent image for agent %s", agent_id + ) + return make_response( + jsonify({"success": False, "message": "Image not found"}), 404 + ) + + if settings.STORAGE_TYPE == "s3": + image_url = storage.generate_presigned_url( + image_path, + expires_in=_AGENT_IMAGE_PRESIGNED_TTL_SECONDS, + content_type=content_type, + ) + response = redirect(image_url, code=302) + response.headers.set( + "Cache-Control", + f"private, max-age={_AGENT_IMAGE_REDIRECT_CACHE_SECONDS}", + ) + response.headers.set("X-Content-Type-Options", "nosniff") + return response + file_obj = storage.get_file(image_path) - extension = image_path.split(".")[-1].lower() - content_type = f"image/{extension}" - if extension == "jpg": - content_type = "image/jpeg" - response = make_response(file_obj.read()) - response.headers.set("Content-Type", content_type) - response.headers.set("Cache-Control", "max-age=86400") + response = Response( + _stream_file(file_obj, size_bytes), + content_type=content_type, + direct_passthrough=True, + ) + response.headers.set("Cache-Control", "public, max-age=86400, immutable") + response.headers.set("X-Content-Type-Options", "nosniff") return response except FileNotFoundError: @@ -657,7 +752,7 @@ class ServeImage(Resource): ) except ValueError: return make_response( - jsonify({"success": False, "message": "Invalid image path"}), 400 + jsonify({"success": False, "message": "Image not found"}), 404 ) except Exception as e: current_app.logger.error(f"Error serving image: {e}") diff --git a/application/api/user/base.py b/application/api/user/base.py index 5aaa478e..82e3c313 100644 --- a/application/api/user/base.py +++ b/application/api/user/base.py @@ -6,9 +6,11 @@ import datetime import os import uuid from functools import wraps +from pathlib import PurePosixPath from typing import Optional, Tuple from flask import current_app, jsonify, make_response, Response +from PIL import Image, UnidentifiedImageError from werkzeug.utils import secure_filename from sqlalchemy import text as _sql_text @@ -18,6 +20,12 @@ from application.storage.db.base_repository import looks_like_uuid, row_to_dict from application.storage.db.repositories.users import UsersRepository from application.storage.db.session import db_readonly, db_session from application.storage.storage_creator import StorageCreator +from application.utils import ( + AGENT_IMAGE_FORMATS, + get_agent_image_content_type, + is_external_image_url, + safe_user_storage_component, +) from application.vectorstore.vector_creator import VectorCreator @@ -28,7 +36,6 @@ current_dir = os.path.dirname( os.path.dirname(os.path.dirname(os.path.abspath(__file__))) ) - def generate_minute_range(start_date, end_date): """Generate a dictionary with minute-level time ranges.""" return { @@ -237,9 +244,112 @@ def get_vector_store(source_id): return store +def _validate_agent_image_upload(file, filename: str) -> None: + """Validate an uploaded agent avatar and rewind it for storage. + + Args: + file: Werkzeug file upload object. + filename: Sanitized upload filename. + + Raises: + OSError: If the upload stream cannot be inspected. + ValueError: If the file exceeds a limit or is not an allowed raster image. + """ + extension = PurePosixPath(filename).suffix.lower() + image_policy = AGENT_IMAGE_FORMATS.get(extension) + if not filename or image_policy is None: + raise ValueError("Unsupported image extension") + + stream = file.stream + stream.seek(0, os.SEEK_END) + size_bytes = stream.tell() + stream.seek(0) + if size_bytes > settings.AGENT_IMAGE_MAX_BYTES: + raise ValueError("Image exceeds the upload size limit") + + with Image.open(stream) as image: + image_format = (image.format or "").upper() + if image_format not in image_policy[0]: + raise ValueError("Image content does not match its extension") + if image.width * image.height > settings.AGENT_IMAGE_MAX_PIXELS: + raise ValueError("Image exceeds the pixel limit") + image.verify() + stream.seek(0) + + +def _safe_agent_image_filename(filename: str) -> str: + """Build an ASCII storage filename without discarding a Unicode file's suffix. + + Args: + filename: Original multipart upload filename. + + Returns: + A traversal-safe filename with a normalized, allow-listed extension. + + Raises: + ValueError: If the original filename has no supported image extension. + """ + normalized_path = filename.replace("\\", "/") + original_path = PurePosixPath(normalized_path) + extension = original_path.suffix.lower() + if extension not in AGENT_IMAGE_FORMATS: + raise ValueError("Unsupported image extension") + + safe_stem = secure_filename(original_path.stem).strip("._") or "avatar" + return f"{safe_stem}{extension}" + + +def copy_agent_image_for_user(image_path: object, user: str, storage) -> str: + """Copy an internal agent avatar into ``user``'s own attachments directory. + + Avatar paths are validated against their owner's upload directory, so an + agent copied to a new owner needs its own blob rather than a reference to + the original owner's. + + Args: + image_path: Source avatar storage path, or an external image URL. + user: Identity that will own the copy. + storage: Storage backend used for both the read and the write. + + Returns: + The newly owned storage path, an unchanged external URL, or an empty + string when the source is unusable. + """ + if not isinstance(image_path, str) or not image_path: + return "" + if is_external_image_url(image_path): + return image_path + if get_agent_image_content_type(image_path) is None: + return "" + + source_name = PurePosixPath(image_path).name + prefix, separator, remainder = source_name.partition("_") + if separator and remainder and looks_like_uuid(prefix): + source_name = remainder + try: + filename = _safe_agent_image_filename(source_name) + except ValueError: + return "" + + owner_component = safe_user_storage_component(user) + destination = ( + f"{settings.UPLOAD_FOLDER.rstrip('/')}/{owner_component}/" + f"attachments/{uuid.uuid4()}_{filename}" + ) + try: + with storage.get_file(image_path) as source_file: + storage.save_file(source_file, destination, storage_class="STANDARD") + except Exception as e: + current_app.logger.warning( + "Could not copy agent image %s for %s: %s", image_path, user, e + ) + return "" + return destination + + def handle_image_upload( - request, existing_url: str, user: str, storage, base_path: str = "attachments/" -) -> Tuple[str, Optional[Response]]: + request, existing_url: str, user: str, storage +) -> Tuple[Optional[str], Optional[Response]]: """ Handle image file upload from request. @@ -248,7 +358,6 @@ def handle_image_upload( existing_url: Existing image URL (fallback) user: User ID storage: Storage instance - base_path: Base path for upload Returns: Tuple of (image_url, error_response) @@ -258,8 +367,26 @@ def handle_image_upload( if "image" in request.files: file = request.files["image"] if file.filename != "": - filename = secure_filename(file.filename) - upload_path = f"{settings.UPLOAD_FOLDER.rstrip('/')}/{user}/{base_path.rstrip('/')}/{uuid.uuid4()}_{filename}" + try: + filename = _safe_agent_image_filename(file.filename) + _validate_agent_image_upload(file, filename) + except ( + Image.DecompressionBombError, + UnidentifiedImageError, + ValueError, + OSError, + ) as e: + current_app.logger.warning(f"Invalid agent image upload: {e}") + return None, make_response( + jsonify({"success": False, "message": "Invalid image upload"}), + 400, + ) + + owner_component = safe_user_storage_component(user) + upload_path = ( + f"{settings.UPLOAD_FOLDER.rstrip('/')}/{owner_component}/" + f"attachments/{uuid.uuid4()}_{filename}" + ) try: storage.save_file(file, upload_path, storage_class="STANDARD") image_url = upload_path diff --git a/application/api/user/sources/upload.py b/application/api/user/sources/upload.py index 1cd41fc5..4fcf9cc5 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,34 @@ class UploadFile(Resource): ), 413, ) + except (UploadTooLargeError, ZipExtractionError) as err: + if isinstance(err, ZipExtractionError): + zip_error_detail = safe_zip_error_message(err) + 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 archive failed safety validation" + 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 +725,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 +1046,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 bf27bf52..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 @@ -28,12 +29,17 @@ from application.api.user.routes import user # noqa: E402 from application.api.connector.routes import connector # noqa: E402 from application.api.v1 import v1_bp # noqa: E402 from application.celery_init import celery # noqa: E402 +from application.core.secret_key import resolve_jwt_secret_key # noqa: E402 from application.core.settings import settings # noqa: E402 from application.storage.db.bootstrap import ensure_database_ready # noqa: E402 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": @@ -84,18 +90,59 @@ app.config.update( celery.config_from_object("application.celeryconfig") api.init_app(app) -if settings.AUTH_TYPE in ("simple_jwt", "session_jwt", "oidc") and not settings.JWT_SECRET_KEY: - key_file = ".jwt_secret_key" - try: - with open(key_file, "r") as f: - settings.JWT_SECRET_KEY = f.read().strip() - except FileNotFoundError: - new_key = os.urandom(32).hex() - with open(key_file, "w") as f: - f.write(new_key) - settings.JWT_SECRET_KEY = new_key - except Exception as e: - raise RuntimeError(f"Failed to setup JWT_SECRET_KEY: {e}") + +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. +settings.JWT_SECRET_KEY = resolve_jwt_secret_key( + settings.JWT_SECRET_KEY, + os.getenv("DEPLOYMENT_TYPE"), +) if settings.AUTH_TYPE == "oidc": _missing_oidc = [ name diff --git a/application/core/secret_key.py b/application/core/secret_key.py new file mode 100644 index 00000000..fdc50f93 --- /dev/null +++ b/application/core/secret_key.py @@ -0,0 +1,89 @@ +"""Resolve stable application signing secrets for local and production use.""" + +from __future__ import annotations + +import os +import tempfile +from pathlib import Path + + +_SHARED_SECRET_DEPLOYMENTS = {"cloud", "production"} + + +def _read_secret_file(key_path: Path) -> str: + """Read and validate a generated local signing secret.""" + secret = key_path.read_text(encoding="utf-8").strip() + if not secret: + raise RuntimeError(f"Signing secret file is empty: {key_path}") + return secret + + +def _create_secret_file_atomically(key_path: Path) -> str: + """Create a complete mode-0600 secret before atomically publishing it.""" + new_secret = os.urandom(32).hex() + descriptor, temporary_name = tempfile.mkstemp( + dir=key_path.parent, + prefix=f".{key_path.name}.", + text=True, + ) + temporary_path = Path(temporary_name) + try: + if hasattr(os, "fchmod"): + os.fchmod(descriptor, 0o600) + with os.fdopen(descriptor, "w", encoding="utf-8") as temporary_file: + descriptor = -1 + temporary_file.write(new_secret) + temporary_file.flush() + os.fsync(temporary_file.fileno()) + + try: + # A hard link publishes the fully written file without replacing a + # secret another process may have won the race to create. + os.link(temporary_path, key_path) + except FileExistsError: + return _read_secret_file(key_path) + return new_secret + finally: + if descriptor >= 0: + os.close(descriptor) + temporary_path.unlink(missing_ok=True) + + +def resolve_jwt_secret_key( + configured_secret: str | None, + deployment_type: str | None, + key_file: str | Path = ".jwt_secret_key", +) -> str: + """Return a shared configured secret or a stable local-development secret. + + Args: + configured_secret: Operator-supplied signing secret. + deployment_type: Deployment class, such as ``cloud`` or ``production``. + key_file: Local fallback file used outside production deployments. + + Returns: + The configured or locally persisted signing secret. + + Raises: + RuntimeError: If production lacks a shared secret or local persistence fails. + """ + if configured_secret and configured_secret.strip(): + return configured_secret + + normalized_deployment = (deployment_type or "").strip().lower() + if normalized_deployment in _SHARED_SECRET_DEPLOYMENTS: + raise RuntimeError( + "JWT_SECRET_KEY must be set to the same strong random value on every " + f"{normalized_deployment} API and worker replica" + ) + + key_path = Path(key_file) + try: + return _read_secret_file(key_path) + except FileNotFoundError: + try: + return _create_secret_file_atomically(key_path) + except Exception as exc: + raise RuntimeError(f"Failed to create signing secret: {exc}") from exc + except Exception as exc: + raise RuntimeError(f"Failed to read signing secret: {exc}") from exc diff --git a/application/core/settings.py b/application/core/settings.py index 12d8924f..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) @@ -99,6 +109,8 @@ class Settings(BaseSettings): DOCLING_TABULAR_MAX_BYTES: int = 2_000_000 DOCLING_MARKUP_MAX_BYTES: int = 8_000_000 ATTACHMENT_TEXT_MAX_BYTES: int = 5_000_000 + AGENT_IMAGE_MAX_BYTES: int = 5_000_000 + AGENT_IMAGE_MAX_PIXELS: int = 16_777_216 VECTOR_STORE: str = "faiss" # "faiss" or "elasticsearch" or "qdrant" or "milvus" or "lancedb" or "pgvector" # Allow-list of retriever keys an agent may use. Values must match the # ``RetrieverCreator.retrievers`` registry keys (``classic`` / ``default``), diff --git a/application/security/zip_archive.py b/application/security/zip_archive.py new file mode 100644 index 00000000..4d3b5527 --- /dev/null +++ b/application/security/zip_archive.py @@ -0,0 +1,401 @@ +"""Safe, bounded extraction for untrusted ZIP archives.""" + +from __future__ import annotations + +import os +import shutil +import stat +import tempfile +import zipfile +from contextlib import suppress +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: + # Rollback is best effort and must not hide the original commit failure. + for target in reversed(moved_files): + with suppress(OSError): + os.remove(target) + for target in reversed(created_directories): + with suppress(OSError): + os.rmdir(target) + 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/storage/base.py b/application/storage/base.py index 1dae6506..af505014 100644 --- a/application/storage/base.py +++ b/application/storage/base.py @@ -1,7 +1,7 @@ """Base storage class for file system abstraction.""" from abc import ABC, abstractmethod -from typing import BinaryIO, List, Callable +from typing import BinaryIO, Callable, List, Optional class BaseStorage(ABC): @@ -37,12 +37,41 @@ class BaseStorage(ABC): """ pass - def generate_presigned_url(self, path: str, expires_in: int = 300) -> str: + def get_file_size(self, path: str) -> int: + """Return a file's encoded size without retaining its bytes in memory. + + Storage backends should override this with a metadata-only lookup. The + compatibility fallback seeks an already-open file and is suitable for + file-like backends that do not provide object metadata. + + Args: + path: Path to the file. + + Returns: + File size in bytes. + """ + file_obj = self.get_file(path) + try: + current = file_obj.tell() + file_obj.seek(0, 2) + size = file_obj.tell() + file_obj.seek(current) + return size + finally: + file_obj.close() + + def generate_presigned_url( + self, + path: str, + expires_in: int = 300, + content_type: Optional[str] = None, + ) -> str: """Return a short-lived presigned download URL; not all backends support it. Args: path: Path to the file expires_in: TTL of the signed URL in seconds + content_type: Optional response Content-Type override Returns: str: A presigned URL granting time-limited read access diff --git a/application/storage/db/repositories/agents.py b/application/storage/db/repositories/agents.py index 22cb5653..cdec8a2a 100644 --- a/application/storage/db/repositories/agents.py +++ b/application/storage/db/repositories/agents.py @@ -176,6 +176,25 @@ class AgentsRepository: row = result.fetchone() return row_to_dict(row) if row is not None else None + def find_image_record(self, agent_id: str) -> Optional[dict]: + """Fetch only the fields needed to verify a public image capability. + + This lookup is intentionally not owner-scoped because the resulting + row is not authorization: callers must verify its HMAC capability + before reading or returning any storage bytes. + """ + if not looks_like_uuid(agent_id): + return None + result = self._conn.execute( + text( + "SELECT id, user_id, image FROM agents " + "WHERE id = CAST(:id AS uuid)" + ), + {"id": agent_id}, + ) + row = result.fetchone() + return row_to_dict(row) if row is not None else None + def list_by_ids(self, agent_ids) -> list[dict]: """Fetch agents whose id is in ``agent_ids`` (team-shared listing path).""" ids = [str(a) for a in agent_ids if looks_like_uuid(str(a))] diff --git a/application/storage/local.py b/application/storage/local.py index a74f6181..a8ec5340 100644 --- a/application/storage/local.py +++ b/application/storage/local.py @@ -60,6 +60,14 @@ class LocalStorage(BaseStorage): return open(full_path, 'rb') + def get_file_size(self, path: str) -> int: + """Return the size of a local file without opening and buffering it.""" + full_path = self._get_full_path(path) + try: + return os.path.getsize(full_path) + except FileNotFoundError as exc: + raise FileNotFoundError(f"File not found: {full_path}") from exc + def delete_file(self, path: str) -> bool: """Delete a file from local storage.""" full_path = self._get_full_path(path) diff --git a/application/storage/s3.py b/application/storage/s3.py index 9c16fab2..74c11975 100644 --- a/application/storage/s3.py +++ b/application/storage/s3.py @@ -121,13 +121,35 @@ class S3Storage(BaseStorage): file_obj.seek(0) return file_obj - def generate_presigned_url(self, path: str, expires_in: int = 300) -> str: + def get_file_size(self, path: str) -> int: + """Return an S3 object's size using HEAD without downloading its body.""" + path = self._validate_path(path) + try: + metadata = self.s3.head_object(Bucket=self.bucket_name, Key=path) + except ClientError as exc: + error_code = str(exc.response.get("Error", {}).get("Code", "")) + # HEAD on a missing key returns 403 instead of 404 when credentials + # lack s3:ListBucket, so access-denied codes also mean "absent". + if error_code in {"404", "NoSuchKey", "NotFound", "403", "AccessDenied", "Forbidden"}: + raise FileNotFoundError(f"File not found: {path}") from exc + raise + return int(metadata["ContentLength"]) + + def generate_presigned_url( + self, + path: str, + expires_in: int = 300, + content_type: Optional[str] = None, + ) -> str: """Return a short-lived presigned GET URL for a private object (TTL <= 1h).""" path = self._validate_path(path) expires_in = min(expires_in, 3600) + params = {"Bucket": self.bucket_name, "Key": path} + if content_type: + params["ResponseContentType"] = content_type return self.s3.generate_presigned_url( "get_object", - Params={"Bucket": self.bucket_name, "Key": path}, + Params=params, ExpiresIn=expires_in, ) diff --git a/application/upload_limits.py b/application/upload_limits.py new file mode 100644 index 00000000..98736c29 --- /dev/null +++ b/application/upload_limits.py @@ -0,0 +1,138 @@ +"""Bound user-controlled upload streams before storage or parsing.""" + +from __future__ import annotations + +import io +import os +from contextlib import suppress +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) + with suppress(AttributeError, OSError): + stream.seek(0) + + 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) + with suppress(AttributeError, OSError): + stream.seek(0) + 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) + with suppress(AttributeError, OSError): + stream.seek(0) + + 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: + # Cleanup is best effort and must not mask the read result or exception. + with suppress(ValueError, OSError): + text_reader.detach() + with suppress(ValueError, OSError): + buffered_reader.detach() diff --git a/application/utils.py b/application/utils.py index 60032893..16dad7b8 100644 --- a/application/utils.py +++ b/application/utils.py @@ -1,12 +1,13 @@ import base64 import hashlib +import hmac import io import logging import os import re import uuid +from pathlib import PurePosixPath from typing import List -from urllib.parse import urlparse import tiktoken from flask import jsonify, make_response @@ -237,31 +238,123 @@ def validate_function_name(function_name): return True -def generate_image_url(image_path): +# Extension -> (accepted Pillow formats, MIME type). JPEG entries also accept +# MPO because Pillow reports multi-picture JPEGs (e.g. iPhone portrait photos, +# Samsung motion photos) as MPO. +AGENT_IMAGE_FORMATS = { + ".gif": (("GIF",), "image/gif"), + ".jpeg": (("JPEG", "MPO"), "image/jpeg"), + ".jpg": (("JPEG", "MPO"), "image/jpeg"), + ".png": (("PNG",), "image/png"), + ".webp": (("WEBP",), "image/webp"), +} + + +def is_external_image_url(image_path: object) -> bool: + """Return whether an image value is an externally hosted HTTP(S) URL.""" + return isinstance(image_path, str) and image_path.startswith(("http://", "https://")) + + +def safe_user_storage_component(user_id: object) -> str: + """Return a deterministic, traversal-safe directory component for a user ID.""" + raw_user_id = str(user_id or "") + component = secure_filename(raw_user_id)[:80] or "user" + digest = hashlib.sha256(raw_user_id.encode("utf-8")).hexdigest()[:16] + return f"{component}-{digest}" + + +def get_agent_image_content_type(image_path: object) -> str | None: + """Return an allow-listed raster MIME type for an internal image path.""" + if not isinstance(image_path, str): + return None + policy = AGENT_IMAGE_FORMATS.get(PurePosixPath(image_path).suffix.lower()) + return policy[1] if policy else None + + +def is_safe_agent_image_path(image_path: object, user_id: object) -> bool: + """Validate that a path is an agent avatar under its owner's upload directory.""" + if not isinstance(image_path, str) or not image_path or not user_id: + return False + if is_external_image_url(image_path) or "\\" in image_path or "\x00" in image_path: + return False + + candidate = PurePosixPath(image_path) + upload_root = PurePosixPath(str(settings.UPLOAD_FOLDER).rstrip("/")) + # An absolute UPLOAD_FOLDER yields absolute stored paths, so the two must + # agree; mismatched anchors mean the path did not come from this root. The + # exact parent match below is what actually contains the path. + if candidate.is_absolute() != upload_root.is_absolute() or ".." in candidate.parts: + return False + if get_agent_image_content_type(image_path) is None: + return False + + owner_components = {safe_user_storage_component(user_id)} + raw_user_id = str(user_id) + raw_owner = PurePosixPath(raw_user_id) + if ( + len(raw_owner.parts) == 1 + and raw_user_id not in {"", ".", ".."} + and "\\" not in raw_user_id + and "\x00" not in raw_user_id + ): + # Compatibility for avatars written before user IDs were sanitized. + owner_components.add(raw_user_id) + + for owner_component in owner_components: + expected_parent = upload_root / owner_component / "attachments" + if candidate.parent == expected_parent and candidate.name: + return True + return False + + +def generate_agent_image_capability( + agent_id: object, image_path: object, user_id: object +) -> str: + """Create an HMAC capability for one agent's current internal image.""" + secret = getattr(settings, "JWT_SECRET_KEY", "") + if not isinstance(secret, str) or not secret: + return "" + try: + canonical_agent_id = str(uuid.UUID(str(agent_id))) + except (TypeError, ValueError, AttributeError): + return "" + if not isinstance(image_path, str) or not user_id: + return "" + payload = ( + f"docsgpt-agent-image-v1\0{canonical_agent_id}\0{user_id}\0{image_path}" + ).encode("utf-8") + return hmac.new(secret.encode("utf-8"), payload, hashlib.sha256).hexdigest() + + +def verify_agent_image_capability( + capability: object, agent_id: object, image_path: object, user_id: object +) -> bool: + """Verify an agent image capability without timing-leaky string comparison.""" + if not isinstance(capability, str) or not re.fullmatch(r"[0-9a-f]{64}", capability): + return False + expected = generate_agent_image_capability(agent_id, image_path, user_id) + return bool(expected) and hmac.compare_digest(capability, expected) + + +def generate_image_url(image_path, agent_id=None, user_id=None): + """Return an external URL or an opaque capability URL for an agent image. + + Internal storage paths are never included in the returned URL. Invalid or + unscoped paths fail closed so a poisoned database row cannot become a file + read capability. + """ if isinstance(image_path, str) and ( image_path.startswith("http://") or image_path.startswith("https://") ): return image_path - strategy = getattr(settings, "URL_STRATEGY", "backend") - if strategy == "s3": - bucket_name = settings.S3_BUCKET_NAME - endpoint_url = settings.S3_ENDPOINT_URL - if endpoint_url: - # S3-compatible service (MinIO, R2, B2, Spaces, ...). - base = endpoint_url.rstrip("/") - if settings.S3_PATH_STYLE: - return f"{base}/{bucket_name}/{image_path}" - parsed = urlparse(base) - return f"{parsed.scheme}://{bucket_name}.{parsed.netloc}/{image_path}" - region_name = ( - settings.S3_REGION - or getattr(settings, "SAGEMAKER_REGION", None) - or "eu-central-1" - ) - return f"https://{bucket_name}.s3.{region_name}.amazonaws.com/{image_path}" - else: - base_url = getattr(settings, "API_URL", "http://localhost:7091") - return f"{base_url}/api/images/{image_path}" + if not is_safe_agent_image_path(image_path, user_id): + return "" + capability = generate_agent_image_capability(agent_id, image_path, user_id) + if not capability: + return "" + canonical_agent_id = str(uuid.UUID(str(agent_id))) + base_url = getattr(settings, "API_URL", "http://localhost:7091").rstrip("/") + return f"{base_url}/api/images/{canonical_agent_id}/{capability}" def calculate_compression_threshold( diff --git a/application/worker.py b/application/worker.py index 19b298c5..c152c25b 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. @@ -404,39 +375,51 @@ def extract_zip_recursive(zip_path, extract_to, current_depth=0, max_depth=5): extract_to (str): Destination path for extracted files. current_depth (int): Current depth of recursion. max_depth (int): Maximum allowed depth of recursion to prevent infinite loops. + + Raises: + ZipExtractionError: If the archive fails safety validation, so ingestion + fails loudly instead of indexing an empty directory. """ if current_depth > max_depth: logging.warning(f"Reached maximum recursion depth of {max_depth}") 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) except OSError: pass - return + raise except Exception as e: - logging.error(f"Error extracting zip file {zip_path}: {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) + logging.error( + "Error extracting zip file %s: %s", + safe_zip_error_message(zip_path), + safe_zip_error_message(e), + exc_info=True, + ) + raise def download_file(url, params, dest_path): @@ -623,7 +606,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/deployment/k8s/deployments/docsgpt-deploy.yaml b/deployment/k8s/deployments/docsgpt-deploy.yaml index 345b1449..fc7631fd 100644 --- a/deployment/k8s/deployments/docsgpt-deploy.yaml +++ b/deployment/k8s/deployments/docsgpt-deploy.yaml @@ -99,6 +99,8 @@ spec: - secretRef: name: docsgpt-secrets env: + - name: DEPLOYMENT_TYPE + value: "cloud" - name: API_URL value: "http://" - name: POSTGRES_URI diff --git a/deployment/k8s/docsgpt-secrets.yaml b/deployment/k8s/docsgpt-secrets.yaml index 6d558a0f..4d125e6c 100644 --- a/deployment/k8s/docsgpt-secrets.yaml +++ b/deployment/k8s/docsgpt-secrets.yaml @@ -12,6 +12,11 @@ type: Opaque # manually add a MONGO_URI key below (base64-encoded) and apply the # opt-in manifests under deployment/k8s/optional-mongo/. Example: # MONGO_URI: +# - JWT_SECRET_KEY below is a known dev-grade default so the stock manifests +# boot out of the box. Replace it with the base64 encoding of a strong random +# value before exposing the API. Every API and worker replica must receive +# the same key. +# Example generation: openssl rand -hex 32 | base64 data: LLM_PROVIDER: ZG9jc2dwdA== INTERNAL_KEY: aW50ZXJuYWw= @@ -23,3 +28,5 @@ data: POSTGRES_URI: cG9zdGdyZXNxbDovL2RvY3NncHQ6ZG9jc2dwdEBwb3N0Z3Jlczo1NDMyL2RvY3NncHQ= postgres-user: ZG9jc2dwdA== postgres-password: ZG9jc2dwdA== + # docsgpt-dev-jwt-secret-change-me — CHANGE before exposing the API + JWT_SECRET_KEY: ZG9jc2dwdC1kZXYtand0LXNlY3JldC1jaGFuZ2UtbWU= diff --git a/docs/content/Deploying/DocsGPT-Settings.mdx b/docs/content/Deploying/DocsGPT-Settings.mdx index 42234d92..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: @@ -217,6 +241,17 @@ STT_ENABLE_DIARIZATION=false `faster_whisper` is an optional backend dependency. Install it in the Python environment used by the DocsGPT API and worker before selecting this provider. +## Agent Image Settings + +Agent avatars uploaded to DocsGPT must be valid PNG, JPEG, GIF, or WebP files. +The decoded image and its file extension must agree; SVG and other formats are +rejected. These limits apply before the image is written to local or S3 storage. + +| Setting | Default | Description | +| --- | --- | --- | +| `AGENT_IMAGE_MAX_BYTES` | `5000000` | Maximum encoded avatar file size in bytes. | +| `AGENT_IMAGE_MAX_PIXELS` | `16777216` | Maximum decoded width × height, limiting decompression-bomb images. | + ## Authentication Settings DocsGPT includes a JWT (JSON Web Token) based authentication feature for managing sessions or securing local deployments while allowing access. @@ -237,27 +272,28 @@ The `AUTH_TYPE` setting in your `.env` file or `settings.py` determines the auth Add the following to your `.env` file (or set in `settings.py`): ```env +# Shared signing key (required in production for every authentication mode) +JWT_SECRET_KEY= + # No authentication (default) AUTH_TYPE=None # OR: Simple JWT (shared token) AUTH_TYPE=simple_jwt -JWT_SECRET_KEY=your_secret_key_here # OR: Session JWT (per-user/session tokens) AUTH_TYPE=session_jwt -JWT_SECRET_KEY=your_secret_key_here # OR: SSO via an OpenID Connect provider (Authentik, Keycloak, Okta, ...) AUTH_TYPE=oidc OIDC_ISSUER=https://auth.example.com/application/o/docsgpt/ OIDC_CLIENT_ID=your_client_id OIDC_FRONTEND_URL=https://docsgpt.example.com -JWT_SECRET_KEY=your_secret_key_here ``` -- If `AUTH_TYPE` is set to `simple_jwt` or `session_jwt`, a `JWT_SECRET_KEY` is required. -- If `JWT_SECRET_KEY` is not set, DocsGPT will generate one and store it in `.jwt_secret_key` in the project root. +- `JWT_SECRET_KEY` signs authentication tokens where applicable and opaque agent-avatar capabilities in every authentication mode, including no-auth mode. +- Cloud and production deployments must set a strong `JWT_SECRET_KEY`, shared unchanged by every API and worker replica. Startup fails rather than creating replica-local keys when it is missing. +- Local development may omit it. DocsGPT atomically generates an owner-readable `.jwt_secret_key` in the project root and reuses it on later starts. #### How Each Method Works @@ -283,7 +319,7 @@ JWT_SECRET_KEY=your_secret_key_here - Always keep your `JWT_SECRET_KEY` secure and private. - If you set it manually, use a strong, random string. -- If not set, DocsGPT will generate a secure key and persist it in `.jwt_secret_key`. +- Keep the value stable. Rotating it invalidates active JWTs and previously generated agent-avatar URLs. #### Checking Current Auth Type @@ -317,7 +353,7 @@ By default DocsGPT stores files locally. Set `STORAGE_TYPE=s3` to use Amazon S3 | `S3_REGION` | Region (use `auto` for Cloudflare R2) | — | | `S3_ENDPOINT_URL` | Custom endpoint for S3-compatible services; leave unset for AWS S3 | — | | `S3_PATH_STYLE` | Use path-style addressing (required by most non-AWS services) | `false` | -| `URL_STRATEGY` | `backend` (proxy through API) or `s3` (direct object URLs) | `backend` | +| `URL_STRATEGY` | Artifact-download delivery: `backend` proxies bytes through the API; `s3` returns short-lived presigned object URLs. Agent avatars always use capability URLs and, with S3 storage, redirect to a size-checked presigned URL. | `backend` | ### AWS S3 diff --git a/docs/content/Deploying/Kubernetes-Deploying.mdx b/docs/content/Deploying/Kubernetes-Deploying.mdx index e671be1a..e0c73633 100644 --- a/docs/content/Deploying/Kubernetes-Deploying.mdx +++ b/docs/content/Deploying/Kubernetes-Deploying.mdx @@ -33,9 +33,11 @@ The `deployment/k8s` folder contains the necessary deployment and service config cd docsgpt/deployment/k8s ``` -2. **Configure Secrets (optional)** +2. **Configure Secrets** - Ensure that you have all the necessary secrets in `docsgpt-secrets.yaml`. Update it with your secrets before applying if you want. By default we will use qdrant as a vectorstore and public docsgpt llm as llm for inference. + Add a strong, base64-encoded `JWT_SECRET_KEY` to `docsgpt-secrets.yaml` before applying it. The API and worker deployments share this value so authentication tokens and agent-avatar capability URLs remain valid across pod restarts and replicas. You can generate a value with `openssl rand -hex 32 | base64`. + + Configure any additional provider secrets you need. By default, the manifests use Qdrant as the vector store and the public DocsGPT LLM for inference. Alternatively, you can use [Neon](https://get.neon.com/docsgpt) as an easy way to set up your vector store with pgvector, which is highly recommended for quick deployments. @@ -106,4 +108,4 @@ If you encounter any issues, you can check the logs of the pods for more details kubectl logs ``` -Replace `` with the actual name of your DocsGPT pod. \ No newline at end of file +Replace `` with the actual name of your DocsGPT pod. diff --git a/docs/content/Deploying/OIDC-SSO.mdx b/docs/content/Deploying/OIDC-SSO.mdx index 377ea929..8851994e 100644 --- a/docs/content/Deploying/OIDC-SSO.mdx +++ b/docs/content/Deploying/OIDC-SSO.mdx @@ -45,7 +45,7 @@ Sessions last `OIDC_SESSION_LIFETIME_SECONDS` (8 hours by default) and renew wit | `OIDC_ALLOWED_GROUPS` | no | — | Comma-separated group allowlist. Unset, any authenticated IdP user may sign in — see [Restricting sign-in by group](#restricting-sign-in-by-group). | | `OIDC_ADMIN_GROUPS` | no | — | Comma-separated groups whose members are granted the global `admin` role. Re-checked at every login and renewal — see [Granting admin via groups](#granting-admin-via-groups). | | `OIDC_GROUPS_CLAIM` | no | `groups` | ID-token/userinfo claim carrying the user's group membership. | -| `JWT_SECRET_KEY` | recommended | auto-generated | Signs DocsGPT session tokens. Set it explicitly in production — required when running multiple API replicas. | +| `JWT_SECRET_KEY` | required in production | auto-generated for local development | Signs DocsGPT session tokens and agent-avatar capabilities. Every API and worker replica must use the same value. | `SCIM_ENABLED` and `SCIM_TOKEN` are listed in the [SCIM section](#scim-user-provisioning). diff --git a/frontend/src/components/DocumentArtifactView.tsx b/frontend/src/components/DocumentArtifactView.tsx index 61db9692..1050c3e2 100644 --- a/frontend/src/components/DocumentArtifactView.tsx +++ b/frontend/src/components/DocumentArtifactView.tsx @@ -1,4 +1,3 @@ -import mermaid from 'mermaid'; import { Download, FileText, History, RotateCcw } from 'lucide-react'; import { useEffect, useMemo, useRef, useState, type ReactNode } from 'react'; import { useSelector } from 'react-redux'; @@ -7,6 +6,7 @@ import userService from '../api/services/userService'; import { useDarkTheme } from '../hooks'; import { selectToken } from '../preferences/preferenceSlice'; import MarkdownPreview from './MarkdownPreview'; +import { renderMermaidDiagram } from './mermaidSecurity'; import Spinner from './Spinner'; import { buildPreviewDocument, @@ -67,13 +67,12 @@ function FramePreview({ let cancelled = false; setRendered(null); setRenderError(false); - mermaid.initialize({ - startOnLoad: false, + renderMermaidDiagram({ + id: `artifact-mermaid-${Date.now()}`, + code: source, + isDarkTheme, securityLevel: 'sandbox', - theme: isDarkTheme ? 'dark' : 'default', - }); - mermaid - .render(`artifact-mermaid-${Date.now()}`, source) + }) .then(({ svg }) => { if (!cancelled) setRendered(svg); }) diff --git a/frontend/src/components/MermaidRenderer.test.tsx b/frontend/src/components/MermaidRenderer.test.tsx new file mode 100644 index 00000000..409235c0 --- /dev/null +++ b/frontend/src/components/MermaidRenderer.test.tsx @@ -0,0 +1,54 @@ +import { act } from 'react'; +import { createRoot, type Root } from 'react-dom/client'; + +Object.assign(globalThis, { IS_REACT_ACT_ENVIRONMENT: true }); + +const renderMermaidDiagramMock = vi.hoisted(() => vi.fn()); + +vi.mock('react-i18next', () => ({ + useTranslation: () => ({ t: (key: string) => key }), +})); +vi.mock('react-redux', () => ({ useSelector: () => 'idle' })); +vi.mock('../hooks', () => ({ useDarkTheme: () => [false] })); +vi.mock('./CopyButton', () => ({ default: () => null })); +vi.mock('./mermaidSecurity', () => ({ + renderMermaidDiagram: renderMermaidDiagramMock, +})); + +import MermaidRenderer from './MermaidRenderer'; + +describe('MermaidRenderer', () => { + let container: HTMLDivElement; + let root: Root; + + beforeEach(() => { + container = document.createElement('div'); + document.body.appendChild(container); + root = createRoot(container); + renderMermaidDiagramMock.mockReset(); + }); + + afterEach(async () => { + await act(async () => root.unmount()); + container.remove(); + }); + + it('keeps its visible host when Mermaid removes its temporary render ID', async () => { + renderMermaidDiagramMock.mockImplementation( + async ({ id }: { id: string }) => { + // Mermaid removes an existing node with the supplied render ID. + document.getElementById(id)?.remove(); + return { svg: '' }; + }, + ); + + await act(async () => { + root.render( + B'} isLoading={false} />, + ); + }); + + expect(container.querySelector('pre.mermaid')).not.toBeNull(); + expect(container.querySelector('svg[data-rendered="true"]')).not.toBeNull(); + }); +}); diff --git a/frontend/src/components/MermaidRenderer.tsx b/frontend/src/components/MermaidRenderer.tsx index 6a9a1ef8..ff86d882 100644 --- a/frontend/src/components/MermaidRenderer.tsx +++ b/frontend/src/components/MermaidRenderer.tsx @@ -1,4 +1,3 @@ -import mermaid from 'mermaid'; import React, { useEffect, useRef, useState } from 'react'; import { useTranslation } from 'react-i18next'; import { useSelector } from 'react-redux'; @@ -11,6 +10,7 @@ import { import { selectStatus } from '../conversation/conversationSlice'; import { useDarkTheme } from '../hooks'; import CopyButton from './CopyButton'; +import { renderMermaidDiagram } from './mermaidSecurity'; import { Button } from './ui/button'; import { MermaidRendererProps } from './types'; @@ -23,12 +23,14 @@ const MermaidRenderer: React.FC = ({ const diagramId = useRef( `mermaid-${Date.now()}-${Math.random().toString(36).substring(2)}`, ); + const renderSequence = useRef(0); const status = useSelector(selectStatus); const [error, setError] = useState(null); const [showCode, setShowCode] = useState(false); const [showDownloadMenu, setShowDownloadMenu] = useState(false); const downloadMenuRef = useRef(null); const containerRef = useRef(null); + const diagramRef = useRef(null); const [hoverPosition, setHoverPosition] = useState<{ x: number; y: number; @@ -83,26 +85,33 @@ const MermaidRenderer: React.FC = ({ return `${hoverPosition.x * 100}% ${hoverPosition.y * 100}%`; }; - useEffect(() => { - const renderDiagram = async () => { - mermaid.initialize({ - startOnLoad: true, - theme: isDarkTheme ? 'dark' : 'default', - securityLevel: 'loose', - suppressErrorRendering: true, - }); + const isCurrentlyLoading = + isLoading !== undefined ? isLoading : status === 'loading'; - const isCurrentlyLoading = - isLoading !== undefined ? isLoading : status === 'loading'; + useEffect(() => { + let cancelled = false; + const renderDiagram = async () => { if (!isCurrentlyLoading && code) { try { - const element = document.getElementById(diagramId.current); + setError(null); + const { svg, bindFunctions } = await renderMermaidDiagram({ + // Mermaid owns and removes this temporary ID while rendering. It + // must never match the visible host's ID. + id: `${diagramId.current}-render-${++renderSequence.current}`, + code, + isDarkTheme, + securityLevel: 'strict', + }); + if (cancelled) return; + + const element = diagramRef.current; if (element) { - element.removeAttribute('data-processed'); - await mermaid.parse(code); //syntax check - mermaid.contentLoaded(); + // Mermaid strict mode sanitizes this generated SVG before it is returned. + element.innerHTML = svg; + bindFunctions?.(element); } } catch (err) { + if (cancelled) return; console.error('Error rendering mermaid diagram:', err); setError( `Failed to render diagram: ${err instanceof Error ? err.message : String(err)}`, @@ -112,7 +121,10 @@ const MermaidRenderer: React.FC = ({ }; renderDiagram(); - }, [code, isDarkTheme, isLoading]); + return () => { + cancelled = true; + }; + }, [code, isDarkTheme, isCurrentlyLoading]); useEffect(() => { const handleClickOutside = (event: MouseEvent) => { @@ -258,8 +270,6 @@ const MermaidRenderer: React.FC = ({ { label: 'Download as MMD', action: downloadMmd }, ]; - const isCurrentlyLoading = - isLoading !== undefined ? isLoading : status === 'loading'; const showDiagramOptions = !isCurrentlyLoading && !error; const errorRender = !isCurrentlyLoading && error; @@ -395,6 +405,7 @@ const MermaidRenderer: React.FC = ({ )} 
 = ({
                 display: 'flex',
                 justifyContent: 'center',
               }}
-            >
-              {code}
-
+ /> {showCode && ( diff --git a/frontend/src/components/mermaidRenderQueue.test.ts b/frontend/src/components/mermaidRenderQueue.test.ts new file mode 100644 index 00000000..80b40fe6 --- /dev/null +++ b/frontend/src/components/mermaidRenderQueue.test.ts @@ -0,0 +1,95 @@ +const mermaidMock = vi.hoisted(() => ({ + initialize: vi.fn(), + render: vi.fn(), +})); + +vi.mock('mermaid', () => ({ default: mermaidMock })); + +import { renderMermaidDiagram } from './mermaidSecurity'; + +function deferred() { + let resolve!: (value: T) => void; + const promise = new Promise((done) => { + resolve = done; + }); + return { promise, resolve }; +} + +describe('renderMermaidDiagram', () => { + beforeEach(() => { + mermaidMock.initialize.mockReset(); + mermaidMock.render.mockReset(); + }); + + it('serializes initialization with rendering across security levels', async () => { + const firstRender = deferred<{ svg: string }>(); + mermaidMock.render + .mockImplementationOnce(() => firstRender.promise) + .mockResolvedValueOnce({ svg: '' }); + + const conversation = renderMermaidDiagram({ + id: 'conversation-diagram', + code: 'flowchart LR\nA --> B', + isDarkTheme: false, + securityLevel: 'strict', + }); + const artifact = renderMermaidDiagram({ + id: 'artifact-diagram', + code: 'flowchart LR\nB --> C', + isDarkTheme: true, + securityLevel: 'sandbox', + }); + + await vi.waitFor(() => expect(mermaidMock.render).toHaveBeenCalledTimes(1)); + expect(mermaidMock.initialize).toHaveBeenLastCalledWith( + expect.objectContaining({ securityLevel: 'strict' }), + ); + + firstRender.resolve({ svg: '' }); + await conversation; + await artifact; + + expect(mermaidMock.initialize.mock.calls.map(([config]) => config)).toEqual( + [ + expect.objectContaining({ securityLevel: 'strict' }), + expect.objectContaining({ securityLevel: 'sandbox' }), + ], + ); + expect(mermaidMock.render.mock.calls).toEqual([ + ['conversation-diagram', 'flowchart LR\nA --> B'], + ['artifact-diagram', 'flowchart LR\nB --> C'], + ]); + }); + + it('renders from the code argument again when only the theme changes', async () => { + mermaidMock.render.mockResolvedValue({ svg: '' }); + const code = 'flowchart LR\nA --> B'; + + await renderMermaidDiagram({ + id: 'theme-diagram', + code, + isDarkTheme: false, + securityLevel: 'strict', + }); + await renderMermaidDiagram({ + id: 'theme-diagram', + code, + isDarkTheme: true, + securityLevel: 'strict', + }); + + expect(mermaidMock.render).toHaveBeenNthCalledWith( + 1, + 'theme-diagram', + code, + ); + expect(mermaidMock.render).toHaveBeenNthCalledWith( + 2, + 'theme-diagram', + code, + ); + expect( + mermaidMock.initialize.mock.calls.map(([config]) => config.theme), + ).toEqual(['default', 'dark']); + }); +}); diff --git a/frontend/src/components/mermaidSecurity.test.ts b/frontend/src/components/mermaidSecurity.test.ts new file mode 100644 index 00000000..312d4a71 --- /dev/null +++ b/frontend/src/components/mermaidSecurity.test.ts @@ -0,0 +1,35 @@ +import mermaid from 'mermaid'; + +import { initializeSecureMermaid } from './mermaidSecurity'; + +describe('initializeSecureMermaid', () => { + beforeEach(() => { + document.body.replaceChildren(); + }); + + it('ignores security-level downgrade attempts in init directives', async () => { + initializeSecureMermaid(false); + const storedPayload = [ + '%%{init: {"securityLevel":"loose"}}%%', + 'flowchart LR', + ' A[Click me]', + ' click A "javascript:alert(document.domain)"', + ].join('\n'); + + const parsed = await mermaid.parse(storedPayload); + + expect(parsed).toMatchObject({ diagramType: 'flowchart-v2' }); + expect(mermaid.mermaidAPI.getConfig().securityLevel).toBe('strict'); + }); + + it('continues to accept ordinary diagrams', async () => { + initializeSecureMermaid(true); + + const parsed = await mermaid.parse( + 'flowchart LR\n A[Start] --> B[Finish]', + ); + + expect(parsed).toMatchObject({ diagramType: 'flowchart-v2' }); + expect(mermaid.mermaidAPI.getConfig().theme).toBe('dark'); + }); +}); diff --git a/frontend/src/components/mermaidSecurity.ts b/frontend/src/components/mermaidSecurity.ts new file mode 100644 index 00000000..57f098b2 --- /dev/null +++ b/frontend/src/components/mermaidSecurity.ts @@ -0,0 +1,46 @@ +import mermaid, { type RenderResult } from 'mermaid'; + +export type MermaidSecurityLevel = 'strict' | 'sandbox'; + +type MermaidRenderOptions = { + id: string; + code: string; + isDarkTheme: boolean; + securityLevel: MermaidSecurityLevel; +}; + +// Mermaid's own render queue does not snapshot its process-global config. +// Keep initialization and its matching render in one application-level task. +let renderQueue: Promise = Promise.resolve(); + +/** Configure Mermaid immediately before one queued render. */ +export function initializeSecureMermaid( + isDarkTheme: boolean, + securityLevel: MermaidSecurityLevel = 'strict', +): void { + mermaid.initialize({ + startOnLoad: false, + theme: isDarkTheme ? 'dark' : 'default', + securityLevel, + suppressErrorRendering: true, + }); +} + +/** Render untrusted Mermaid source with configuration isolated from other views. */ +export function renderMermaidDiagram({ + id, + code, + isDarkTheme, + securityLevel, +}: MermaidRenderOptions): Promise { + const render = () => { + initializeSecureMermaid(isDarkTheme, securityLevel); + return mermaid.render(id, code); + }; + const result = renderQueue.then(render, render); + renderQueue = result.then( + () => undefined, + () => undefined, + ); + return result; +} diff --git a/tests/api/user/agents/test_routes_happy.py b/tests/api/user/agents/test_routes_happy.py index 24483dd2..12e6197e 100644 --- a/tests/api/user/agents/test_routes_happy.py +++ b/tests/api/user/agents/test_routes_happy.py @@ -465,6 +465,28 @@ class TestUpdateAgent: got = AgentsRepository(pg_conn).get(str(agent["id"]), user) assert got["name"] == "new name" + def test_ignores_client_supplied_image_path(self, app, pg_conn): + from application.api.user.agents.routes import UpdateAgent + from application.storage.db.repositories.agents import AgentsRepository + + user = "u-upd-image-path" + agent = _seed_agent( + pg_conn, user=user, status="draft", with_source=False, + ) + + with _patch_db(pg_conn), app.test_request_context( + f"/api/update_agent/{agent['id']}", method="PUT", + json={"description": "safe update", "image": ".env"}, + ): + from flask import request + request.decoded_token = {"sub": user} + response = UpdateAgent().put(str(agent["id"])) + + assert response.status_code == 200 + got = AgentsRepository(pg_conn).get(str(agent["id"]), user) + assert got["description"] == "safe update" + assert got["image"] is None + def test_invalid_status_returns_400(self, app, pg_conn): from application.api.user.agents.routes import UpdateAgent @@ -1221,6 +1243,120 @@ class TestAdoptAgentMore: mine = repo.list_for_user("u-adopter") assert any(a["name"] == "Template X" for a in mine) + def test_adopt_copies_image_into_adopter_storage(self, app, pg_conn): + import io + from unittest.mock import MagicMock + + from application.api.user.agents.routes import AdoptAgent + from application.core.settings import settings + from application.storage.db.repositories.agents import AgentsRepository + from application.utils import ( + is_safe_agent_image_path, + safe_user_storage_component, + ) + + repo = AgentsRepository(pg_conn) + template = repo.create( + "__system__", + "Template Img", + "template", + image="inputs/__system__/attachments/abc_logo.png", + ) + + storage = MagicMock() + storage.get_file.return_value = io.BytesIO(b"png-bytes") + storage.save_file.side_effect = lambda data, path, **kw: {"path": path} + + with _patch_db(pg_conn), patch( + "application.api.user.agents.routes.storage", storage + ), app.test_request_context( + f"/api/adopt_agent?id={template['id']}", method="POST" + ): + from flask import request + request.decoded_token = {"sub": "u-adopter"} + response = AdoptAgent().post() + + assert response.status_code == 200 + adopted = next( + a for a in repo.list_for_user("u-adopter") if a["name"] == "Template Img" + ) + owner = safe_user_storage_component("u-adopter") + expected_prefix = f"{settings.UPLOAD_FOLDER}/{owner}/attachments/" + assert adopted["image"].startswith(expected_prefix) + assert adopted["image"].endswith("_logo.png") + assert is_safe_agent_image_path(adopted["image"], "u-adopter") + storage.get_file.assert_called_once_with( + "inputs/__system__/attachments/abc_logo.png" + ) + + def test_adopt_drops_image_when_copy_fails(self, app, pg_conn): + from unittest.mock import MagicMock + + from application.api.user.agents.routes import AdoptAgent + from application.storage.db.repositories.agents import AgentsRepository + + repo = AgentsRepository(pg_conn) + template = repo.create( + "__system__", + "Template Broken Img", + "template", + image="inputs/__system__/attachments/gone_logo.png", + ) + + storage = MagicMock() + storage.get_file.side_effect = FileNotFoundError("missing") + + with _patch_db(pg_conn), patch( + "application.api.user.agents.routes.storage", storage + ), app.test_request_context( + f"/api/adopt_agent?id={template['id']}", method="POST" + ): + from flask import request + request.decoded_token = {"sub": "u-adopter2"} + response = AdoptAgent().post() + + assert response.status_code == 200 + adopted = next( + a + for a in repo.list_for_user("u-adopter2") + if a["name"] == "Template Broken Img" + ) + assert not adopted.get("image") + + def test_adopt_keeps_external_image_url(self, app, pg_conn): + from unittest.mock import MagicMock + + from application.api.user.agents.routes import AdoptAgent + from application.storage.db.repositories.agents import AgentsRepository + + repo = AgentsRepository(pg_conn) + template = repo.create( + "__system__", + "Template Ext Img", + "template", + image="https://cdn.example.com/logo.png", + ) + + storage = MagicMock() + + with _patch_db(pg_conn), patch( + "application.api.user.agents.routes.storage", storage + ), app.test_request_context( + f"/api/adopt_agent?id={template['id']}", method="POST" + ): + from flask import request + request.decoded_token = {"sub": "u-adopter3"} + response = AdoptAgent().post() + + assert response.status_code == 200 + adopted = next( + a + for a in repo.list_for_user("u-adopter3") + if a["name"] == "Template Ext Img" + ) + assert adopted["image"] == "https://cdn.example.com/logo.png" + storage.get_file.assert_not_called() + def test_adopt_template_missing_returns_404(self, app, pg_conn): from application.api.user.agents.routes import AdoptAgent diff --git a/tests/api/user/agents/test_routes_helpers.py b/tests/api/user/agents/test_routes_helpers.py index a7572a07..ecbf2bb1 100644 --- a/tests/api/user/agents/test_routes_helpers.py +++ b/tests/api/user/agents/test_routes_helpers.py @@ -258,3 +258,11 @@ class TestBuildCreateKwargs: {}, image_url="/upload/img.png", agent_type="classic", ) assert out.get("image") == "/upload/img.png" + + def test_client_image_path_is_never_persisted(self): + from application.api.user.agents.routes import _build_create_kwargs + + out = _build_create_kwargs( + {"image": ".env"}, image_url="", agent_type="classic", + ) + assert "image" not in out diff --git a/tests/api/user/attachments/test_routes.py b/tests/api/user/attachments/test_routes.py index c010918e..d28989af 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 @@ -1520,23 +1543,54 @@ class TestLiveSpeechToTextAdditional: class TestServeImage: """Tests for ServeImage endpoint.""" + agent_id = "00000000-0000-0000-0000-000000000001" + image_path = "inputs/user123/attachments/avatar.png" + + @staticmethod + def _capability(agent_id, image_path, user_id="user123"): + from application.utils import generate_agent_image_capability + + return generate_agent_image_capability(agent_id, image_path, user_id) + def test_serve_image_success(self, flask_app): from application.api.user.attachments.routes import ServeImage app = Flask(__name__) mock_storage = MagicMock() mock_file_obj = io.BytesIO(b"\x89PNG\r\n") + mock_storage.get_file_size.return_value = len(mock_file_obj.getvalue()) mock_storage.get_file.return_value = mock_file_obj - with patch("application.api.user.base.storage", mock_storage): + agent = { + "id": self.agent_id, + "user_id": "user123", + "image": self.image_path, + } + + with patch( + "application.utils.settings.JWT_SECRET_KEY", "test-image-secret" + ), patch( + "application.api.user.attachments.routes.AgentsRepository.find_image_record", + return_value=agent, + ), patch( + "application.api.user.attachments.routes.db_readonly", _fake_readonly + ), patch( + "application.api.user.attachments.routes.settings.STORAGE_TYPE", "local" + ), patch("application.api.user.base.storage", mock_storage): + capability = self._capability(self.agent_id, self.image_path) with app.test_request_context( - "/api/images/test/image.png", + f"/api/images/{self.agent_id}/{capability}", method="GET", ): resource = ServeImage() - response = resource.get("test/image.png") + response = resource.get(self.agent_id, capability) assert _get_response_status(response) == 200 assert response.headers.get("Content-Type") == "image/png" + assert response.headers.get("X-Content-Type-Options") == "nosniff" + assert response.is_streamed + assert response.content_length is None + assert b"".join(response.response) == b"\x89PNG\r\n" + mock_storage.get_file.assert_called_once_with(self.image_path) def test_serve_image_jpg_content_type(self, flask_app): from application.api.user.attachments.routes import ServeImage @@ -1544,32 +1598,184 @@ class TestServeImage: app = Flask(__name__) mock_storage = MagicMock() mock_file_obj = io.BytesIO(b"\xff\xd8\xff\xe0") + mock_storage.get_file_size.return_value = len(mock_file_obj.getvalue()) mock_storage.get_file.return_value = mock_file_obj - with patch("application.api.user.base.storage", mock_storage): + image_path = "inputs/user123/attachments/photo.jpg" + agent = { + "id": self.agent_id, + "user_id": "user123", + "image": image_path, + } + + with patch( + "application.utils.settings.JWT_SECRET_KEY", "test-image-secret" + ), patch( + "application.api.user.attachments.routes.AgentsRepository.find_image_record", + return_value=agent, + ), patch( + "application.api.user.attachments.routes.db_readonly", _fake_readonly + ), patch( + "application.api.user.attachments.routes.settings.STORAGE_TYPE", "local" + ), patch("application.api.user.base.storage", mock_storage): + capability = self._capability(self.agent_id, image_path) with app.test_request_context( - "/api/images/test/photo.jpg", + f"/api/images/{self.agent_id}/{capability}", method="GET", ): resource = ServeImage() - response = resource.get("test/photo.jpg") + response = resource.get(self.agent_id, capability) assert _get_response_status(response) == 200 assert response.headers.get("Content-Type") == "image/jpeg" - def test_serve_image_not_found(self, flask_app): + def test_s3_avatar_redirects_without_downloading_object(self, flask_app): from application.api.user.attachments.routes import ServeImage app = Flask(__name__) mock_storage = MagicMock() - mock_storage.get_file.side_effect = FileNotFoundError("not found") + mock_storage.get_file_size.return_value = 1024 + mock_storage.generate_presigned_url.return_value = ( + "https://bucket.example/avatar.png?signature=short-lived" + ) + agent = { + "id": self.agent_id, + "user_id": "user123", + "image": self.image_path, + } - with patch("application.api.user.base.storage", mock_storage): + with patch( + "application.utils.settings.JWT_SECRET_KEY", "test-image-secret" + ), patch( + "application.api.user.attachments.routes.AgentsRepository.find_image_record", + return_value=agent, + ), patch( + "application.api.user.attachments.routes.db_readonly", _fake_readonly + ), patch( + "application.api.user.attachments.routes.settings.STORAGE_TYPE", "s3" + ), patch("application.api.user.base.storage", mock_storage): + capability = self._capability(self.agent_id, self.image_path) with app.test_request_context( - "/api/images/missing/image.png", + f"/api/images/{self.agent_id}/{capability}", method="GET" + ): + response = ServeImage().get(self.agent_id, capability) + + assert response.status_code == 302 + assert response.headers["Location"].startswith("https://bucket.example/") + assert response.headers["Cache-Control"] == "private, max-age=240" + mock_storage.generate_presigned_url.assert_called_once_with( + self.image_path, expires_in=300, content_type="image/png" + ) + mock_storage.get_file.assert_not_called() + + def test_legacy_oversized_avatar_is_rejected_before_read(self, flask_app): + from application.api.user.attachments.routes import ServeImage + + app = Flask(__name__) + mock_storage = MagicMock() + mock_storage.get_file_size.return_value = 5001 + agent = { + "id": self.agent_id, + "user_id": "user123", + "image": self.image_path, + } + + with patch( + "application.utils.settings.JWT_SECRET_KEY", "test-image-secret" + ), patch( + "application.api.user.attachments.routes.AgentsRepository.find_image_record", + return_value=agent, + ), patch( + "application.api.user.attachments.routes.db_readonly", _fake_readonly + ), patch( + "application.api.user.attachments.routes.settings.AGENT_IMAGE_MAX_BYTES", 5000 + ), patch("application.api.user.base.storage", mock_storage): + capability = self._capability(self.agent_id, self.image_path) + with app.test_request_context( + f"/api/images/{self.agent_id}/{capability}", method="GET" + ): + response = ServeImage().get(self.agent_id, capability) + + assert response.status_code == 404 + mock_storage.get_file.assert_not_called() + mock_storage.generate_presigned_url.assert_not_called() + + def test_unknown_capability_never_reads_storage(self, flask_app): + from application.api.user.attachments.routes import ServeImage + + app = Flask(__name__) + mock_storage = MagicMock() + + agent = { + "id": self.agent_id, + "user_id": "user123", + "image": self.image_path, + } + + with patch( + "application.utils.settings.JWT_SECRET_KEY", "test-image-secret" + ), patch( + "application.api.user.attachments.routes.AgentsRepository.find_image_record", + return_value=agent, + ), patch( + "application.api.user.attachments.routes.db_readonly", _fake_readonly + ), patch("application.api.user.base.storage", mock_storage): + with app.test_request_context( + f"/api/images/{self.agent_id}/{'0' * 64}", method="GET", ): resource = ServeImage() - response = resource.get("missing/image.png") + response = resource.get(self.agent_id, "0" * 64) + assert _get_response_status(response) == 404 + mock_storage.get_file.assert_not_called() + + def test_poisoned_agent_path_never_reads_storage(self, flask_app): + from application.api.user.attachments.routes import ServeImage + + app = Flask(__name__) + mock_storage = MagicMock() + agent = {"id": self.agent_id, "user_id": "user123", "image": ".env"} + + with patch( + "application.utils.settings.JWT_SECRET_KEY", "test-image-secret" + ), patch( + "application.api.user.attachments.routes.AgentsRepository.find_image_record", + return_value=agent, + ), patch( + "application.api.user.attachments.routes.db_readonly", _fake_readonly + ), patch("application.api.user.base.storage", mock_storage): + capability = self._capability(self.agent_id, ".env") + with app.test_request_context( + f"/api/images/{self.agent_id}/{capability}", method="GET" + ): + response = ServeImage().get(self.agent_id, capability) + assert _get_response_status(response) == 404 + mock_storage.get_file.assert_not_called() + + def test_serve_image_file_not_found(self, flask_app): + from application.api.user.attachments.routes import ServeImage + + app = Flask(__name__) + mock_storage = MagicMock() + mock_storage.get_file_size.side_effect = FileNotFoundError("not found") + agent = { + "id": self.agent_id, + "user_id": "user123", + "image": self.image_path, + } + + with patch( + "application.utils.settings.JWT_SECRET_KEY", "test-image-secret" + ), patch( + "application.api.user.attachments.routes.AgentsRepository.find_image_record", + return_value=agent, + ), patch( + "application.api.user.attachments.routes.db_readonly", _fake_readonly + ), patch("application.api.user.base.storage", mock_storage): + capability = self._capability(self.agent_id, self.image_path) + with app.test_request_context( + f"/api/images/{self.agent_id}/{capability}", method="GET" + ): + response = ServeImage().get(self.agent_id, capability) assert _get_response_status(response) == 404 def test_serve_image_generic_error(self, flask_app): @@ -1577,15 +1783,29 @@ class TestServeImage: app = Flask(__name__) mock_storage = MagicMock() - mock_storage.get_file.side_effect = Exception("storage error") + mock_storage.get_file_size.side_effect = Exception("storage error") - with patch("application.api.user.base.storage", mock_storage): + agent = { + "id": self.agent_id, + "user_id": "user123", + "image": self.image_path, + } + + with patch( + "application.utils.settings.JWT_SECRET_KEY", "test-image-secret" + ), patch( + "application.api.user.attachments.routes.AgentsRepository.find_image_record", + return_value=agent, + ), patch( + "application.api.user.attachments.routes.db_readonly", _fake_readonly + ), patch("application.api.user.base.storage", mock_storage): + capability = self._capability(self.agent_id, self.image_path) with app.test_request_context( - "/api/images/broken/image.png", + f"/api/images/{self.agent_id}/{capability}", method="GET", ): resource = ServeImage() - response = resource.get("broken/image.png") + response = resource.get(self.agent_id, capability) assert _get_response_status(response) == 500 diff --git a/tests/api/user/sources/test_upload.py b/tests/api/user/sources/test_upload.py index fbe84a31..566cf4c0 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,144 @@ 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 response.json["message"] == ( + "documents.zip: ZIP archive failed safety validation" + ) + assert "bad" not in response.json["message"] + assert "\n" not in response.json["message"] + logged_values = " ".join(str(value) for value in warning.call_args.args) + assert "bad\\x0aFORGED" in logged_values + 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 +614,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_base.py b/tests/api/user/test_base.py index 06a1545d..75145f6f 100644 --- a/tests/api/user/test_base.py +++ b/tests/api/user/test_base.py @@ -4,6 +4,7 @@ from contextlib import contextmanager from unittest.mock import Mock, patch import pytest +from PIL import Image from werkzeug.datastructures import FileStorage @@ -112,7 +113,19 @@ class TestGetVectorStore: @pytest.mark.unit class TestHandleImageUpload: - pass + @staticmethod + def _image_bytes( + width: int = 1, height: int = 1, image_format: str = "PNG" + ) -> bytes: + buffer = io.BytesIO() + Image.new("RGB", (width, height), color="white").save( + buffer, format=image_format + ) + return buffer.getvalue() + + @classmethod + def _png_bytes(cls) -> bytes: + return cls._image_bytes() def test_returns_existing_url_when_no_file(self, flask_app): from application.api.user.base import handle_image_upload @@ -135,7 +148,7 @@ class TestHandleImageUpload: with flask_app.test_request_context(): mock_file = FileStorage( - stream=io.BytesIO(b"fake image data"), filename="test_image.png" + stream=io.BytesIO(self._png_bytes()), filename="test_image.png" ) mock_request = Mock() mock_request.files = {"image": mock_file} @@ -152,6 +165,46 @@ class TestHandleImageUpload: assert "user123" in url mock_storage.save_file.assert_called_once() + def test_accepts_multi_picture_jpeg(self, flask_app): + from application.api.user.base import handle_image_upload + + # Pillow reports multi-picture JPEGs (e.g. iPhone portrait photos) as + # MPO; they must still be accepted under a .jpg extension. + buffer = io.BytesIO() + frame = Image.new("RGB", (1, 1), color="white") + frame.save(buffer, format="MPO", save_all=True, append_images=[frame]) + + with flask_app.test_request_context(): + mock_file = FileStorage( + stream=io.BytesIO(buffer.getvalue()), filename="photo.jpg" + ) + mock_storage = Mock() + + url, error = handle_image_upload( + Mock(files={"image": mock_file}), "", "user123", mock_storage + ) + + assert error is None + assert url is not None + mock_storage.save_file.assert_called_once() + + def test_uploads_image_with_non_ascii_basename(self, flask_app): + from application.api.user.base import handle_image_upload + + with flask_app.test_request_context(): + mock_file = FileStorage( + stream=io.BytesIO(self._png_bytes()), filename="上传.png" + ) + mock_storage = Mock() + + url, error = handle_image_upload( + Mock(files={"image": mock_file}), "", "user123", mock_storage + ) + + assert error is None + assert url.endswith("_avatar.png") + mock_storage.save_file.assert_called_once() + def test_ignores_empty_filename(self, flask_app): from application.api.user.base import handle_image_upload @@ -175,7 +228,9 @@ class TestHandleImageUpload: from application.api.user.base import handle_image_upload with flask_app.app_context(): - mock_file = FileStorage(stream=io.BytesIO(b"data"), filename="test.png") + mock_file = FileStorage( + stream=io.BytesIO(self._png_bytes()), filename="test.png" + ) mock_request = Mock() mock_request.files = {"image": mock_file} mock_storage = Mock() @@ -189,6 +244,135 @@ class TestHandleImageUpload: assert error is not None assert error.status_code == 400 + def test_rejects_non_image_content(self, flask_app): + from application.api.user.base import handle_image_upload + + with flask_app.app_context(): + mock_file = FileStorage( + stream=io.BytesIO(b"not an image"), filename="test.png" + ) + mock_request = Mock(files={"image": mock_file}) + mock_storage = Mock() + + url, error = handle_image_upload( + mock_request, "old.jpg", "user123", mock_storage + ) + + assert url is None + assert error.status_code == 400 + mock_storage.save_file.assert_not_called() + + def test_accepts_image_at_encoded_byte_limit(self, flask_app): + from application.api.user.base import handle_image_upload + + payload = self._png_bytes() + with patch( + "application.api.user.base.settings.AGENT_IMAGE_MAX_BYTES", + len(payload), + ), flask_app.app_context(): + mock_file = FileStorage( + stream=io.BytesIO(payload), filename="at-limit.png" + ) + mock_storage = Mock() + + url, error = handle_image_upload( + Mock(files={"image": mock_file}), "", "user123", mock_storage + ) + + assert error is None + assert url.endswith("_at-limit.png") + mock_storage.save_file.assert_called_once() + + def test_rejects_image_over_encoded_byte_limit(self, flask_app): + from application.api.user.base import handle_image_upload + + payload = self._png_bytes() + with patch( + "application.api.user.base.settings.AGENT_IMAGE_MAX_BYTES", + len(payload) - 1, + ), flask_app.app_context(): + mock_file = FileStorage( + stream=io.BytesIO(payload), filename="too-large.png" + ) + mock_storage = Mock() + + url, error = handle_image_upload( + Mock(files={"image": mock_file}), "old.png", "user123", mock_storage + ) + + assert url is None + assert error.status_code == 400 + mock_storage.save_file.assert_not_called() + + @pytest.mark.parametrize( + ("width", "height", "expected_status"), + [(4, 4, None), (5, 4, 400)], + ) + def test_enforces_decoded_pixel_limit( + self, flask_app, width, height, expected_status + ): + from application.api.user.base import handle_image_upload + + payload = self._image_bytes(width, height) + with patch( + "application.api.user.base.settings.AGENT_IMAGE_MAX_PIXELS", 16 + ), flask_app.app_context(): + mock_file = FileStorage( + stream=io.BytesIO(payload), filename="dimensions.png" + ) + mock_storage = Mock() + + url, error = handle_image_upload( + Mock(files={"image": mock_file}), "", "user123", mock_storage + ) + + if expected_status is None: + assert error is None + assert url + mock_storage.save_file.assert_called_once() + else: + assert url is None + assert error.status_code == expected_status + mock_storage.save_file.assert_not_called() + + def test_rejects_image_whose_content_does_not_match_extension( + self, flask_app + ): + from application.api.user.base import handle_image_upload + + with flask_app.app_context(): + mock_file = FileStorage( + stream=io.BytesIO(self._png_bytes()), filename="disguised.jpg" + ) + mock_storage = Mock() + + url, error = handle_image_upload( + Mock(files={"image": mock_file}), "", "user123", mock_storage + ) + + assert url is None + assert error.status_code == 400 + mock_storage.save_file.assert_not_called() + + def test_sanitizes_user_directory_component(self, flask_app): + from application.api.user.base import handle_image_upload + + with flask_app.app_context(): + mock_file = FileStorage( + stream=io.BytesIO(self._png_bytes()), filename="test.png" + ) + mock_request = Mock(files={"image": mock_file}) + mock_storage = Mock() + + url, error = handle_image_upload( + mock_request, "", "../../secrets", mock_storage + ) + + assert error is None + assert ".." not in url + assert url.startswith("inputs/secrets-") + assert "/attachments/" in url + @pytest.mark.unit class TestRequireAgentDecorator: @@ -321,4 +505,3 @@ class TestRequireAgentPgConn: with _patch_base_db(pg_conn), flask_app.app_context(): result = fn(webhook_token="webhook-123") assert result["got"] == str(agent["id"]) - 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/core/test_secret_key.py b/tests/core/test_secret_key.py new file mode 100644 index 00000000..75195cde --- /dev/null +++ b/tests/core/test_secret_key.py @@ -0,0 +1,36 @@ +import stat + +import pytest + + +def test_configured_secret_is_used_without_touching_the_filesystem(tmp_path): + from application.core.secret_key import resolve_jwt_secret_key + + key_file = tmp_path / "missing" / "jwt-secret" + + assert ( + resolve_jwt_secret_key("configured-secret", None, key_file) + == "configured-secret" + ) + assert not key_file.exists() + + +def test_cloud_deployment_requires_an_explicit_shared_secret(tmp_path): + from application.core.secret_key import resolve_jwt_secret_key + + with pytest.raises(RuntimeError, match="JWT_SECRET_KEY must be set"): + resolve_jwt_secret_key("", "cloud", tmp_path / "jwt-secret") + + +def test_local_secret_is_created_once_with_owner_only_permissions(tmp_path): + from application.core.secret_key import resolve_jwt_secret_key + + key_file = tmp_path / "jwt-secret" + + first = resolve_jwt_secret_key("", None, key_file) + second = resolve_jwt_secret_key("", None, key_file) + + assert second == first + assert key_file.read_text(encoding="utf-8") == first + assert len(first) == 64 + assert stat.S_IMODE(key_file.stat().st_mode) == 0o600 diff --git a/tests/e2e/specs/tier-b/images.spec.ts b/tests/e2e/specs/tier-b/images.spec.ts index 6c09e59d..8eaa1669 100644 --- a/tests/e2e/specs/tier-b/images.spec.ts +++ b/tests/e2e/specs/tier-b/images.spec.ts @@ -1,34 +1,6 @@ -// Tier-B · B17 image serve endpoint. -/** - * Exercises `GET /api/images/` in - * `application/api/user/attachments/routes.py` — the static-ish endpoint - * that fronts `storage.get_file()` for user-uploaded agent images and - * anything else the UI references via `generate_image_url(...)`. The route - * is not authenticated (no `@token_required` wrapping) by design — the - * image URLs are shared publicly e.g. in the agents marketplace — so we - * don't set up a `newUserContext` here, unlike every other Tier-B spec. - * - * Storage contract gotcha: `LocalStorage.__init__` sets `base_dir` to the - * repo root (three `os.path.dirname()` calls up from - * `application/storage/local.py`). The e2e env sets `UPLOAD_FOLDER=.e2e-tmp/inputs` - * but that is the directory the upload handlers write INTO; it is not the - * storage base. So an image request for `/api/images/foo/bar.png` resolves - * to `/foo/bar.png`. The tests below write their fixture files - * under `/.e2e-tmp/inputs/e2e-images/` and request - * `/api/images/.e2e-tmp/inputs/e2e-images/` — the `.` prefix is - * NOT `..` so it passes the route's literal-substring guard. - * - * Covered: - * 1. served file: write a small PNG under storage, GET returns bytes + - * correct Content-Type + Cache-Control header. - * 2. path traversal: `/api/images/../../etc/passwd` — rejected 400 before - * hitting storage. Also URL-encoded %2E%2E variant. - * 3. missing file: GET against a known-absent path → 404. - * 4. non-image extensions serve with a predictable content-type (the - * route trusts the extension and returns `image/`). - * 5. jpg alias: extension `.jpg` is served as `image/jpeg` (not `image/jpg`). - */ +/** Tier-B · public agent-image capability endpoint. */ +import { createHmac } from 'node:crypto'; import { mkdir, rm, writeFile } from 'node:fs/promises'; import { dirname, resolve } from 'node:path'; import { fileURLToPath } from 'node:url'; @@ -36,146 +8,139 @@ import { fileURLToPath } from 'node:url'; import * as playwright from '@playwright/test'; const { expect, test } = playwright; +import { pg } from '../../helpers/db.js'; +import { resetDb } from '../../helpers/reset.js'; + const HERE = dirname(fileURLToPath(import.meta.url)); -// Repo root = tests/e2e/specs/tier-b/ → up 4 = repo root. const REPO_ROOT = resolve(HERE, '..', '..', '..', '..'); -// Scratch dir we own — writes go here, TRUNCATE doesn't touch the filesystem. -const IMAGES_DIR = resolve(REPO_ROOT, '.e2e-tmp', 'inputs', 'e2e-images'); -// Path used in URLs — relative-to-repo-root, matches what `storage.get_file` -// will resolve via its `base_dir` join. -const IMAGES_URL_PREFIX = '.e2e-tmp/inputs/e2e-images'; - +const IMAGE_USER = 'e2e-images'; +const IMAGES_DIR = resolve( + REPO_ROOT, + '.e2e-tmp', + 'inputs', + IMAGE_USER, + 'attachments', +); +const STORAGE_PREFIX = `.e2e-tmp/inputs/${IMAGE_USER}/attachments`; const API_URL = process.env.API_URL ?? 'http://127.0.0.1:7099'; +const IMAGE_SECRET = + process.env.JWT_SECRET_KEY ?? 'e2e-fixed-secret-never-use-in-prod'; -// Minimal valid PNG — 1x1 transparent pixel. Borrowed from the commonly- -// reproduced public-domain snippet; it's a real PNG that decoders accept. -// Using a real PNG (vs. a fake 8-byte header) means a future test that -// actually tries to render it won't break. -const TINY_PNG_BASE64 = - 'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk+A8AAQUBAScY42YAAAAASUVORK5CYII='; +const TINY_PNG = Buffer.from( + 'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk+A8AAQUBAScY42YAAAAASUVORK5CYII=', + 'base64', +); -async function seedImage( - relName: string, - data: Buffer = Buffer.from(TINY_PNG_BASE64, 'base64'), -): Promise { - await mkdir(IMAGES_DIR, { recursive: true }); - const fullPath = resolve(IMAGES_DIR, relName); - await writeFile(fullPath, data); - return `${IMAGES_URL_PREFIX}/${relName}`; +function imageCapability( + agentId: string, + imagePath: string, + userId = IMAGE_USER, +): string { + const payload = `docsgpt-agent-image-v1\0${agentId}\0${userId}\0${imagePath}`; + return createHmac('sha256', IMAGE_SECRET).update(payload).digest('hex'); } -test.describe('tier-b · images (public image serve)', () => { - // Filesystem is shared across tests in this file; each test writes under - // a unique filename to avoid collision. After the suite we clean the dir. +async function seedAgentImage( + filename: string, + data: Buffer = TINY_PNG, +): Promise<{ agentId: string; imagePath: string; url: string }> { + await mkdir(IMAGES_DIR, { recursive: true }); + await writeFile(resolve(IMAGES_DIR, filename), data); + const imagePath = `${STORAGE_PREFIX}/${filename}`; + const { rows } = await pg.query<{ id: string }>( + `INSERT INTO agents (user_id, name, status, image) + VALUES ($1, 'image-capability-e2e', 'draft', $2) + RETURNING id::text AS id`, + [IMAGE_USER, imagePath], + ); + const agentId = rows[0]?.id; + if (!agentId) throw new Error('Failed to seed image agent'); + const capability = imageCapability(agentId, imagePath); + return { + agentId, + imagePath, + url: `/api/images/${agentId}/${capability}`, + }; +} + +test.describe('tier-b · agent image capabilities', () => { + test.beforeEach(async () => { + await resetDb(); + }); + test.afterAll(async () => { - await rm(IMAGES_DIR, { recursive: true, force: true }); + await rm(resolve(IMAGES_DIR, '..'), { recursive: true, force: true }); }); - test('serves a seeded PNG with correct content-type and cache header', async () => { + test('valid capability serves exact PNG bytes with hardened headers', async () => { const api = await playwright.request.newContext({ baseURL: API_URL }); try { - const imagePath = await seedImage('served.png'); - const res = await api.get(`/api/images/${imagePath}`); + const seeded = await seedAgentImage('served.png'); + const res = await api.get(seeded.url); expect(res.status()).toBe(200); - - const headers = res.headers(); - expect(headers['content-type']).toBe('image/png'); - - // The route sets `Cache-Control: max-age=86400`. Spec asserts presence - // (don't pin the specific value — a future 1h/1d swap shouldn't fail). - const cacheControl = headers['cache-control']; - expect(cacheControl, 'Cache-Control header should be present').toBeTruthy(); - expect(cacheControl).toMatch(/max-age=\d+/); - - // Body is the exact bytes we wrote. - const expected = Buffer.from(TINY_PNG_BASE64, 'base64'); - const body = await res.body(); - expect(body.equals(expected)).toBe(true); + expect(res.headers()['content-type']).toBe('image/png'); + expect(res.headers()['cache-control']).toMatch(/max-age=\d+/); + expect(res.headers()['x-content-type-options']).toBe('nosniff'); + expect((await res.body()).equals(TINY_PNG)).toBe(true); } finally { await api.dispose(); } }); - test('path traversal rejected (literal ..) → 400', async () => { - // Playwright's APIRequestContext resolves URLs through the WHATWG URL - // parser which collapses `..` segments before sending. To actually hit - // Flask with `..` in the path we must URL-encode the dots so Playwright - // passes them through unmodified — Werkzeug decodes them back to `..` - // and the route's substring guard then triggers 400. + test('raw storage paths are no longer accepted', async () => { const api = await playwright.request.newContext({ baseURL: API_URL }); try { - const encoded = await api.get( - '/api/images/%2e%2e%2f%2e%2e%2fetc%2fpasswd', - ); - expect(encoded.status()).toBe(400); - const body = (await encoded.json()) as { success: boolean; message?: string }; - expect(body.success).toBe(false); - expect(body.message ?? '').toMatch(/invalid/i); - - // Also the mixed-path traversal case where the literal `..` sits - // inside the trailing path (no leading-slash collapse because of the - // non-dot first segment). Encoded so Playwright can't normalise it away. - const encoded2 = await api.get( - '/api/images/e2e-images/%2e%2e%2fbadfile', - ); - expect(encoded2.status()).toBe(400); - } finally { - await api.dispose(); - } - }); - - test('missing file returns 404', async () => { - const api = await playwright.request.newContext({ baseURL: API_URL }); - try { - // Ensure the directory exists (so the error is truly "file not found", - // not a directory-missing surfaced as FileNotFoundError from a missing - // intermediate segment). - await mkdir(IMAGES_DIR, { recursive: true }); - const res = await api.get( - `/api/images/${IMAGES_URL_PREFIX}/definitely-not-here-${Date.now()}.png`, - ); + const res = await api.get('/api/images/.env'); expect(res.status()).toBe(404); - const body = (await res.json()) as { success: boolean; message?: string }; - expect(body.success).toBe(false); } finally { await api.dispose(); } }); - test('non-image extension is served with image/ content-type (route trusts the extension)', async () => { - // The handler does not inspect bytes — it splits the path by `.` and - // returns `image/`. This is a deliberate contract (legacy agents - // may reference `.webp`, `.gif`, etc.) but a weird one: `.txt` would - // also be served as `image/txt`. Pin that behaviour so a future MIME- - // sniffer refactor trips this test loudly rather than silently changing - // what browsers see. + test('forged capability is rejected before storage access', async () => { const api = await playwright.request.newContext({ baseURL: API_URL }); try { - const imagePath = await seedImage( - 'not-really.webp', - Buffer.from('RIFFxxxxWEBP', 'utf8'), + const seeded = await seedAgentImage('forged.png'); + const res = await api.get(`/api/images/${seeded.agentId}/${'0'.repeat(64)}`); + expect(res.status()).toBe(404); + } finally { + await api.dispose(); + } + }); + + test('even a signed poisoned database path is rejected', async () => { + const api = await playwright.request.newContext({ baseURL: API_URL }); + try { + const { rows } = await pg.query<{ id: string }>( + `INSERT INTO agents (user_id, name, status, image) + VALUES ($1, 'poisoned-image-e2e', 'draft', '.env') + RETURNING id::text AS id`, + [IMAGE_USER], ); - const res = await api.get(`/api/images/${imagePath}`); - expect(res.status()).toBe(200); - expect(res.headers()['content-type']).toBe('image/webp'); + const agentId = rows[0]?.id; + if (!agentId) throw new Error('Failed to seed poisoned image agent'); + const capability = imageCapability(agentId, '.env'); + const res = await api.get(`/api/images/${agentId}/${capability}`); + expect(res.status()).toBe(404); } finally { await api.dispose(); } }); - test('jpg extension is aliased to image/jpeg', async () => { + test('jpg uses image/jpeg while non-image extensions fail closed', async () => { const api = await playwright.request.newContext({ baseURL: API_URL }); try { - const imagePath = await seedImage( + const jpg = await seedAgentImage( 'aliased.jpg', - Buffer.from('not-a-real-jpeg-just-bytes', 'utf8'), + Buffer.from([0xff, 0xd8, 0xff, 0xe0]), ); - const res = await api.get(`/api/images/${imagePath}`); - expect(res.status()).toBe(200); - // The route explicitly branches `jpg -> image/jpeg` — verifying here - // so a refactor to a generic mime-type map can't quietly regress to - // `image/jpg` (which Safari rejects). - expect(res.headers()['content-type']).toBe('image/jpeg'); + const jpgRes = await api.get(jpg.url); + expect(jpgRes.status()).toBe(200); + expect(jpgRes.headers()['content-type']).toBe('image/jpeg'); + + const text = await seedAgentImage('not-an-image.txt', Buffer.from('secret')); + const textRes = await api.get(text.url); + expect(textRes.status()).toBe(404); } finally { await api.dispose(); } diff --git a/tests/storage/db/repositories/test_agents.py b/tests/storage/db/repositories/test_agents.py index 394e5364..4a20c6b5 100644 --- a/tests/storage/db/repositories/test_agents.py +++ b/tests/storage/db/repositories/test_agents.py @@ -71,6 +71,18 @@ class TestGet: created = repo.create("user-1", "a", "draft") assert repo.get(created["id"], "user-other") is None + def test_find_image_record_returns_only_capability_fields(self, pg_conn): + repo = _repo(pg_conn) + created = repo.create( + "user-1", "a", "draft", image="inputs/user-1/attachments/a.png" + ) + + image_record = repo.find_image_record(str(created["id"])) + + assert set(image_record) >= {"id", "user_id", "image"} + assert image_record["image"] == "inputs/user-1/attachments/a.png" + assert "name" not in image_record + def test_get_by_legacy_id(self, pg_conn): repo = _repo(pg_conn) created = repo.create( diff --git a/tests/storage/test_local_storage.py b/tests/storage/test_local_storage.py index 196b0621..0c68de0e 100644 --- a/tests/storage/test_local_storage.py +++ b/tests/storage/test_local_storage.py @@ -123,6 +123,23 @@ class TestLocalStorageGetFile: ) +@pytest.mark.unit +class TestLocalStorageFileSize: + + @patch("os.path.getsize", return_value=42) + def test_get_file_size_uses_metadata_without_opening( + self, mock_getsize, local_storage + ): + assert local_storage.get_file_size("avatars/test.png") == 42 + + expected_path = os.path.join( + os.path.realpath("/tmp/test_storage"), "avatars/test.png" + ) + assert os.path.normpath(mock_getsize.call_args[0][0]) == os.path.normpath( + expected_path + ) + + @pytest.mark.unit class TestLocalStorageDeleteFile: diff --git a/tests/storage/test_s3_storage.py b/tests/storage/test_s3_storage.py index aed2a9b2..5cdbaa53 100644 --- a/tests/storage/test_s3_storage.py +++ b/tests/storage/test_s3_storage.py @@ -242,6 +242,73 @@ class TestS3StorageGetFile: s3_storage.get_file(path) +class TestS3StorageFileSize: + """Test metadata-only size lookup used by bounded public delivery.""" + + @pytest.mark.unit + def test_get_file_size_uses_head_without_downloading( + self, s3_storage, mock_boto3_client + ): + mock_boto3_client.head_object.return_value = {"ContentLength": 1234} + + assert s3_storage.get_file_size("avatars/test.png") == 1234 + + mock_boto3_client.head_object.assert_called_once_with( + Bucket="test-bucket", Key="avatars/test.png" + ) + mock_boto3_client.download_fileobj.assert_not_called() + + @pytest.mark.unit + def test_get_file_size_maps_missing_object_to_file_not_found( + self, s3_storage, mock_boto3_client + ): + mock_boto3_client.head_object.side_effect = ClientError( + {"Error": {"Code": "NoSuchKey", "Message": "Not found"}}, + "head_object", + ) + + with pytest.raises(FileNotFoundError, match="File not found"): + s3_storage.get_file_size("avatars/missing.png") + + @pytest.mark.unit + def test_get_file_size_maps_denied_head_to_file_not_found( + self, s3_storage, mock_boto3_client + ): + # Without s3:ListBucket, HEAD on a missing key returns 403, not 404. + mock_boto3_client.head_object.side_effect = ClientError( + {"Error": {"Code": "403", "Message": "Forbidden"}}, + "head_object", + ) + + with pytest.raises(FileNotFoundError, match="File not found"): + s3_storage.get_file_size("avatars/missing.png") + + +class TestS3StoragePresignedUrl: + """Test response metadata on short-lived direct avatar delivery URLs.""" + + @pytest.mark.unit + def test_generate_presigned_url_can_override_content_type( + self, s3_storage, mock_boto3_client + ): + mock_boto3_client.generate_presigned_url.return_value = "https://signed" + + result = s3_storage.generate_presigned_url( + "avatars/test.png", expires_in=300, content_type="image/png" + ) + + assert result == "https://signed" + mock_boto3_client.generate_presigned_url.assert_called_once_with( + "get_object", + Params={ + "Bucket": "test-bucket", + "Key": "avatars/test.png", + "ResponseContentType": "image/png", + }, + ExpiresIn=300, + ) + + class TestS3StorageDeleteFile: """Test file deletion functionality.""" 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_utils.py b/tests/test_utils.py index b8e15048..16ec3852 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -16,6 +16,7 @@ from application.utils import ( get_hash, get_missing_fields, generate_image_url, + is_safe_agent_image_path, limit_chat_history, num_tokens_from_object_or_list, num_tokens_from_string, @@ -317,42 +318,114 @@ class TestGenerateImageUrl: assert generate_image_url("http://example.com/img.png") == "http://example.com/img.png" @pytest.mark.unit - def test_s3_strategy(self): + def test_internal_image_uses_opaque_capability(self): with patch("application.utils.settings") as s: - s.URL_STRATEGY = "s3" - s.S3_BUCKET_NAME = "my-bucket" - s.S3_ENDPOINT_URL = None - s.S3_REGION = "us-west-2" - result = generate_image_url("path/to/img.png") - assert result == "https://my-bucket.s3.us-west-2.amazonaws.com/path/to/img.png" + s.JWT_SECRET_KEY = "test-image-secret" + s.UPLOAD_FOLDER = "inputs" + s.API_URL = "https://api.example.com" + result = generate_image_url( + "inputs/user-1/attachments/avatar.png", + "00000000-0000-0000-0000-000000000001", + "user-1", + ) + assert result.startswith( + "https://api.example.com/api/images/" + "00000000-0000-0000-0000-000000000001/" + ) + assert "inputs" not in result + assert "avatar.png" not in result @pytest.mark.unit - def test_s3_strategy_custom_endpoint_path_style(self): + def test_internal_image_fails_closed_without_agent_identity(self): with patch("application.utils.settings") as s: - s.URL_STRATEGY = "s3" - s.S3_BUCKET_NAME = "my-bucket" - s.S3_ENDPOINT_URL = "https://account.r2.cloudflarestorage.com" - s.S3_PATH_STYLE = True - result = generate_image_url("path/to/img.png") - assert result == "https://account.r2.cloudflarestorage.com/my-bucket/path/to/img.png" + s.JWT_SECRET_KEY = "test-image-secret" + s.UPLOAD_FOLDER = "inputs" + assert generate_image_url("inputs/user-1/attachments/avatar.png") == "" @pytest.mark.unit - def test_s3_strategy_custom_endpoint_virtual_host(self): + def test_internal_image_fails_closed_without_secret(self): with patch("application.utils.settings") as s: - s.URL_STRATEGY = "s3" - s.S3_BUCKET_NAME = "my-bucket" - s.S3_ENDPOINT_URL = "https://minio.example.com" - s.S3_PATH_STYLE = False - result = generate_image_url("path/to/img.png") - assert result == "https://my-bucket.minio.example.com/path/to/img.png" + s.JWT_SECRET_KEY = "" + s.UPLOAD_FOLDER = "inputs" + result = generate_image_url( + "inputs/user-1/attachments/avatar.png", + "00000000-0000-0000-0000-000000000001", + "user-1", + ) + assert result == "" @pytest.mark.unit - def test_backend_strategy(self): + def test_internal_image_rejects_path_outside_owner_uploads(self): with patch("application.utils.settings") as s: - s.URL_STRATEGY = "backend" + s.JWT_SECRET_KEY = "test-image-secret" + s.UPLOAD_FOLDER = "inputs" s.API_URL = "http://localhost:7091" - result = generate_image_url("path/to/img.png") - assert result == "http://localhost:7091/api/images/path/to/img.png" + result = generate_image_url( + ".env", + "00000000-0000-0000-0000-000000000001", + "user-1", + ) + assert result == "" + + @pytest.mark.unit + def test_absolute_upload_folder_still_serves_owned_images(self): + with patch("application.utils.settings") as s: + s.JWT_SECRET_KEY = "test-image-secret" + s.UPLOAD_FOLDER = "/data/inputs" + s.API_URL = "https://api.example.com" + result = generate_image_url( + "/data/inputs/user-1/attachments/avatar.png", + "00000000-0000-0000-0000-000000000001", + "user-1", + ) + assert result.startswith( + "https://api.example.com/api/images/" + "00000000-0000-0000-0000-000000000001/" + ) + + +class TestIsSafeAgentImagePath: + + @pytest.mark.unit + def test_absolute_upload_folder_accepts_owned_path(self): + with patch("application.utils.settings") as s: + s.UPLOAD_FOLDER = "/data/inputs" + assert is_safe_agent_image_path( + "/data/inputs/user-1/attachments/avatar.png", "user-1" + ) + + @pytest.mark.unit + def test_absolute_upload_folder_rejects_other_owner(self): + with patch("application.utils.settings") as s: + s.UPLOAD_FOLDER = "/data/inputs" + assert not is_safe_agent_image_path( + "/data/inputs/user-2/attachments/avatar.png", "user-1" + ) + + @pytest.mark.unit + def test_absolute_upload_folder_rejects_escape(self): + with patch("application.utils.settings") as s: + s.UPLOAD_FOLDER = "/data/inputs" + assert not is_safe_agent_image_path( + "/data/inputs/user-1/attachments/../../../etc/passwd.png", "user-1" + ) + assert not is_safe_agent_image_path("/etc/passwd.png", "user-1") + + @pytest.mark.unit + def test_relative_upload_folder_rejects_absolute_path(self): + with patch("application.utils.settings") as s: + s.UPLOAD_FOLDER = "inputs" + assert not is_safe_agent_image_path( + "/inputs/user-1/attachments/avatar.png", "user-1" + ) + + @pytest.mark.unit + def test_relative_upload_folder_accepts_owned_path(self): + with patch("application.utils.settings") as s: + s.UPLOAD_FOLDER = "inputs" + assert is_safe_agent_image_path( + "inputs/user-1/attachments/avatar.png", "user-1" + ) class TestCalculateCompressionThreshold: @@ -641,17 +714,20 @@ class TestGenerateImageUrlEdgeCases: @pytest.mark.unit def test_non_string_input(self): result = generate_image_url(123) - # Not a string, not starting with http, uses default strategy - assert "/api/images/" in result or "s3" in result + assert result == "" @pytest.mark.unit - def test_default_strategy_is_backend(self): + def test_email_owner_legacy_path_is_supported(self): with patch("application.utils.settings") as s: - # Simulate missing URL_STRATEGY attribute - del s.URL_STRATEGY + s.JWT_SECRET_KEY = "test-image-secret" + s.UPLOAD_FOLDER = "inputs" s.API_URL = "http://localhost:7091" - result = generate_image_url("img.png") - assert "localhost:7091" in result + result = generate_image_url( + "inputs/person@example.com/attachments/img.png", + "00000000-0000-0000-0000-000000000001", + "person@example.com", + ) + assert result.startswith("http://localhost:7091/api/images/") class TestGetHashEdgeCases: diff --git a/tests/test_worker_utils.py b/tests/test_worker_utils.py index e8c11007..bd6e4021 100644 --- a/tests/test_worker_utils.py +++ b/tests/test_worker_utils.py @@ -203,6 +203,30 @@ 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, ZipExtractionError + + 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, + ): + with pytest.raises(ZipExtractionError): + 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..1b0b6dca 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,14 +280,14 @@ 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) assert not os.path.exists(os.path.join(extract_to, "inner.zip")) def test_respects_max_depth(self): - """Extraction should stop at max recursion depth.""" + """Nesting beyond max depth should fail the extraction loudly.""" with tempfile.TemporaryDirectory() as temp_dir: extract_to = os.path.join(temp_dir, "extract") os.makedirs(extract_to) @@ -192,10 +311,8 @@ class TestExtractZipRecursive: f.write(current_content) # Extract with max_depth=2 - extract_zip_recursive(zip_path, extract_to, max_depth=2) - - # The deepest nested zips should remain unextracted - # (we can't easily verify the exact behavior, but the function should not crash) + with pytest.raises(ZipExtractionError): + extract_zip_recursive(zip_path, extract_to, max_depth=2) def test_rejects_path_traversal(self): """Zip with path traversal should be rejected and removed.""" @@ -208,7 +325,8 @@ class TestExtractZipRecursive: with zipfile.ZipFile(zip_path, "w") as zf: zf.writestr("../../../tmp/malicious.txt", "malicious") - extract_zip_recursive(zip_path, extract_to) + with pytest.raises(ZipExtractionError): + extract_zip_recursive(zip_path, extract_to) # Zip should be removed assert not os.path.exists(zip_path) @@ -216,8 +334,26 @@ class TestExtractZipRecursive: # Malicious file should NOT exist outside extract_to assert not os.path.exists("/tmp/malicious.txt") - def test_handles_corrupted_zip_gracefully(self): - """Corrupted zip should be handled gracefully without crashing.""" + 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") + + with pytest.raises(ZipExtractionError): + extract_zip_recursive(zip_path, extract_to) + + assert not os.path.lexists(os.path.join(extract_to, "escape-link")) + + def test_corrupted_zip_fails_loudly(self): + """Corrupted zip should raise so ingestion fails instead of indexing nothing.""" with tempfile.TemporaryDirectory() as temp_dir: zip_path = os.path.join(temp_dir, "corrupted.zip") extract_to = os.path.join(temp_dir, "extract") @@ -227,10 +363,11 @@ class TestExtractZipRecursive: with open(zip_path, "wb") as f: f.write(b"This is not a valid zip file") - # Should not raise, just log error - extract_zip_recursive(zip_path, extract_to) + with pytest.raises(ZipExtractionError): + extract_zip_recursive(zip_path, extract_to) - # Function should complete without exception + # The rejected zip should be removed + assert not os.path.exists(zip_path) class TestZipBombProtection: