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.
This commit is contained in:
Vito Sansevero
+224
-141
@@ -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
|
||||
|
||||
+8
-5
@@ -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):
|
||||
|
||||
+367
-2
@@ -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()
|
||||
|
||||
@@ -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}")
|
||||
|
||||
Reference in New Issue
Block a user