Vito Sansevero
b0e60ec4b8
Settings, diagnostics, duplicate/output scans and the logging endpoints returned absolute filesystem paths, and several error envelopes echoed str(e) which for OSError carries the full path. - _public_path() renders a path relative to the ComfyUI base directory (folder_paths.base_path, else the output dir's parent) or to the parent of an opted-in extra gallery root, falling back to the basename; _public_error() reduces OSError to strerror + basename. - GET /settings returns gallery roots and monitored directories in that form; POST /settings resolves relative roots against the same anchors (GalleryConfig.resolve_gallery_root) so the GET output round-trips. - Diagnostics output_dirs, monitored_directories and the missing-DB message no longer carry absolute paths; the image monitor is reached through the package-or-top-level import fallback instead of a three-dot relative import that failed outside the package. - Duplicate scan 'path' is now relative to the output root and the delete endpoint accepts relative paths. - /logs/files 'path' and /logs/stats 'log_directory' use the public form; unexpected errors answer with a generic message.
270 lines
10 KiB
Python
270 lines
10 KiB
Python
"""Tests for gallery root validation (GalleryConfig.validate_gallery_root).
|
|
|
|
A gallery root may only point inside one of ComfyUI's own directories
|
|
(output/input/temp/user, as reported by ``folder_paths``) or inside a
|
|
directory listed in ``PROMPT_MANAGER_EXTRA_GALLERY_ROOTS``. Anything else,
|
|
including filesystem roots and the home directory, is rejected so the
|
|
image-serving routes can never be pointed at arbitrary files.
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import tempfile
|
|
import types
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import MagicMock
|
|
|
|
_mock_server = MagicMock()
|
|
_mock_server.PromptServer.instance.routes = MagicMock()
|
|
sys.modules["server"] = _mock_server
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
from py.config import GalleryConfig # noqa: E402
|
|
|
|
EXTRA_ROOTS_ENV = "PROMPT_MANAGER_EXTRA_GALLERY_ROOTS"
|
|
|
|
|
|
class GalleryRootTestCase(unittest.TestCase):
|
|
"""Shared fixture: a fake ComfyUI tree exposed through a stub folder_paths."""
|
|
|
|
def setUp(self):
|
|
self.tmpdir = tempfile.mkdtemp()
|
|
self.addCleanup(self._rmtree, self.tmpdir)
|
|
|
|
self.comfy_dir = Path(self.tmpdir) / "ComfyUI"
|
|
self.output_dir = self.comfy_dir / "output"
|
|
self.input_dir = self.comfy_dir / "input"
|
|
self.output_dir.mkdir(parents=True)
|
|
self.input_dir.mkdir(parents=True)
|
|
|
|
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),
|
|
get_input_directory=lambda: str(self.input_dir),
|
|
get_temp_directory=self._raise_runtime_error,
|
|
# get_user_directory intentionally absent
|
|
)
|
|
self.addCleanup(self._restore_folder_paths)
|
|
|
|
self._orig_extra = os.environ.pop(EXTRA_ROOTS_ENV, None)
|
|
self.addCleanup(self._restore_extra_env)
|
|
|
|
@staticmethod
|
|
def _raise_runtime_error():
|
|
raise RuntimeError("temp directory unavailable")
|
|
|
|
@staticmethod
|
|
def _rmtree(path):
|
|
import shutil
|
|
|
|
shutil.rmtree(path, ignore_errors=True)
|
|
|
|
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
|
|
|
|
|
|
class TestValidateGalleryRootRejects(GalleryRootTestCase):
|
|
|
|
def test_filesystem_root_rejected(self):
|
|
ok, reason = GalleryConfig.validate_gallery_root(os.path.abspath(os.sep))
|
|
self.assertFalse(ok)
|
|
self.assertIn("root", reason.lower())
|
|
|
|
@unittest.skipUnless(os.name == "nt", "drive roots only exist on Windows")
|
|
def test_windows_drive_root_rejected(self):
|
|
ok, _ = GalleryConfig.validate_gallery_root("C:\\")
|
|
self.assertFalse(ok)
|
|
|
|
def test_home_directory_rejected(self):
|
|
ok, reason = GalleryConfig.validate_gallery_root(os.path.expanduser("~"))
|
|
self.assertFalse(ok)
|
|
self.assertIn("home", reason.lower())
|
|
|
|
def test_directory_outside_comfyui_rejected(self):
|
|
outside = Path(self.tmpdir) / "elsewhere"
|
|
outside.mkdir()
|
|
ok, _ = GalleryConfig.validate_gallery_root(str(outside))
|
|
self.assertFalse(ok)
|
|
|
|
def test_nonexistent_path_rejected(self):
|
|
missing = self.output_dir / "does_not_exist"
|
|
ok, reason = GalleryConfig.validate_gallery_root(str(missing))
|
|
self.assertFalse(ok)
|
|
self.assertIn("exist", reason.lower())
|
|
|
|
def test_file_rejected(self):
|
|
file_path = self.output_dir / "image.png"
|
|
file_path.write_bytes(b"PNG")
|
|
ok, reason = GalleryConfig.validate_gallery_root(str(file_path))
|
|
self.assertFalse(ok)
|
|
self.assertIn("directory", reason.lower())
|
|
|
|
def test_empty_and_non_string_rejected(self):
|
|
self.assertFalse(GalleryConfig.validate_gallery_root("")[0])
|
|
self.assertFalse(GalleryConfig.validate_gallery_root(" ")[0])
|
|
self.assertFalse(GalleryConfig.validate_gallery_root(None)[0])
|
|
self.assertFalse(GalleryConfig.validate_gallery_root(42)[0])
|
|
|
|
def test_sibling_with_shared_prefix_rejected(self):
|
|
"""'output2' must not pass because it string-starts with 'output'."""
|
|
sibling = self.comfy_dir / "output2"
|
|
sibling.mkdir()
|
|
ok, _ = GalleryConfig.validate_gallery_root(str(sibling))
|
|
self.assertFalse(ok)
|
|
|
|
def test_traversal_out_of_output_rejected(self):
|
|
outside = Path(self.tmpdir) / "elsewhere"
|
|
outside.mkdir()
|
|
sneaky = os.path.join(str(self.output_dir), "..", "..", "elsewhere")
|
|
ok, _ = GalleryConfig.validate_gallery_root(sneaky)
|
|
self.assertFalse(ok)
|
|
|
|
def test_symlink_escaping_output_rejected(self):
|
|
outside = Path(self.tmpdir) / "elsewhere"
|
|
outside.mkdir()
|
|
link = self.output_dir / "escape"
|
|
try:
|
|
os.symlink(str(outside), str(link), target_is_directory=True)
|
|
except (OSError, NotImplementedError):
|
|
self.skipTest("symlinks not available on this platform")
|
|
ok, _ = GalleryConfig.validate_gallery_root(str(link))
|
|
self.assertFalse(ok)
|
|
|
|
def test_env_root_that_is_filesystem_root_is_ignored(self):
|
|
"""Listing '/' in the env var must not re-open the whole disk."""
|
|
os.environ[EXTRA_ROOTS_ENV] = os.path.abspath(os.sep)
|
|
outside = Path(self.tmpdir) / "elsewhere"
|
|
outside.mkdir()
|
|
ok, _ = GalleryConfig.validate_gallery_root(str(outside))
|
|
self.assertFalse(ok)
|
|
|
|
|
|
class TestValidateGalleryRootAccepts(GalleryRootTestCase):
|
|
|
|
def test_output_directory_itself_accepted(self):
|
|
ok, reason = GalleryConfig.validate_gallery_root(str(self.output_dir))
|
|
self.assertTrue(ok, reason)
|
|
|
|
def test_subdirectory_of_output_accepted(self):
|
|
sub = self.output_dir / "renders" / "2026"
|
|
sub.mkdir(parents=True)
|
|
ok, reason = GalleryConfig.validate_gallery_root(str(sub))
|
|
self.assertTrue(ok, reason)
|
|
|
|
def test_subdirectory_of_input_accepted(self):
|
|
sub = self.input_dir / "refs"
|
|
sub.mkdir()
|
|
ok, reason = GalleryConfig.validate_gallery_root(str(sub))
|
|
self.assertTrue(ok, reason)
|
|
|
|
def test_directory_listed_in_env_accepted(self):
|
|
extra_a = Path(self.tmpdir) / "extra_a"
|
|
extra_b = Path(self.tmpdir) / "extra_b"
|
|
extra_a.mkdir()
|
|
extra_b.mkdir()
|
|
os.environ[EXTRA_ROOTS_ENV] = os.pathsep.join([str(extra_a), str(extra_b)])
|
|
|
|
ok_a, reason_a = GalleryConfig.validate_gallery_root(str(extra_a))
|
|
ok_b, reason_b = GalleryConfig.validate_gallery_root(str(extra_b / "."))
|
|
self.assertTrue(ok_a, reason_a)
|
|
self.assertTrue(ok_b, reason_b)
|
|
|
|
def test_relative_path_resolves_against_comfyui_base_not_cwd(self):
|
|
sub = self.output_dir / "rel"
|
|
sub.mkdir()
|
|
orig_cwd = os.getcwd()
|
|
os.chdir(self.tmpdir)
|
|
self.addCleanup(os.chdir, orig_cwd)
|
|
|
|
ok, reason = GalleryConfig.validate_gallery_root("output/rel")
|
|
self.assertTrue(ok, reason)
|
|
self.assertEqual(
|
|
GalleryConfig.resolve_gallery_root("output/rel"),
|
|
os.path.normcase(os.path.realpath(sub)),
|
|
)
|
|
# A path that only exists relative to the CWD is not accepted.
|
|
(Path(self.tmpdir) / "rel").mkdir()
|
|
self.assertFalse(GalleryConfig.validate_gallery_root("rel")[0])
|
|
|
|
def test_accepts_without_folder_paths_when_env_lists_dir(self):
|
|
sys.modules.pop("folder_paths", None)
|
|
extra = Path(self.tmpdir) / "extra"
|
|
extra.mkdir()
|
|
os.environ[EXTRA_ROOTS_ENV] = str(extra)
|
|
ok, reason = GalleryConfig.validate_gallery_root(str(extra))
|
|
self.assertTrue(ok, reason)
|
|
|
|
def test_rejects_everything_without_folder_paths_and_env(self):
|
|
sys.modules.pop("folder_paths", None)
|
|
ok, _ = GalleryConfig.validate_gallery_root(str(self.output_dir))
|
|
self.assertFalse(ok)
|
|
|
|
|
|
class TestGetAllOutputDirsFiltersRoots(GalleryRootTestCase):
|
|
"""PromptManagerAPI._get_all_output_dirs() only returns validated roots."""
|
|
|
|
def setUp(self):
|
|
super().setUp()
|
|
from py.api import PromptManagerAPI
|
|
|
|
self.api = PromptManagerAPI()
|
|
self.api._cached_output_dir = None
|
|
self._orig_dirs = list(GalleryConfig.MONITORING_DIRECTORIES)
|
|
self.addCleanup(self._restore_dirs)
|
|
|
|
def _restore_dirs(self):
|
|
GalleryConfig.MONITORING_DIRECTORIES = self._orig_dirs
|
|
|
|
def _dirs(self):
|
|
return [os.path.normcase(str(p)) for p in self.api._get_all_output_dirs()]
|
|
|
|
def test_filesystem_root_is_never_served(self):
|
|
GalleryConfig.MONITORING_DIRECTORIES = [os.path.abspath(os.sep)]
|
|
dirs = self._dirs()
|
|
self.assertNotIn(os.path.normcase(os.path.abspath(os.sep)), dirs)
|
|
self.assertEqual(dirs, [os.path.normcase(os.path.realpath(self.output_dir))])
|
|
|
|
def test_outside_directory_is_dropped_and_valid_one_kept(self):
|
|
outside = Path(self.tmpdir) / "elsewhere"
|
|
outside.mkdir()
|
|
good = self.output_dir / "good"
|
|
good.mkdir()
|
|
GalleryConfig.MONITORING_DIRECTORIES = [str(outside), str(good)]
|
|
self.assertEqual(self._dirs(), [os.path.normcase(os.path.realpath(good))])
|
|
|
|
def test_all_rejected_falls_back_to_comfyui_output_only(self):
|
|
outside = Path(self.tmpdir) / "elsewhere"
|
|
outside.mkdir()
|
|
GalleryConfig.MONITORING_DIRECTORIES = [str(outside)]
|
|
self.assertEqual(
|
|
self._dirs(), [os.path.normcase(os.path.realpath(self.output_dir))]
|
|
)
|
|
|
|
def test_no_configured_dirs_uses_auto_detection(self):
|
|
GalleryConfig.MONITORING_DIRECTORIES = []
|
|
self.api._find_comfyui_output_dir = lambda: str(self.output_dir)
|
|
self.assertEqual(self._dirs(), [os.path.normcase(str(self.output_dir))])
|
|
|
|
def test_nothing_available_returns_empty_list(self):
|
|
sys.modules.pop("folder_paths", None)
|
|
outside = Path(self.tmpdir) / "elsewhere"
|
|
outside.mkdir()
|
|
GalleryConfig.MONITORING_DIRECTORIES = [str(outside)]
|
|
self.assertEqual(self._dirs(), [])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|