From f9da6488d69fc301f83382d18e930cb17a4df6d0 Mon Sep 17 00:00:00 2001 From: Vito Sansevero Date: Wed, 7 Oct 2026 19:55:00 -0700 Subject: [PATCH] fix(autotag): contain batch rows, require media extensions, bound the scan - The batch tagger runs every generated_images row through _resolve_allowed_image_path; rows that exist but lie outside the served output directories are skipped and counted in skipped_outside_roots on progress and complete events, and the engine receives the resolved path. - _resolve_allowed_image_path requires a media extension (.png .jpg .jpeg .webp .gif .bmp .tiff .tif), so autotag/single returns 403 for other files. - autotag/start holds a module-level lock; a concurrent start gets 409 {"success": false, "error": "Auto-tagging is already running"}. - scan_output_dir walks with a 50,000-file and 12-level cap, skips symlinked files and thumbnails, returns path through _public_path, relative_path in POSIX form, and a truncated flag. - Autotag handlers report exceptions via the shared _public_error; the local copy is removed. Tests close the database before unlinking it. --- py/api/autotag_routes.py | 211 ++++++++++++++++++++++-------------- tests/test_api_autotag.py | 223 +++++++++++++++++++++++++++++++++++++- 2 files changed, 349 insertions(+), 85 deletions(-) diff --git a/py/api/autotag_routes.py b/py/api/autotag_routes.py index b9068c7..edd8a52 100644 --- a/py/api/autotag_routes.py +++ b/py/api/autotag_routes.py @@ -3,22 +3,33 @@ import asyncio import json import os +import threading from pathlib import Path from typing import Optional +from urllib.parse import quote from aiohttp import web _INTERNAL_ERROR = "An internal error occurred. Check server logs for details." +_BUSY_ERROR = "Auto-tagging is already running" + +# Only these may be handed to the tagging engine or listed by the scanner. +_MEDIA_EXTENSIONS = frozenset( + {".png", ".jpg", ".jpeg", ".webp", ".gif", ".bmp", ".tiff", ".tif"} +) +# Bounds for scan_output_dir: a runaway output tree must not pin the server. +_SCAN_MAX_FILES = 50_000 +_SCAN_MAX_DEPTH = 12 + +# One batch run at a time; a second autotag/start gets 409 while this is held. +_autotag_batch_lock = threading.Lock() -def _public_error(exc: Exception) -> str: - """Describe an exception without leaking absolute server paths.""" - if isinstance(exc, OSError) and exc.strerror: - filename = exc.filename - if filename: - return f"{exc.strerror}: {os.path.basename(str(filename))}" - return exc.strerror - return str(exc) +def _file_exists(path) -> bool: + try: + return Path(str(path)).exists() + except (OSError, ValueError): + return False def path_is_within(candidate, root) -> bool: @@ -91,7 +102,9 @@ class AutotagRoutesMixin: except Exception as e: self.logger.error(f"Get autotag models error: {e}") - return web.json_response({"success": False, "error": str(e)}, status=500) + return web.json_response( + {"success": False, "error": self._public_error(e)}, status=500 + ) async def _stream_sse(self, request, events): """Write an async iterator of dict events as server-sent events. @@ -250,6 +263,8 @@ class AutotagRoutesMixin: "processed": 0, "tagged": 0, "skipped": 0, + "errors": 0, + "skipped_outside_roots": 0, "status": "No images with linked prompts found", } service.unload_model() @@ -261,7 +276,13 @@ class AutotagRoutesMixin: "status": f"Found {total} images. Processing...", } - counts = {"processed": 0, "tagged": 0, "skipped": 0, "errors": 0} + counts = { + "processed": 0, + "tagged": 0, + "skipped": 0, + "errors": 0, + "skipped_outside_roots": 0, + } tagged_prompt_ids = set() last_update = _time.monotonic() finished = False @@ -274,6 +295,7 @@ class AutotagRoutesMixin: "processed": counts["processed"], "tagged": counts["tagged"], "skipped": counts["skipped"], + "skipped_outside_roots": counts["skipped_outside_roots"], } try: @@ -281,17 +303,26 @@ class AutotagRoutesMixin: image_path = image_data.get("image_path") prompt_id = image_data.get("prompt_id") skip_reason = None + allowed_path = None if not image_path or not prompt_id: skip_reason = "unlinked" elif prompt_id in tagged_prompt_ids: skip_reason = "prompt already processed" - elif not Path(image_path).exists(): - skip_reason = "file missing" - elif opts["skip_tagged"] and self._real_tags( - image_data.get("prompt_tags", []) - ): - tagged_prompt_ids.add(prompt_id) - skip_reason = "already tagged" + else: + # Rows can point anywhere (restored DB, hand edits): only + # media files inside the served output roots are tagged. + allowed_path = self._resolve_allowed_image_path(str(image_path)) + if allowed_path is None: + if _file_exists(image_path): + skip_reason = "outside allowed directories" + counts["skipped_outside_roots"] += 1 + else: + skip_reason = "file missing" + elif opts["skip_tagged"] and self._real_tags( + image_data.get("prompt_tags", []) + ): + tagged_prompt_ids.add(prompt_id) + skip_reason = "already tagged" if skip_reason is not None: counts["skipped"] += 1 @@ -307,7 +338,7 @@ class AutotagRoutesMixin: try: tags = await self._run_in_executor( service.generate_tags, - str(image_path), + allowed_path, general_threshold=opts["general_threshold"], character_threshold=opts["character_threshold"], ) @@ -385,10 +416,19 @@ class AutotagRoutesMixin: "message": _INTERNAL_ERROR, } - return await self._stream_sse(request, events()) + if not _autotag_batch_lock.acquire(blocking=False): + return web.json_response( + {"success": False, "error": _BUSY_ERROR}, status=409 + ) + try: + return await self._stream_sse(request, events()) + finally: + _autotag_batch_lock.release() def _resolve_allowed_image_path(self, raw_path: str) -> Optional[str]: - """Return the real path when it is a file inside an output directory.""" + """Return the real path when it is a media file inside an output directory.""" + if Path(str(raw_path)).suffix.lower() not in _MEDIA_EXTENSIONS: + return None for output_dir in self._get_all_output_dirs(): if path_is_within(raw_path, output_dir): real = os.path.realpath(raw_path) @@ -522,7 +562,7 @@ class AutotagRoutesMixin: except Exception as e: self.logger.error(f"AutoTag single error: {e}") return web.json_response( - {"success": False, "error": _public_error(e)}, status=500 + {"success": False, "error": self._public_error(e)}, status=500 ) async def apply_autotag(self, request): @@ -568,7 +608,9 @@ class AutotagRoutesMixin: except Exception as e: self.logger.error(f"Apply autotag error: {e}") - return web.json_response({"success": False, "error": str(e)}, status=500) + return web.json_response( + {"success": False, "error": self._public_error(e)}, status=500 + ) async def unload_autotag_model(self, request): """Manually unload the AutoTag model from memory.""" @@ -595,7 +637,37 @@ class AutotagRoutesMixin: except Exception as e: self.logger.error(f"Unload autotag model error: {e}") - return web.json_response({"success": False, "error": str(e)}, status=500) + return web.json_response( + {"success": False, "error": self._public_error(e)}, status=500 + ) + + @staticmethod + def _collect_scan_files(output_path: Path): + """Media files under ``output_path`` from a bounded, symlink-free walk. + + Descends at most ``_SCAN_MAX_DEPTH`` levels, skips ``thumbnails`` + trees and symlinked files, and stops after ``_SCAN_MAX_FILES`` files. + Returns ``(paths, truncated)``. + """ + found = [] + base_depth = len(output_path.parts) + for dirpath, dirnames, filenames in os.walk(output_path): + depth = len(Path(dirpath).parts) - base_depth + dirnames[:] = sorted( + name + for name in dirnames + if name != "thumbnails" and depth + 1 <= _SCAN_MAX_DEPTH + ) + for name in sorted(filenames): + if Path(name).suffix.lower() not in _MEDIA_EXTENSIONS: + continue + candidate = Path(dirpath) / name + if candidate.is_symlink(): + continue + if len(found) >= _SCAN_MAX_FILES: + return found, True + found.append(candidate) + return found, False async def scan_output_dir(self, request): """Scan ComfyUI output directory for images.""" @@ -608,73 +680,50 @@ class AutotagRoutesMixin: ) output_path = Path(output_dir) - image_extensions = [ - ".png", - ".jpg", - ".jpeg", - ".gif", - ".webp", - ".bmp", - ".tiff", - ] + files, truncated = await self._run_in_executor( + self._collect_scan_files, output_path + ) + has_thumbnails = (output_path / "thumbnails").is_dir() images = [] - seen_paths = set() - - for ext in image_extensions: - for pattern in [f"*{ext.lower()}", f"*{ext.upper()}"]: - for image_path in output_path.rglob(pattern): - if "thumbnails" not in image_path.parts: - normalized_path = str(image_path).lower() - if normalized_path not in seen_paths: - seen_paths.add(normalized_path) - - rel_path = image_path.relative_to(output_path) - - thumbnail_url = None - thumbnails_dir = output_path / "thumbnails" - if thumbnails_dir.exists(): - rel_path_no_ext = rel_path.with_suffix("") - thumbnail_rel_path = ( - "thumbnails/" - f"{rel_path_no_ext.as_posix()}" - f"_thumb{image_path.suffix}" - ) - thumbnail_abs_path = thumbnails_dir / ( - f"{rel_path_no_ext.as_posix()}" - f"_thumb{image_path.suffix}" - ) - if thumbnail_abs_path.exists(): - from urllib.parse import quote - - thumbnail_url = ( - "/prompt_manager/images/serve/" - f'{quote(thumbnail_rel_path, safe="/")}' - ) - - from urllib.parse import quote as url_quote - - images.append( - { - "filename": image_path.name, - "path": str(image_path), - "relative_path": str(rel_path), - "url": ( - "/prompt_manager/images/serve/" - + url_quote(rel_path.as_posix(), safe="/") - ), - "thumbnail_url": thumbnail_url, - } - ) + for image_path in files: + rel_path = image_path.relative_to(output_path) + thumbnail_url = None + if has_thumbnails: + thumb_rel = ( + f"thumbnails/{rel_path.with_suffix('').as_posix()}" + f"_thumb{image_path.suffix}" + ) + if (output_path / thumb_rel).exists(): + thumbnail_url = "/prompt_manager/images/serve/" + quote( + thumb_rel, safe="/" + ) + images.append( + { + "filename": image_path.name, + "path": self._public_path(image_path), + "relative_path": rel_path.as_posix(), + "url": "/prompt_manager/images/serve/" + + quote(rel_path.as_posix(), safe="/"), + "thumbnail_url": thumbnail_url, + } + ) images.sort(key=lambda x: x["filename"]) self.logger.info(f"Found {len(images)} images in output directory") return web.json_response( - {"success": True, "images": images, "count": len(images)} + { + "success": True, + "images": images, + "count": len(images), + "truncated": truncated, + } ) except Exception as e: self.logger.error(f"Scan output dir error: {e}") - return web.json_response({"success": False, "error": str(e)}, status=500) + return web.json_response( + {"success": False, "error": self._public_error(e)}, status=500 + ) diff --git a/tests/test_api_autotag.py b/tests/test_api_autotag.py index 00407d0..5c90f17 100644 --- a/tests/test_api_autotag.py +++ b/tests/test_api_autotag.py @@ -4,11 +4,13 @@ The ML engine is replaced by a fake service so no model is ever loaded and no network request is made. """ +import asyncio import json import os import shutil import sys import tempfile +import threading import types import unittest from pathlib import Path @@ -121,6 +123,7 @@ class AutotagAPITestCase(AioHTTPTestCase): async def tearDownAsync(self): self._service_patch.stop() + self.api.db.close_all() for path in ( self._temp_db.name, self._temp_db.name + "-wal", @@ -287,6 +290,41 @@ class TestAutotagSingleContainment(AutotagAPITestCase): self.assertNotIn(str(self.secret_dir), data["error"]) self.assertEqual(self.service.generate_calls, []) + async def test_non_media_extension_inside_output_dir_is_forbidden(self): + notes = self.output_dir / "notes.txt" + notes.write_text("not an image") + + resp = await self._post_json( + "/prompt_manager/autotag/single", {"path": str(notes)} + ) + + self.assertEqual(resp.status, 403) + self.assertFalse((await resp.json())["success"]) + self.assertEqual(self.service.generate_calls, []) + + async def test_non_media_extension_via_image_id_is_forbidden(self): + prompt_id = self._save_prompt() + notes = self.output_dir / "notes.txt" + notes.write_text("not an image") + image_id = self._link_image(prompt_id, notes) + + resp = await self._post_json( + "/prompt_manager/autotag/single", {"image_id": image_id} + ) + + self.assertEqual(resp.status, 403) + self.assertEqual(self.service.generate_calls, []) + + async def test_uppercase_media_extension_is_accepted(self): + image = self._make_image(self.output_dir, "shout.JPEG") + + resp = await self._post_json( + "/prompt_manager/autotag/single", {"path": str(image)} + ) + + self.assertEqual(resp.status, 200) + self.assertEqual(len(self.service.generate_calls), 1) + async def test_missing_selector_is_400(self): resp = await self._post_json("/prompt_manager/autotag/single", {}) self.assertEqual(resp.status, 400) @@ -424,6 +462,86 @@ class TestAutotagBatch(AutotagAPITestCase): self.assertTrue(all(0 <= e["progress"] <= 100 for e in progress)) self.assertTrue(any("processed" in e for e in progress)) + async def test_rows_outside_output_dirs_are_skipped_and_counted(self): + inside = self._save_prompt("inside") + outside = self._save_prompt("outside") + self._link_image(inside, self._make_image(self.output_dir, "in.png")) + secret = self._make_image(self.secret_dir, "password.png") + self._link_image(outside, secret) + + _, events = await self._start() + + done = events[-1] + self.assertEqual(done["type"], "complete") + self.assertEqual(done["processed"], 1) + self.assertEqual(done["tagged"], 1) + self.assertEqual(done["skipped"], 1) + self.assertEqual(done["skipped_outside_roots"], 1) + self.assertEqual(len(self.service.generate_calls), 1) + self.assertEqual( + os.path.realpath(self.service.generate_calls[0][0]), + os.path.realpath(str(self.output_dir / "in.png")), + ) + self.assertEqual(self.api.db.get_prompt_by_id(outside)["tags"], []) + progress = [e for e in events if e["type"] == "progress" and "skipped" in e] + self.assertTrue(progress) + self.assertTrue(all("skipped_outside_roots" in e for e in progress)) + body = json.dumps(events) + self.assertNotIn(str(self.secret_dir), body) + + async def test_second_start_while_running_is_409(self): + self._seed() + started = threading.Event() + release = threading.Event() + + def slow_generate(image_path, **kwargs): + started.set() + release.wait(5) + return ["slow"] + + self.service.generate_tags = slow_generate + first = asyncio.ensure_future(self._start()) + try: + await asyncio.get_event_loop().run_in_executor(None, started.wait, 5) + self.assertTrue(started.is_set()) + + resp = await self.client.request( + "POST", "/prompt_manager/autotag/start?model_type=gguf" + ) + + self.assertEqual(resp.status, 409) + self.assertEqual( + await resp.json(), + {"success": False, "error": "Auto-tagging is already running"}, + ) + finally: + release.set() + resp1, events = await first + self.assertEqual(resp1.status, 200) + self.assertEqual(events[-1]["type"], "complete") + + async def test_lock_is_released_after_a_run(self): + self._seed() + + resp1, _ = await self._start() + resp2, events = await self._start() + + self.assertEqual(resp1.status, 200) + self.assertEqual(resp2.status, 200) + self.assertEqual(events[-1]["type"], "complete") + + async def test_lock_is_released_when_the_run_fails(self): + self._seed() + self.service.load_error = RuntimeError("no model") + + _, events = await self._start() + self.assertEqual(events[-1]["type"], "error") + + self.service.load_error = None + resp, events = await self._start() + self.assertEqual(resp.status, 200) + self.assertEqual(events[-1]["type"], "complete") + async def test_already_tagged_prompts_are_skipped_unless_asked(self): p = self._save_prompt("tagged", tags=["cat"]) self._link_image(p, self._make_image(self.output_dir, "t.png")) @@ -603,10 +721,11 @@ class TestAutotagSingleEdgeCases(AutotagAPITestCase): self.assertEqual((await resp.json())["error"], "model crashed") async def test_oserror_without_filename_uses_strerror(self): - from py.api.autotag_routes import _public_error + from py.api import _public_error self.assertEqual(_public_error(OSError(5, "I/O error")), "I/O error") self.assertEqual(_public_error(ValueError("plain")), "plain") + self.assertIs(self.api._public_error, _public_error) async def test_path_with_nul_byte_is_forbidden(self): from py.api.autotag_routes import path_is_within @@ -759,6 +878,45 @@ class TestApplyAutotag(AutotagAPITestCase): self.assertEqual(resp.status, 500) +class TestAutotagErrorBodiesHideServerPaths(AutotagAPITestCase): + """Every autotag handler reports exceptions through the shared public helper.""" + + def _locked(self): + return PermissionError( + 13, "Permission denied", str(self.secret_dir / "locked.db") + ) + + async def _assert_hidden(self, resp): + self.assertEqual(resp.status, 500) + body = await resp.text() + self.assertFalse(json.loads(body)["success"]) + self.assertNotIn(str(self.secret_dir), body) + self.assertIn("locked.db", body) + + async def test_models_error(self): + def boom(): + raise self._locked() + + self.service.get_models_status = boom + resp = await self.client.request("GET", "/prompt_manager/autotag/models") + await self._assert_hidden(resp) + + async def test_apply_error(self): + with patch.object(self.api.db, "get_prompt_by_id", side_effect=self._locked()): + resp = await self._post_json( + "/prompt_manager/autotag/apply", {"prompt_id": 1, "tags": ["a"]} + ) + await self._assert_hidden(resp) + + async def test_unload_error(self): + def boom(): + raise self._locked() + + self.service.is_model_loaded = boom + resp = await self.client.request("POST", "/prompt_manager/autotag/unload") + await self._assert_hidden(resp) + + class TestUnloadAutotag(AutotagAPITestCase): async def test_nothing_loaded(self): @@ -799,10 +957,12 @@ class TestScanOutputDir(AutotagAPITestCase): resp = await self.client.request("GET", "/prompt_manager/scan_output_dir") - data = await resp.json() + body = await resp.text() + data = json.loads(body) self.assertTrue(data["success"]) self.assertEqual([i["filename"] for i in data["images"]], ["a.png", "b.JPG"]) self.assertEqual(data["count"], 2) + self.assertIs(data["truncated"], False) self.assertIsNone(data["images"][0]["thumbnail_url"]) self.assertEqual( data["images"][1]["thumbnail_url"], @@ -811,19 +971,74 @@ class TestScanOutputDir(AutotagAPITestCase): self.assertEqual( data["images"][1]["url"], "/prompt_manager/images/serve/sub/b.JPG" ) + self.assertEqual(data["images"][1]["relative_path"], "sub/b.JPG") + self.assertEqual( + data["images"][1]["path"], self.api._public_path(sub / "b.JPG") + ) + self.assertFalse(os.path.isabs(data["images"][1]["path"])) + self.assertNotIn(str(self.tmp_root), body) + + async def test_symlinked_files_are_skipped(self): + secret = self._make_image(self.secret_dir, "password.png") + self._make_image(self.output_dir, "real.png") + link = self.output_dir / "link.png" + try: + link.symlink_to(secret) + except (OSError, NotImplementedError): + self.skipTest("symlinks not supported on this platform") + self.api._find_comfyui_output_dir = lambda: str(self.output_dir) + + data = await ( + await self.client.request("GET", "/prompt_manager/scan_output_dir") + ).json() + + self.assertEqual([i["filename"] for i in data["images"]], ["real.png"]) + + async def test_file_cap_marks_the_listing_truncated(self): + for name in ("a.png", "b.png", "c.png"): + self._make_image(self.output_dir, name) + self.api._find_comfyui_output_dir = lambda: str(self.output_dir) + + with patch("py.api.autotag_routes._SCAN_MAX_FILES", 2): + data = await ( + await self.client.request("GET", "/prompt_manager/scan_output_dir") + ).json() + + self.assertEqual(data["count"], 2) + self.assertIs(data["truncated"], True) + + async def test_depth_cap_stops_the_walk(self): + deep = self.output_dir / "d1" / "d2" / "d3" + deep.mkdir(parents=True) + self._make_image(self.output_dir / "d1", "shallow.png") + self._make_image(deep, "deep.png") + self.api._find_comfyui_output_dir = lambda: str(self.output_dir) + + with patch("py.api.autotag_routes._SCAN_MAX_DEPTH", 1): + data = await ( + await self.client.request("GET", "/prompt_manager/scan_output_dir") + ).json() + + self.assertEqual([i["filename"] for i in data["images"]], ["shallow.png"]) async def test_missing_output_dir_is_404(self): self.api._find_comfyui_output_dir = lambda: None resp = await self.client.request("GET", "/prompt_manager/scan_output_dir") self.assertEqual(resp.status, 404) - async def test_error_is_500(self): + async def test_error_is_500_without_the_server_path(self): def boom(): - raise RuntimeError("x") + raise PermissionError( + 13, "Permission denied", str(self.secret_dir / "locked.png") + ) self.api._find_comfyui_output_dir = boom resp = await self.client.request("GET", "/prompt_manager/scan_output_dir") self.assertEqual(resp.status, 500) + body = await resp.text() + self.assertFalse(json.loads(body)["success"]) + self.assertNotIn(str(self.secret_dir), body) + self.assertIn("locked.png", body) if __name__ == "__main__":