From 027fdc09bb7a7a47ac30bb25dc57b37bed9c8fde Mon Sep 17 00:00:00 2001 From: Vito Sansevero Date: Wed, 7 Oct 2026 18:46:16 -0700 Subject: [PATCH] fix: cap every request-driven resource in the image routes - /images/recent, /images/all and /images/output parse limit/offset through parse_page_params on every path (absent, 0, negative, float strings), and the clamped values are what reach the DB. /images/all gets its own MAX_BULK_LIMIT ceiling instead of an unbounded query; /images/recent caps its fetch window at MAX_RECENT_IMAGES_WINDOW. - parse_page_params also clamps offset to MAX_PAGE_OFFSET so values past int64 no longer surface as a 500 from SQLite. - Gallery and thumbnail scans share _iter_media_files: no symlink following, MAX_SCAN_DEPTH, and MAX_GALLERY_FILES / MAX_THUMBNAIL_FILES caps. Thumbnail size comes only from named presets. - JSON bodies on /images/link and /images/generate-thumbnails are read in bounded chunks; over MAX_JSON_BODY_BYTES answers 413, malformed or non-object JSON answers 400. - /images/link only accepts media paths inside the allowed directories (403 otherwise), closing a path-existence oracle. --- py/api/images.py | 365 ++++++++++++++++++++++--------------- py/api/prompts.py | 13 +- tests/test_api_images.py | 369 +++++++++++++++++++++++++++++++++++++- tests/test_api_prompts.py | 15 +- 4 files changed, 613 insertions(+), 149 deletions(-) diff --git a/py/api/images.py b/py/api/images.py index 18602e3..af73623 100644 --- a/py/api/images.py +++ b/py/api/images.py @@ -147,6 +147,120 @@ def _thumbnail_complete_payload(total, stats, start): } +# Caps on request-driven work (list sizes, scans, request bodies). +MAX_GALLERY_FILES = 50_000 +MAX_RECENT_IMAGES_WINDOW = 5_000 +MAX_BULK_LIMIT = 5_000 +MAX_JSON_BODY_BYTES = 1024 * 1024 +_READ_CHUNK = 64 * 1024 + + +class _BodyTooLarge(Exception): + """Raised when a request body exceeds MAX_JSON_BODY_BYTES.""" + + +def _json_error(message, status): + return web.json_response({"success": False, "error": message}, status=status) + + +async def _read_json_body(request): + """Read a JSON object body in bounded chunks. + + Raises _BodyTooLarge past MAX_JSON_BODY_BYTES (declared or streamed) and + ValueError when the body is not a JSON object. + """ + limit = MAX_JSON_BODY_BYTES + if request.content_length is not None and request.content_length > limit: + raise _BodyTooLarge() + chunks, size = [], 0 + while True: + chunk = await request.content.read(_READ_CHUNK) + if not chunk: + break + size += len(chunk) + if size > limit: + raise _BodyTooLarge() + chunks.append(chunk) + data = json.loads(b"".join(chunks).decode("utf-8")) + if not isinstance(data, dict): + raise ValueError("JSON body must be an object") + return data + + +async def _json_body_or_error(request): + """Return ``(data, None)`` or ``(None, error_response)``.""" + try: + return await _read_json_body(request), None + except _BodyTooLarge: + return None, _json_error("Request body too large", 413) + except (ValueError, UnicodeDecodeError): + return None, _json_error("Invalid JSON body", 400) + + +def _iter_media_files(output_path, max_files): + """Yield media files under *output_path* with bounded effort. + + Symlinks are never followed, ``thumbnails/`` is skipped, descent stops + at MAX_SCAN_DEPTH and at most *max_files* paths are produced. + """ + output_path = Path(output_path) + root_depth = len(output_path.parts) + produced = 0 + for root, dirs, files in os.walk(output_path): + root_path = Path(root) + rel_parts = root_path.parts[root_depth:] + if "thumbnails" in rel_parts: + dirs[:] = [] + continue + if len(rel_parts) >= MAX_SCAN_DEPTH: + dirs[:] = [] + for name in sorted(files): + if Path(name).suffix.lower() not in MEDIA_EXTENSIONS: + continue + if produced >= max_files: + return + produced += 1 + yield root_path / name + + +def _output_image_entry(media_path, output_path, root_index): + """Build one gallery entry for *media_path* (blocking stat).""" + stat = media_path.stat() + rel_path = media_path.relative_to(output_path) + extension = media_path.suffix.lower() + is_video = extension in VIDEO_EXTENSIONS + + thumbnail_url = None + thumb_rel = ( + f"thumbnails/{rel_path.with_suffix('').as_posix()}_thumb" + f"{'.jpg' if is_video else extension}" + ) + if (output_path / thumb_rel).exists(): + thumbnail_url = ( + f"/prompt_manager/images/serve/{urllib.parse.quote(thumb_rel, safe='/')}" + ) + + return { + "id": str(hash(str(media_path))), + "filename": media_path.name, + "path": str(media_path), + "relative_path": str(rel_path), + "root_dir": str(output_path), + "root_index": root_index, + "url": f"/prompt_manager/images/serve/{rel_path.as_posix()}", + "thumbnail_url": thumbnail_url, + "size": stat.st_size, + "modified_time": stat.st_mtime, + "extension": extension, + "media_type": "video" if is_video else "image", + "is_video": is_video, + } + + +def _in_subfolder(rel_dir, subfolder): + return rel_dir == subfolder or rel_dir.startswith(subfolder + os.sep) + + class ImageRoutesMixin: """Mixin providing image and gallery-related API endpoints.""" @@ -246,14 +360,16 @@ class ImageRoutesMixin: return web.json_response({"success": False, "error": str(e)}, status=500) async def get_recent_images(self, request): - """Get recently generated images.""" + """Get recently generated images (bounded window, newest first).""" try: try: limit, offset = parse_page_params(request.query) except ValueError: return bad_request("limit and offset must be integers") - # The DB query has no offset; fetch one bounded window and slice. + # The DB query has no offset: fetch one bounded window and slice, + # so the offset is capped to keep the window small. + offset = min(offset, MAX_RECENT_IMAGES_WINDOW - limit) rows = await self._run_in_executor( self.db.get_recent_images, limit + offset ) @@ -277,16 +393,39 @@ class ImageRoutesMixin: ) async def get_all_images(self, request): - """Get all generated images with linked prompts.""" + """Get generated images with linked prompts (bulk-capped pages).""" try: - images = await self._run_in_executor(self.db.get_all_images) + try: + limit, offset = parse_page_params( + request.query, + default_limit=MAX_BULK_LIMIT, + max_limit=MAX_BULK_LIMIT, + ) + except ValueError: + return bad_request("limit and offset must be integers") + + images = await self._run_in_executor( + self.db.get_all_images, limit=limit, offset=offset + ) return web.json_response( - {"success": True, "images": images, "count": len(images)} + { + "success": True, + "images": images, + "count": len(images), + "pagination": { + "limit": limit, + "offset": offset, + "count": len(images), + "has_more": len(images) == limit, + }, + } ) except Exception as e: self.logger.error(f"Get all images error: {e}") - return web.json_response({"success": False, "error": str(e)}, status=500) + return web.json_response( + {"success": False, "error": _safe_error(e)}, status=500 + ) async def search_images(self, request): """Search images by prompt text.""" @@ -309,29 +448,19 @@ class ImageRoutesMixin: def _scan_gallery_files_sync(self, output_path): """Scan output directory for media files (blocking I/O, run in executor). - Returns list of (path, mtime) tuples sorted by mtime descending. + Returns list of (path, mtime) tuples sorted by mtime descending; the + walk is bounded by MAX_SCAN_DEPTH and MAX_GALLERY_FILES. """ - image_extensions = [".png", ".jpg", ".jpeg", ".webp", ".gif"] - video_extensions = [".mp4", ".webm", ".avi", ".mov", ".mkv", ".m4v", ".wmv"] - media_extensions = image_extensions + video_extensions - all_images = [] - - seen_paths = set() - for ext in media_extensions: - for pattern in [f"*{ext}", f"*{ext.upper()}"]: - for media_path in output_path.rglob(pattern): - if "thumbnails" not in media_path.parts: - normalized_path = str(media_path).lower() - if normalized_path not in seen_paths: - seen_paths.add(normalized_path) - try: - mtime = media_path.stat().st_mtime - all_images.append((media_path, mtime)) - except OSError: - continue - - all_images.sort(key=lambda x: x[1], reverse=True) - return all_images + found = [] + for media_path in _iter_media_files(output_path, MAX_GALLERY_FILES): + try: + found.append((media_path, media_path.stat().st_mtime)) + except OSError: + continue + if len(found) >= MAX_GALLERY_FILES: + self.logger.warning(f"Gallery scan capped at {MAX_GALLERY_FILES} files") + found.sort(key=lambda x: x[1], reverse=True) + return found async def _get_gallery_files(self, output_path): """Get gallery files with per-directory TTL cache.""" @@ -348,8 +477,6 @@ class ImageRoutesMixin: async def get_output_images(self, request): """Get all images from ComfyUI output folder(s).""" try: - from urllib.parse import quote - output_dirs = self._get_all_output_dirs() if not output_dirs: return web.json_response( @@ -359,78 +486,16 @@ class ImageRoutesMixin: "images": [], }, ) - - limit = int(request.query.get("limit", 100)) - offset = int(request.query.get("offset", 0)) + try: + limit, offset = parse_page_params(request.query, default_limit=100) + except ValueError: + return bad_request("limit and offset must be integers") subfolder = request.query.get("subfolder", "").strip() - # Collect (path, mtime, root, root_index) from all directories - all_images = [] - for root_idx, output_path in enumerate(output_dirs): - dir_images = await self._get_gallery_files(output_path) - for img_path, mtime in dir_images: - all_images.append((img_path, mtime, output_path, root_idx)) - - # Sort combined results by mtime (newest first) — no .stat() calls - all_images.sort(key=lambda x: x[1], reverse=True) - - # Apply subfolder filter if provided - if subfolder: - filtered = [] - for img_path, mtime, root, ridx in all_images: - rel_dir = str(img_path.relative_to(root).parent) - if rel_dir == subfolder or rel_dir.startswith(subfolder + os.sep): - filtered.append((img_path, mtime, root, ridx)) - all_images = filtered - + all_images = await self._collect_output_images(output_dirs, subfolder) total = len(all_images) - paginated = all_images[offset : offset + limit] - - video_extensions = [".mp4", ".webm", ".avi", ".mov", ".mkv", ".m4v", ".wmv"] - - def _format_page(): - images = [] - for media_path, mtime, output_path, root_index in paginated: - try: - stat = media_path.stat() - rel_path = media_path.relative_to(output_path) - extension = media_path.suffix.lower() - is_video = extension in video_extensions - media_type = "video" if is_video else "image" - - thumbnail_url = None - thumbnails_dir = output_path / "thumbnails" - if thumbnails_dir.exists(): - thumbnail_ext = ".jpg" if is_video else extension - rel_path_no_ext = rel_path.with_suffix("") - thumbnail_rel_path = f"thumbnails/{rel_path_no_ext.as_posix()}_thumb{thumbnail_ext}" - thumbnail_abs_path = output_path / thumbnail_rel_path - if thumbnail_abs_path.exists(): - thumbnail_url = f'/prompt_manager/images/serve/{quote(thumbnail_rel_path, safe="/")}' - - images.append( - { - "id": str(hash(str(media_path))), - "filename": media_path.name, - "path": str(media_path), - "relative_path": str(rel_path), - "root_dir": str(output_path), - "root_index": root_index, - "url": f"/prompt_manager/images/serve/{rel_path.as_posix()}", - "thumbnail_url": thumbnail_url, - "size": stat.st_size, - "modified_time": stat.st_mtime, - "extension": extension, - "media_type": media_type, - "is_video": is_video, - } - ) - except Exception as e: - self.logger.error(f"Error processing media {media_path}: {e}") - continue - return images - - images = await self._run_in_executor(_format_page) + page = all_images[offset : offset + limit] + images = await self._run_in_executor(self._format_output_page, page) return web.json_response( { @@ -444,7 +509,9 @@ class ImageRoutesMixin: ) except Exception as e: self.logger.error(f"Output images error: {e}") - return web.json_response({"success": False, "error": str(e)}, status=500) + return web.json_response( + {"success": False, "error": _safe_error(e)}, status=500 + ) async def serve_image(self, request): """Serve the actual image file using streamed FileResponse.""" @@ -590,12 +657,17 @@ class ImageRoutesMixin: @classmethod def _thumbnail_size(cls, quality): + """Map a quality preset name to a pixel size; anything else is medium.""" + if not isinstance(quality, str): + return cls.THUMBNAIL_SIZES["medium"] return cls.THUMBNAIL_SIZES.get(quality, cls.THUMBNAIL_SIZES["medium"]) async def generate_thumbnails(self, request): """Generate every thumbnail in one executor job (POST, blocking).""" try: - data = await request.json() + data, error = await _json_body_or_error(request) + if error is not None: + return error thumbnail_size = self._thumbnail_size(data.get("quality", "medium")) output_dir = self._find_comfyui_output_dir() @@ -639,38 +711,23 @@ class ImageRoutesMixin: def _thumbnail_targets(self, output_path, thumbnails_dir): """Scan *output_path* for media (blocking); returns (src, dst, is_video). - The walk never follows symlinks, stops at MAX_SCAN_DEPTH and returns - at most MAX_THUMBNAIL_FILES entries per request. + Bounded by _iter_media_files (no symlinks, MAX_SCAN_DEPTH) and + MAX_THUMBNAIL_FILES per request. """ output_path = Path(output_path) - root_depth = len(output_path.parts) targets = [] - for root, dirs, files in os.walk(output_path): - root_path = Path(root) - rel_parts = root_path.parts[root_depth:] - if "thumbnails" in rel_parts: - dirs[:] = [] + for src in _iter_media_files(output_path, MAX_THUMBNAIL_FILES): + suffix = src.suffix.lower() + is_video = suffix in VIDEO_EXTENSIONS + rel_no_ext = src.relative_to(output_path).with_suffix("") + thumb_suffix = ".jpg" if is_video else suffix + dst = thumbnails_dir / f"{rel_no_ext.as_posix()}_thumb{thumb_suffix}" + if not _is_within(dst, thumbnails_dir): + self.logger.warning(f"Skipping thumbnail outside safe dir: {src.name}") continue - if len(rel_parts) >= MAX_SCAN_DEPTH: - dirs[:] = [] - for name in sorted(files): - src = root_path / name - suffix = src.suffix.lower() - if suffix not in MEDIA_EXTENSIONS: - continue - is_video = suffix in VIDEO_EXTENSIONS - rel_no_ext = src.relative_to(output_path).with_suffix("") - thumb_suffix = ".jpg" if is_video else suffix - dst = thumbnails_dir / f"{rel_no_ext.as_posix()}_thumb{thumb_suffix}" - if not _is_within(dst, thumbnails_dir): - self.logger.warning(f"Skipping thumbnail outside safe dir: {name}") - continue - if len(targets) >= MAX_THUMBNAIL_FILES: - self.logger.warning( - f"Thumbnail scan capped at {MAX_THUMBNAIL_FILES} files" - ) - return targets - targets.append((src, dst, is_video)) + targets.append((src, dst, is_video)) + if len(targets) >= MAX_THUMBNAIL_FILES: + self.logger.warning(f"Thumbnail scan capped at {MAX_THUMBNAIL_FILES}") return targets def _generate_one(self, src, dst, thumbnail_size, is_video): @@ -1025,23 +1082,23 @@ class ImageRoutesMixin: return web.json_response({"success": False, "error": str(e)}, status=500) async def link_image_to_prompt(self, request): - """Link a generated image to a prompt.""" + """Link a generated image (inside an allowed directory) to a prompt.""" try: - data = await request.json() + data, error = await _json_body_or_error(request) + if error is not None: + return error prompt_id = data.get("prompt_id") image_path = data.get("image_path") metadata = data.get("metadata", {}) - if not prompt_id or not image_path: - return web.json_response( - { - "success": False, - "error": "prompt_id and image_path are required", - }, - status=400, - ) + if not prompt_id or not isinstance(image_path, str) or not image_path: + return bad_request("prompt_id and image_path are required") - if not os.path.exists(image_path): + denied = _media_access_error(image_path, self._get_all_output_dirs()) + if denied is not None: + return denied + + if not os.path.isfile(image_path): return web.json_response( {"success": False, "error": "Image file not found"}, status=404 ) @@ -1060,7 +1117,9 @@ class ImageRoutesMixin: except Exception as e: self.logger.error(f"Link image error: {e}") - return web.json_response({"success": False, "error": str(e)}, status=500) + return web.json_response( + {"success": False, "error": _safe_error(e)}, status=500 + ) async def get_image_prompt(self, request): """Get prompt information for a specific image path.""" @@ -1130,3 +1189,27 @@ class ImageRoutesMixin: except Exception as e: self.logger.error(f"Delete image error: {e}") return web.json_response({"success": False, "error": str(e)}, status=500) + + async def _collect_output_images(self, output_dirs, subfolder): + """Merge cached per-root listings, newest first, optionally filtered.""" + merged = [] + for root_idx, output_path in enumerate(output_dirs): + for img_path, mtime in await self._get_gallery_files(output_path): + merged.append((img_path, mtime, output_path, root_idx)) + merged.sort(key=lambda x: x[1], reverse=True) + if not subfolder: + return merged + return [ + entry + for entry in merged + if _in_subfolder(str(entry[0].relative_to(entry[2]).parent), subfolder) + ] + + def _format_output_page(self, page): + entries = [] + for media_path, _mtime, output_path, root_index in page: + try: + entries.append(_output_image_entry(media_path, output_path, root_index)) + except Exception as e: + self.logger.error(f"Error processing media {media_path.name}: {e}") + return entries diff --git a/py/api/prompts.py b/py/api/prompts.py index fd908da..dd75ffb 100644 --- a/py/api/prompts.py +++ b/py/api/prompts.py @@ -6,19 +6,22 @@ import os from aiohttp import web -# Upper bound on page sizes for every list endpoint. +# Upper bounds on page sizes and offsets for every list endpoint. The offset +# ceiling keeps values inside SQLite's 64-bit range and bounds skip work. MAX_PAGE_LIMIT = 500 +MAX_PAGE_OFFSET = 10_000_000 -def parse_page_params(query, default_limit=50): +def parse_page_params(query, default_limit=50, max_limit=MAX_PAGE_LIMIT): """Return ``(limit, offset)`` from a query mapping, clamped to safe bounds. - ``limit`` is clamped to ``[1, MAX_PAGE_LIMIT]`` and ``offset`` to ``>= 0``. - Non-integer values raise ``ValueError`` so the caller can answer 400. + ``limit`` is clamped to ``[1, max_limit]`` and ``offset`` to + ``[0, MAX_PAGE_OFFSET]``. Non-integer values raise ``ValueError`` so the + caller can answer 400. """ limit = int(query.get("limit", default_limit)) offset = int(query.get("offset", 0)) - return max(1, min(limit, MAX_PAGE_LIMIT)), max(0, offset) + return max(1, min(limit, max_limit)), max(0, min(offset, MAX_PAGE_OFFSET)) def bad_request(message): diff --git a/tests/test_api_images.py b/tests/test_api_images.py index b274976..d84bca4 100644 --- a/tests/test_api_images.py +++ b/tests/test_api_images.py @@ -11,7 +11,7 @@ import sys import tempfile import unittest from pathlib import Path -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) @@ -24,6 +24,7 @@ from aiohttp.test_utils import AioHTTPTestCase # noqa: E402 from PIL import Image # noqa: E402 from database.operations import PromptDatabase # noqa: E402 +import py.api.images as images_module # noqa: E402 from py.api import PromptManagerAPI # noqa: E402 from py.api.prompts import MAX_PAGE_LIMIT # noqa: E402 from utils.hashing import generate_prompt_hash # noqa: E402 @@ -154,7 +155,7 @@ class TestServeOutputImage(ImageAPITestCase): encoded = str(outside).replace(os.sep, "%2f").replace("/", "%2f") resp = await self.client.request( - f"GET", f"/prompt_manager/images/serve/{encoded}" + "GET", f"/prompt_manager/images/serve/{encoded}" ) self.assertEqual(resp.status, 403) @@ -494,5 +495,369 @@ class TestGenerateThumbnailsPost(ThumbnailTestCase): self.assertEqual(data["total_images"], 0) +class TestRecentImagesCaps(ImageAPITestCase): + """Every path through /images/recent hands the DB a bounded window.""" + + def _spy_db(self): + seen = [] + original = self.api.db.get_recent_images + + def spy(limit): + seen.append(limit) + return original(limit) + + self.api.db.get_recent_images = spy + return seen + + async def test_absent_params_use_bounded_default(self): + seen = self._spy_db() + + resp = await self.client.request("GET", "/prompt_manager/images/recent") + + data = await resp.json() + self.assertEqual(data["pagination"]["limit"], 50) + self.assertEqual(seen, [50]) + + async def test_huge_offset_is_clamped_before_reaching_db(self): + seen = self._spy_db() + + resp = await self.client.request( + "GET", "/prompt_manager/images/recent?limit=10&offset=999999999999" + ) + + self.assertEqual(resp.status, 200) + data = await resp.json() + self.assertLessEqual(seen[0], images_module.MAX_RECENT_IMAGES_WINDOW) + self.assertEqual( + data["pagination"]["offset"], images_module.MAX_RECENT_IMAGES_WINDOW - 10 + ) + + async def test_float_string_limit_is_400(self): + resp = await self.client.request( + "GET", "/prompt_manager/images/recent?limit=1e9" + ) + + self.assertEqual(resp.status, 400) + + async def test_zero_limit_becomes_one(self): + seen = self._spy_db() + + await self.client.request("GET", "/prompt_manager/images/recent?limit=0") + + self.assertEqual(seen, [1]) + + +class TestAllImagesCaps(ImageAPITestCase): + """/images/all never asks the DB for an unbounded result.""" + + def _spy_db(self): + seen = [] + original = self.api.db.get_all_images + + def spy(limit=0, offset=0): + seen.append((limit, offset)) + return original(limit=limit, offset=offset) + + self.api.db.get_all_images = spy + return seen + + async def test_absent_params_use_bulk_ceiling(self): + seen = self._spy_db() + + resp = await self.client.request("GET", "/prompt_manager/images/all") + + data = await resp.json() + self.assertTrue(data["success"]) + self.assertEqual(data["pagination"]["limit"], images_module.MAX_BULK_LIMIT) + self.assertEqual(seen, [(images_module.MAX_BULK_LIMIT, 0)]) + + async def test_limit_above_bulk_ceiling_is_clamped(self): + seen = self._spy_db() + + await self.client.request("GET", "/prompt_manager/images/all?limit=99999999") + + self.assertEqual(seen[0][0], images_module.MAX_BULK_LIMIT) + + async def test_negative_offset_is_clamped(self): + seen = self._spy_db() + + await self.client.request("GET", "/prompt_manager/images/all?offset=-9") + + self.assertEqual(seen[0][1], 0) + + async def test_non_integer_limit_is_400(self): + resp = await self.client.request("GET", "/prompt_manager/images/all?limit=abc") + + self.assertEqual(resp.status, 400) + + async def test_returns_linked_images(self): + self._link_image(make_png(self.output_dir / "one.png")) + + resp = await self.client.request("GET", "/prompt_manager/images/all?limit=1") + + data = await resp.json() + self.assertEqual(data["count"], 1) + self.assertEqual(data["pagination"]["count"], 1) + + +class TestOutputImagesCaps(ImageAPITestCase): + """GET /prompt_manager/images/output clamps paging on every path.""" + + URL = "/prompt_manager/images/output" + + async def test_absent_params_use_default_limit(self): + resp = await self.client.request("GET", self.URL) + + data = await resp.json() + self.assertTrue(data["success"]) + self.assertEqual(data["limit"], 100) + self.assertEqual(data["offset"], 0) + + async def test_limit_is_clamped_to_max_page_limit(self): + resp = await self.client.request("GET", f"{self.URL}?limit=999999") + + data = await resp.json() + self.assertEqual(data["limit"], MAX_PAGE_LIMIT) + + async def test_zero_limit_becomes_one(self): + resp = await self.client.request("GET", f"{self.URL}?limit=0") + + data = await resp.json() + self.assertEqual(data["limit"], 1) + + async def test_float_string_limit_is_400(self): + resp = await self.client.request("GET", f"{self.URL}?limit=1e9") + + self.assertEqual(resp.status, 400) + data = await resp.json() + self.assertFalse(data["success"]) + + async def test_negative_offset_is_clamped(self): + resp = await self.client.request("GET", f"{self.URL}?offset=-5") + + data = await resp.json() + self.assertEqual(data["offset"], 0) + + async def test_pages_and_filters_by_subfolder(self): + make_png(self.output_dir / "sub" / "a.png") + make_png(self.output_dir / "sub" / "b.png") + make_png(self.output_dir / "top.png") + make_png(self.output_dir / "thumbnails" / "top_thumb.png") + + page = await self.client.request("GET", f"{self.URL}?subfolder=sub&limit=1") + rest = await self.client.request("GET", f"{self.URL}?limit=10") + + data = await page.json() + self.assertEqual(data["total"], 2) + self.assertEqual(len(data["images"]), 1) + self.assertTrue(data["has_more"]) + self.assertEqual(data["images"][0]["root_index"], 0) + all_data = await rest.json() + names = sorted(i["filename"] for i in all_data["images"]) + self.assertEqual(names, ["a.png", "b.png", "top.png"]) + top = next(i for i in all_data["images"] if i["filename"] == "top.png") + self.assertIn("top_thumb.png", top["thumbnail_url"]) + + async def test_no_output_dirs_reports_failure(self): + self.api._get_all_output_dirs = lambda: [] + + resp = await self.client.request("GET", self.URL) + + data = await resp.json() + self.assertFalse(data["success"]) + + +class TestMediaScanCaps(ImageAPITestCase): + """Request-triggered filesystem walks are depth- and count-capped.""" + + def test_gallery_scan_honours_file_cap(self): + for name in ("a", "b", "c"): + make_png(self.output_dir / f"{name}.png") + + with patch.object(images_module, "MAX_GALLERY_FILES", 2): + found = self.api._scan_gallery_files_sync(self.output_dir) + + self.assertEqual(len(found), 2) + + def test_gallery_scan_honours_depth_cap(self): + make_png(self.output_dir / "a" / "shallow.png") + make_png(self.output_dir / "a" / "b" / "deep.png") + + with patch.object(images_module, "MAX_SCAN_DEPTH", 1): + found = [ + p.name for p, _ in self.api._scan_gallery_files_sync(self.output_dir) + ] + + self.assertEqual(found, ["shallow.png"]) + + def test_gallery_scan_skips_symlinked_directories(self): + with tempfile.TemporaryDirectory() as other: + make_png(Path(other) / "outside.png") + self._make_symlink(other, self.output_dir / "link") + + found = self.api._scan_gallery_files_sync(self.output_dir) + + self.assertEqual(found, []) + + def test_gallery_scan_skips_thumbnails_and_non_media(self): + make_png(self.output_dir / "keep.png") + make_png(self.output_dir / "thumbnails" / "keep_thumb.png") + (self.output_dir / "notes.txt").write_text("x") + + found = [p.name for p, _ in self.api._scan_gallery_files_sync(self.output_dir)] + + self.assertEqual(found, ["keep.png"]) + + def test_thumbnail_targets_honour_file_cap(self): + for name in ("a", "b", "c"): + make_png(self.output_dir / f"{name}.png") + + with patch.object(images_module, "MAX_THUMBNAIL_FILES", 2): + targets = self.api._thumbnail_targets( + self.output_dir, self.output_dir / "thumbnails" + ) + + self.assertEqual(len(targets), 2) + + def test_thumbnail_targets_honour_depth_cap(self): + make_png(self.output_dir / "a" / "shallow.png") + make_png(self.output_dir / "a" / "b" / "deep.png") + + with patch.object(images_module, "MAX_SCAN_DEPTH", 1): + targets = self.api._thumbnail_targets( + self.output_dir, self.output_dir / "thumbnails" + ) + + self.assertEqual([src.name for src, _, _ in targets], ["shallow.png"]) + + def test_thumbnail_targets_skip_symlinked_directories(self): + with tempfile.TemporaryDirectory() as other: + make_png(Path(other) / "outside.png") + self._make_symlink(other, self.output_dir / "link") + + targets = self.api._thumbnail_targets( + self.output_dir, self.output_dir / "thumbnails" + ) + + self.assertEqual(targets, []) + + def test_thumbnail_size_ignores_non_string_quality(self): + self.assertEqual(self.api._thumbnail_size({"w": 99999}), (300, 300)) + self.assertEqual(self.api._thumbnail_size(None), (300, 300)) + self.assertEqual(self.api._thumbnail_size("low"), (150, 150)) + + +class TestJsonBodyCaps(ImageAPITestCase): + """JSON bodies are read in chunks and rejected past MAX_JSON_BODY_BYTES.""" + + LINK_URL = "/prompt_manager/images/link" + THUMBS_URL = "/prompt_manager/images/generate-thumbnails" + + def _oversized(self, cap): + return b'{"image_path": "' + b"a" * (cap + 16) + b'"}' + + async def test_oversized_content_length_is_413(self): + with patch.object(images_module, "MAX_JSON_BODY_BYTES", 512): + resp = await self.client.request( + "POST", + self.LINK_URL, + data=self._oversized(512), + headers={"Content-Type": "application/json"}, + ) + + self.assertEqual(resp.status, 413) + data = await resp.json() + self.assertFalse(data["success"]) + + async def test_oversized_chunked_body_is_413(self): + with patch.object(images_module, "MAX_JSON_BODY_BYTES", 512): + resp = await self.client.request( + "POST", + self.THUMBS_URL, + data=self._oversized(512), + chunked=True, + headers={"Content-Type": "application/json"}, + ) + + self.assertEqual(resp.status, 413) + + async def test_invalid_json_is_400(self): + resp = await self.client.request( + "POST", + self.LINK_URL, + data=b"{not json", + headers={"Content-Type": "application/json"}, + ) + + self.assertEqual(resp.status, 400) + data = await resp.json() + self.assertFalse(data["success"]) + + async def test_non_object_json_is_400(self): + resp = await self.client.request("POST", self.THUMBS_URL, json=[1, 2, 3]) + + self.assertEqual(resp.status, 400) + + +class TestLinkImage(ImageAPITestCase): + """POST /prompt_manager/images/link only links media inside allowed dirs.""" + + URL = "/prompt_manager/images/link" + + async def test_links_png_inside_output_dir(self): + prompt_id = self._save_prompt("to link") + png = make_png(self.output_dir / "gen.png") + + resp = await self.client.request( + "POST", self.URL, json={"prompt_id": prompt_id, "image_path": str(png)} + ) + + self.assertEqual(resp.status, 200) + data = await resp.json() + self.assertTrue(data["success"]) + self.assertGreater(data["image_id"], 0) + + async def test_missing_fields_is_400(self): + resp = await self.client.request("POST", self.URL, json={"prompt_id": 1}) + + self.assertEqual(resp.status, 400) + + async def test_path_outside_allowed_dirs_is_403(self): + prompt_id = self._save_prompt("to link") + with tempfile.TemporaryDirectory() as other: + png = make_png(Path(other) / "outside.png") + + resp = await self.client.request( + "POST", self.URL, json={"prompt_id": prompt_id, "image_path": str(png)} + ) + + self.assertEqual(resp.status, 403) + + async def test_non_media_path_is_403(self): + prompt_id = self._save_prompt("to link") + txt = self.output_dir / "secret.txt" + txt.write_text("x") + + resp = await self.client.request( + "POST", self.URL, json={"prompt_id": prompt_id, "image_path": str(txt)} + ) + + self.assertEqual(resp.status, 403) + + async def test_missing_file_is_404(self): + prompt_id = self._save_prompt("to link") + + resp = await self.client.request( + "POST", + self.URL, + json={ + "prompt_id": prompt_id, + "image_path": str(self.output_dir / "no.png"), + }, + ) + + self.assertEqual(resp.status, 404) + + if __name__ == "__main__": unittest.main() diff --git a/tests/test_api_prompts.py b/tests/test_api_prompts.py index 9bb0879..db847e6 100644 --- a/tests/test_api_prompts.py +++ b/tests/test_api_prompts.py @@ -18,7 +18,11 @@ from aiohttp.test_utils import AioHTTPTestCase # noqa: E402 from database.operations import PromptDatabase # noqa: E402 from py.api import PromptManagerAPI # noqa: E402 -from py.api.prompts import MAX_PAGE_LIMIT, safe_error_message # noqa: E402 +from py.api.prompts import ( # noqa: E402 + MAX_PAGE_LIMIT, + MAX_PAGE_OFFSET, + safe_error_message, +) from utils.hashing import generate_prompt_hash # noqa: E402 @@ -119,6 +123,15 @@ class TestSearchBounds(PromptAPITestCase): data = await resp.json() self.assertEqual(data["pagination"]["offset"], 0) + async def test_offset_beyond_int64_is_clamped_not_500(self): + resp = await self.client.request( + "GET", "/prompt_manager/search?offset=99999999999999999999999" + ) + + self.assertEqual(resp.status, 200) + data = await resp.json() + self.assertEqual(data["pagination"]["offset"], MAX_PAGE_OFFSET) + async def test_offset_skips_results(self): for i in range(3): self._save_prompt(f"searchable {i}")