diff --git a/py/api/__init__.py b/py/api/__init__.py index 6d831a3..bc1e715 100644 --- a/py/api/__init__.py +++ b/py/api/__init__.py @@ -535,18 +535,18 @@ class PromptManagerAPI( from ..config import GalleryConfig if GalleryConfig.MONITORING_DIRECTORIES: - configured_dir = Path(GalleryConfig.MONITORING_DIRECTORIES[0]).resolve() - if configured_dir.exists() and configured_dir.is_dir(): + configured = GalleryConfig.MONITORING_DIRECTORIES[0] + # Same check as the settings endpoint: a root hand-edited + # into config.json must not open a directory outside ComfyUI + ok, reason = GalleryConfig.validate_gallery_root(configured) + if ok: + configured_dir = Path(configured).resolve() self.logger.info( f"Using configured monitoring directory: {configured_dir}" ) self._cached_output_dir = str(configured_dir) return self._cached_output_dir - else: - self.logger.warning( - "Configured directory does not exist: " - f"{GalleryConfig.MONITORING_DIRECTORIES[0]}" - ) + self.logger.warning(f"Ignoring configured directory: {reason}") except ImportError: self.logger.debug( "GalleryConfig not available, skipping configured directory check" diff --git a/py/api/admin.py b/py/api/admin.py index 12e1162..bf508eb 100644 --- a/py/api/admin.py +++ b/py/api/admin.py @@ -31,26 +31,51 @@ IMAGE_EXTENSIONS = (".png", ".jpg", ".jpeg", ".gif", ".webp", ".bmp", ".tiff") VIDEO_EXTENSIONS = (".mp4", ".avi", ".mov", ".mkv", ".webm", ".m4v", ".wmv") SCAN_BATCH_SIZE = 50 # Files per executor call during the output scan MAX_SAFETY_BACKUPS = 10 # '.backup_*' files kept next to the live database +MAX_SCAN_FILES = 50_000 # Upper bound on media files one scan will look at +MAX_SCAN_DEPTH = 12 # Directory levels below an output root that are scanned +# Files the duplicate cleaner may delete: everything the scan reports, plus +# the common single-f TIFF spelling. Anything else inside the output +# directory (databases, scripts, notes) is never a "duplicate image". +DELETABLE_MEDIA_EXTENSIONS = frozenset(IMAGE_EXTENSIONS + VIDEO_EXTENSIONS + (".tif",)) def _collect_media_files(output_dirs, extensions): - """All media files under ``output_dirs`` (blocking), skipping thumbnails. + """Media files under ``output_dirs`` (blocking), bounded and symlink-free. - Matches extensions case-insensitively and de-duplicates by - case-folded path so case-insensitive filesystems do not list a file twice. + Walks each root without following symlinks, skips ``thumbnails`` + directories and symlinked files, never descends more than MAX_SCAN_DEPTH + levels and stops after MAX_SCAN_FILES files in total. Extensions match + case-insensitively and each path is listed once. """ + wanted = {ext.lower() for ext in extensions} found = [] seen = set() for output_dir in output_dirs: - for ext in extensions: - for pattern in (f"*{ext}", f"*{ext.upper()}"): - for media_path in Path(output_dir).rglob(pattern): - if "thumbnails" in media_path.parts: - continue - key = os.path.normcase(str(media_path)) - if key not in seen: - seen.add(key) - found.append(media_path) + root = os.path.abspath(str(output_dir)) + if not os.path.isdir(root): + continue + for current, dirnames, filenames in os.walk(root, followlinks=False): + depth = 0 if current == root else len(Path(current).relative_to(root).parts) + dirnames[:] = sorted( + d + for d in dirnames + if d != "thumbnails" + and depth < MAX_SCAN_DEPTH + and not os.path.islink(os.path.join(current, d)) + ) + for name in sorted(filenames): + if os.path.splitext(name)[1].lower() not in wanted: + continue + full = os.path.join(current, name) + if os.path.islink(full) or not os.path.isfile(full): + continue + key = os.path.normcase(full) + if key in seen: + continue + seen.add(key) + found.append(Path(full)) + if len(found) >= MAX_SCAN_FILES: + return found return found @@ -521,6 +546,11 @@ class AdminRoutesMixin: ) return "outside output directory" + if file_path.suffix.lower() not in DELETABLE_MEDIA_EXTENSIONS: + return "not a media file" + if "thumbnails" in rel_path.parts: + return "thumbnails are managed by the gallery" + if not file_path.is_file(): return "file not found" diff --git a/py/config.py b/py/config.py index 0bf59d3..02f7024 100644 --- a/py/config.py +++ b/py/config.py @@ -249,7 +249,9 @@ class GalleryConfig: if "enabled" in monitoring: cls.MONITORING_ENABLED = monitoring["enabled"] if "directories" in monitoring: - cls.MONITORING_DIRECTORIES = monitoring["directories"] + directories = cls._validated_directories(monitoring["directories"]) + if directories is not None: + cls.MONITORING_DIRECTORIES = directories if "extensions" in monitoring: cls.SUPPORTED_EXTENSIONS = monitoring["extensions"] if "processing_delay" in monitoring: @@ -283,6 +285,30 @@ class GalleryConfig: if "metadata_extraction_timeout" in performance: cls.METADATA_EXTRACTION_TIMEOUT = performance["metadata_extraction_timeout"] + @classmethod + def _validated_directories(cls, directories: Any) -> Optional[List[str]]: + """Gallery roots from a config payload that pass validate_gallery_root. + + A hand-edited config.json is as untrusted as a settings request: + entries that fail validation are dropped with a warning (naming only + the basename), and a value that is not a list is ignored entirely + (``None`` is returned so the current roots stay as they are). + """ + if not isinstance(directories, list): + config_logger.warning( + "Ignoring gallery monitoring directories: expected a list" + ) + return None + kept = [] + for entry in directories: + ok, reason = cls.validate_gallery_root(entry) + if ok: + kept.append(entry) + continue + label = os.path.basename(entry) if isinstance(entry, str) else entry + config_logger.warning(f"Ignoring gallery root {label!r}: {reason}") + return kept + @classmethod def allowed_gallery_parents(cls) -> List[str]: """Canonical directories under which gallery roots may live. diff --git a/tests/test_api_admin.py b/tests/test_api_admin.py index 5e45f99..3bea626 100644 --- a/tests/test_api_admin.py +++ b/tests/test_api_admin.py @@ -561,6 +561,130 @@ class TestDeleteDuplicatesSyncBody(AdminAPITestCase): self.assertEqual(result["failed_count"], 1) self.assertEqual(result["deleted_count"], 0) + def test_refuses_non_media_files_inside_output(self): + for name in ("notes.txt", "prompts.db", "script.py", "noext"): + (self.output_dir / name).write_bytes(b"x") + + result = self.api._delete_duplicate_images_sync( + ["notes.txt", "prompts.db", "script.py", "noext"], str(self.output_dir) + ) + + self.assertEqual(result["deleted_count"], 0) + self.assertEqual(result["failed_count"], 4) + for name in ("notes.txt", "prompts.db", "script.py", "noext"): + self.assertTrue((self.output_dir / name).exists(), name) + self.assertTrue(all("not a media file" in f for f in result["failed_files"])) + + def test_refuses_files_under_thumbnails(self): + thumb = self.output_dir / "thumbnails" / "sub" / "x_thumb.png" + thumb.parent.mkdir(parents=True) + thumb.write_bytes(b"t") + + result = self.api._delete_duplicate_images_sync( + [os.path.join("thumbnails", "sub", "x_thumb.png")], str(self.output_dir) + ) + + self.assertEqual(result["deleted_count"], 0) + self.assertTrue(thumb.exists()) + self.assertIn("thumbnail", result["failed_files"][0]) + + def test_every_scanned_media_extension_is_deletable(self): + from py.api.admin import ( + DELETABLE_MEDIA_EXTENSIONS, + IMAGE_EXTENSIONS, + VIDEO_EXTENSIONS, + ) + + for ext in IMAGE_EXTENSIONS + VIDEO_EXTENSIONS + (".tif",): + self.assertIn(ext, DELETABLE_MEDIA_EXTENSIONS) + for ext in (".txt", ".db", ".py", ".json", ".html", ""): + self.assertNotIn(ext, DELETABLE_MEDIA_EXTENSIONS) + (self.output_dir / "a.WEBM").write_bytes(b"v") + result = self.api._delete_duplicate_images_sync( + ["a.WEBM"], str(self.output_dir) + ) + self.assertEqual(result["deleted_count"], 1) + + +class TestCollectMediaFiles(AdminAPITestCase): + """_collect_media_files: bounded walk that never follows symlinked files.""" + + def _collect(self): + from py.api.admin import ( + IMAGE_EXTENSIONS, + VIDEO_EXTENSIONS, + _collect_media_files, + ) + + found = _collect_media_files( + [self.output_dir], IMAGE_EXTENSIONS + VIDEO_EXTENSIONS + ) + return sorted(p.relative_to(self.output_dir).as_posix() for p in found) + + def test_matches_case_insensitively_and_skips_thumbnails(self): + (self.output_dir / "a.png").write_bytes(b"x") + (self.output_dir / "b.PNG").write_bytes(b"x") + (self.output_dir / "c.Mp4").write_bytes(b"x") + (self.output_dir / "d.txt").write_bytes(b"x") + (self.output_dir / "sub").mkdir() + (self.output_dir / "sub" / "e.jpg").write_bytes(b"x") + (self.output_dir / "thumbnails").mkdir() + (self.output_dir / "thumbnails" / "a_thumb.png").write_bytes(b"x") + (self.output_dir / "sub" / "thumbnails").mkdir() + (self.output_dir / "sub" / "thumbnails" / "e_thumb.jpg").write_bytes(b"x") + self.assertEqual(self._collect(), ["a.png", "b.PNG", "c.Mp4", "sub/e.jpg"]) + + def test_symlinked_files_and_directories_are_skipped(self): + secret = Path(self.tmpdir) / "secret.png" + secret.write_bytes(b"s") + secret_dir = Path(self.tmpdir) / "secret_dir" + secret_dir.mkdir() + (secret_dir / "inner.png").write_bytes(b"s") + (self.output_dir / "real.png").write_bytes(b"r") + try: + os.symlink(secret, self.output_dir / "link.png") + os.symlink(secret_dir, self.output_dir / "linkdir") + except (OSError, NotImplementedError): + self.skipTest("symlinks not available") + self.assertEqual(self._collect(), ["real.png"]) + + def test_depth_is_capped(self): + from py.api.admin import MAX_SCAN_DEPTH + + self.assertEqual(MAX_SCAN_DEPTH, 12) + deep = self.output_dir + for level in range(1, MAX_SCAN_DEPTH + 2): + deep = deep / f"d{level}" + deep.mkdir() + (deep / f"f{level}.png").write_bytes(b"x") + found = self._collect() + self.assertIn( + "/".join(f"d{i}" for i in range(1, MAX_SCAN_DEPTH + 1)) + + f"/f{MAX_SCAN_DEPTH}.png", + found, + ) + self.assertFalse( + any(name.endswith(f"f{MAX_SCAN_DEPTH + 1}.png") for name in found) + ) + self.assertEqual(len(found), MAX_SCAN_DEPTH) + + def test_file_count_is_capped(self): + import py.api.admin as admin_module + from unittest.mock import patch + + self.assertEqual(admin_module.MAX_SCAN_FILES, 50_000) + for i in range(6): + (self.output_dir / f"{i}.png").write_bytes(b"x") + with patch.object(admin_module, "MAX_SCAN_FILES", 4): + self.assertEqual(len(self._collect()), 4) + + def test_missing_root_is_skipped(self): + from py.api.admin import _collect_media_files + + self.assertEqual( + _collect_media_files([Path(self.tmpdir) / "nope"], (".png",)), [] + ) + class TestSettingsReadMonitorThroughModule(AdminAPITestCase): """get_settings reaches the image monitor through utils.image_monitor.""" diff --git a/tests/test_api_core.py b/tests/test_api_core.py index 31ceb33..86d4ca7 100644 --- a/tests/test_api_core.py +++ b/tests/test_api_core.py @@ -634,6 +634,16 @@ class TestInstanceHelpers(FolderPathsFixture): GalleryConfig.MONITORING_DIRECTORIES = [] self.assertEqual(self.api._find_comfyui_output_dir(), first) + def test_find_output_dir_skips_configured_root_outside_comfyui(self): + outside = Path(self.tmpdir) / "outside" + outside.mkdir() + GalleryConfig.MONITORING_DIRECTORIES = [str(outside)] + self.api._cached_output_dir = None + result = self.api._find_comfyui_output_dir() + self.assertNotEqual( + os.path.normcase(str(result or "")), os.path.normcase(str(outside)) + ) + def test_find_output_dir_skips_missing_configured(self): GalleryConfig.MONITORING_DIRECTORIES = [os.path.join(self.tmpdir, "nope")] self.api._cached_output_dir = None diff --git a/tests/test_gallery_roots.py b/tests/test_gallery_roots.py index 144372d..9d1dfef 100644 --- a/tests/test_gallery_roots.py +++ b/tests/test_gallery_roots.py @@ -7,7 +7,9 @@ including filesystem roots and the home directory, is rejected so the image-serving routes can never be pointed at arbitrary files. """ +import json import os +import shutil import sys import tempfile import types @@ -21,11 +23,89 @@ sys.modules.setdefault("server", _mock_server) sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) -from py.config import GalleryConfig # noqa: E402 +from py.config import GalleryConfig, PromptManagerConfig # noqa: E402 EXTRA_ROOTS_ENV = "PROMPT_MANAGER_EXTRA_GALLERY_ROOTS" +class TestUpdateConfigValidatesRoots(unittest.TestCase): + """Roots hand-edited into config.json go through validate_gallery_root.""" + + def setUp(self): + self.tmpdir = tempfile.mkdtemp() + self.addCleanup(shutil.rmtree, self.tmpdir, True) + self.comfy_dir = Path(self.tmpdir) / "ComfyUI" + self.output_dir = self.comfy_dir / "output" + self.output_dir.mkdir(parents=True) + self.outside = Path(self.tmpdir) / "outside" + self.outside.mkdir() + self._orig_folder_paths = sys.modules.get("folder_paths") + sys.modules["folder_paths"] = types.SimpleNamespace( + base_path=str(self.comfy_dir), + get_output_directory=lambda: str(self.output_dir), + ) + self.addCleanup(self._restore_folder_paths) + self._orig_dirs = list(GalleryConfig.MONITORING_DIRECTORIES) + self.addCleanup( + setattr, GalleryConfig, "MONITORING_DIRECTORIES", self._orig_dirs + ) + self._orig_extra = os.environ.pop(EXTRA_ROOTS_ENV, None) + self.addCleanup(self._restore_extra_env) + + def _restore_folder_paths(self): + if self._orig_folder_paths is None: + sys.modules.pop("folder_paths", None) + else: + sys.modules["folder_paths"] = self._orig_folder_paths + + def _restore_extra_env(self): + if self._orig_extra is None: + os.environ.pop(EXTRA_ROOTS_ENV, None) + else: + os.environ[EXTRA_ROOTS_ENV] = self._orig_extra + + def test_invalid_roots_are_dropped_with_a_warning(self): + good = self.output_dir / "renders" + good.mkdir() + with self.assertLogs("prompt_manager.config", level="WARNING") as logs: + GalleryConfig.update_config( + { + "monitoring": { + "directories": [ + str(self.outside), + str(good), + os.path.abspath(os.sep), + 42, + ] + } + } + ) + self.assertEqual(GalleryConfig.MONITORING_DIRECTORIES, [str(good)]) + self.assertEqual(len(logs.output), 3) + self.assertNotIn(str(self.outside), "".join(logs.output)) + + def test_non_list_directories_are_ignored(self): + GalleryConfig.MONITORING_DIRECTORIES = [str(self.output_dir)] + GalleryConfig.update_config({"monitoring": {"directories": "/etc"}}) + self.assertEqual(GalleryConfig.MONITORING_DIRECTORIES, [str(self.output_dir)]) + + def test_load_from_file_drops_invalid_roots(self): + config_path = os.path.join(self.tmpdir, "config.json") + with open(config_path, "w") as f: + json.dump( + { + "gallery": { + "monitoring": { + "directories": [str(self.outside), str(self.output_dir)] + } + } + }, + f, + ) + PromptManagerConfig.load_from_file(config_path) + self.assertEqual(GalleryConfig.MONITORING_DIRECTORIES, [str(self.output_dir)]) + + class GalleryRootTestCase(unittest.TestCase): """Shared fixture: a fake ComfyUI tree exposed through a stub folder_paths.""" diff --git a/tests/test_image_monitor_singleton.py b/tests/test_image_monitor_singleton.py index 867f5cc..1867bc2 100644 --- a/tests/test_image_monitor_singleton.py +++ b/tests/test_image_monitor_singleton.py @@ -232,6 +232,38 @@ class TestStartStop(unittest.TestCase): observer_cls.assert_not_called() self.assertFalse(self.monitor.is_monitoring) + def test_directories_failing_validation_are_not_watched(self): + allowed = os.path.join(self.tmp.name, "allowed") + os.mkdir(allowed) + config = types.SimpleNamespace( + MONITORING_ENABLED=True, + MONITORING_DIRECTORIES=[self.tmp.name, allowed], + validate_gallery_root=lambda p: (p == allowed, "outside ComfyUI"), + ) + fake = FakeObserver() + with patch.object(im_mod, "Observer", return_value=fake): + with patch.object( + im_mod.ImageMonitor, "_gallery_config", return_value=config + ): + self.monitor.start_monitoring() + self.assertEqual([p for _, p, _ in fake.scheduled], [allowed]) + self.assertEqual(self.monitor.monitored_directories, [allowed]) + self.monitor.stop_monitoring() + + def test_explicit_directories_are_validated_too(self): + config = types.SimpleNamespace( + MONITORING_ENABLED=True, + MONITORING_DIRECTORIES=[], + validate_gallery_root=lambda p: (False, "outside ComfyUI"), + ) + with patch.object(im_mod, "Observer") as observer_cls: + with patch.object( + im_mod.ImageMonitor, "_gallery_config", return_value=config + ): + self.monitor.start_monitoring([self.tmp.name]) + observer_cls.assert_not_called() + self.assertFalse(self.monitor.is_monitoring) + def test_configured_directories_are_used_when_none_are_given(self): config = types.SimpleNamespace( MONITORING_ENABLED=True, MONITORING_DIRECTORIES=[self.tmp.name] diff --git a/utils/image_monitor.py b/utils/image_monitor.py index ae38baa..48d3324 100644 --- a/utils/image_monitor.py +++ b/utils/image_monitor.py @@ -573,10 +573,19 @@ class ImageMonitor: valid_directories = [] for output_dir in output_directories or []: - if os.path.isdir(output_dir): - valid_directories.append(output_dir) - else: + if not os.path.isdir(output_dir): self.logger.warning(f"Directory does not exist: {output_dir}") + continue + # Same containment rule as the settings endpoint, so a root + # hand-edited into config.json is never watched (and its files + # never linked) when it lies outside ComfyUI's directories. + validate = getattr(config, "validate_gallery_root", None) + if callable(validate): + ok, reason = validate(output_dir) + if not ok: + self.logger.warning(f"Not watching {output_dir}: {reason}") + continue + valid_directories.append(output_dir) if not valid_directories: self.logger.warning("No valid directories to monitor") return