From 672b4a255f3b081b2469012e087cf88279e2f5cb Mon Sep 17 00:00:00 2001 From: DrJKL Date: Thu, 11 Jun 2026 14:24:44 -0700 Subject: [PATCH] Simplify auto-managed frontend cleanup with on-disk markers Replace the .auto_managed.json sidecar with a .auto_managed/ marker directory. Tracked version names now come from real single-component dirents, removing the untrusted-input parsing, path-traversal checks, and rmtree boundary guards that the JSON design required. Also fix the dead Optional handling in _prune_auto_managed_versions, extract _ensure_release_downloaded to drop the download_succeeded flag, and reuse _provider_dir in init_frontend_unsafe. Amp-Thread-ID: https://ampcode.com/threads/T-019eb879-1e6f-77aa-abb3-4d229d18061f Co-authored-by: Amp --- app/frontend_management.py | 251 +++++++------------ tests-unit/app_test/frontend_manager_test.py | 84 ++----- 2 files changed, 123 insertions(+), 212 deletions(-) diff --git a/app/frontend_management.py b/app/frontend_management.py index fb2c639a8..7e456bcda 100644 --- a/app/frontend_management.py +++ b/app/frontend_management.py @@ -1,5 +1,4 @@ import argparse -import json import logging import os import re @@ -208,153 +207,95 @@ class FrontendManager: CUSTOM_FRONTENDS_ROOT = str(Path(__file__).parents[1] / "web_custom_versions") AUTO_MANAGED_VERSION_SPECIFIERS = ("latest", "prerelease") - AUTO_MANAGED_METADATA_FILENAME = ".auto_managed.json" - _VERSION_DIRNAME_PATTERN = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*$") + AUTO_MANAGED_MARKER_DIRNAME = ".auto_managed" @classmethod def _provider_dir(cls, repo_owner: str, repo_name: str) -> Path: return Path(cls.CUSTOM_FRONTENDS_ROOT) / f"{repo_owner}_{repo_name}" @classmethod - def _auto_managed_metadata_path(cls, repo_owner: str, repo_name: str) -> Path: - return cls._provider_dir(repo_owner, repo_name) / cls.AUTO_MANAGED_METADATA_FILENAME - - @classmethod - def _is_safe_version_dirname(cls, name: str) -> bool: - if not isinstance(name, str): - return False - if name in (".", "..") or "/" in name or "\\" in name or "\x00" in name: - return False - return bool(cls._VERSION_DIRNAME_PATTERN.match(name)) + def _auto_managed_marker_dir(cls, repo_owner: str, repo_name: str) -> Path: + return cls._provider_dir(repo_owner, repo_name) / cls.AUTO_MANAGED_MARKER_DIRNAME @classmethod def _read_auto_managed_versions(cls, repo_owner: str, repo_name: str) -> list[str]: - metadata_path = cls._auto_managed_metadata_path(repo_owner, repo_name) - if not metadata_path.exists(): - return [] - try: - with open(metadata_path, "r", encoding="utf-8") as fh: - data = json.load(fh) - except (OSError, ValueError) as exc: - logging.warning( - "Could not read frontend auto-managed metadata at %s: %s", - metadata_path, - exc, - ) - return [] - if not isinstance(data, dict): - logging.warning( - "Frontend auto-managed metadata at %s has unexpected shape; ignoring.", - metadata_path, - ) - return [] - versions = data.get("auto_managed", []) - if not isinstance(versions, list): - return [] - return [v for v in versions if cls._is_safe_version_dirname(v)] + """Return versions ComfyUI auto-downloaded for @latest / @prerelease. - @classmethod - def _write_auto_managed_versions( - cls, repo_owner: str, repo_name: str, versions: list[str] - ) -> None: - metadata_path = cls._auto_managed_metadata_path(repo_owner, repo_name) - metadata_path.parent.mkdir(parents=True, exist_ok=True) - safe_versions = [v for v in versions if cls._is_safe_version_dirname(v)] - payload = {"auto_managed": sorted(set(safe_versions))} - tmp_path = metadata_path.with_suffix(metadata_path.suffix + ".tmp") + Each tracked version is an empty marker file under the provider's + ``.auto_managed`` directory. Because the names come straight from real + single-component directory entries, there is no untrusted parsing and no + path-traversal surface to defend against. + """ + marker_dir = cls._auto_managed_marker_dir(repo_owner, repo_name) try: - with open(tmp_path, "w", encoding="utf-8") as fh: - json.dump(payload, fh, indent=2, sort_keys=True) - os.replace(tmp_path, metadata_path) + return sorted(entry.name for entry in marker_dir.iterdir() if entry.is_file()) + except FileNotFoundError: + return [] except OSError as exc: logging.warning( - "Could not write frontend auto-managed metadata at %s: %s", - metadata_path, + "Could not read frontend auto-managed markers at %s: %s", + marker_dir, + exc, + ) + return [] + + @classmethod + def _mark_auto_managed(cls, repo_owner: str, repo_name: str, version: str) -> None: + marker_dir = cls._auto_managed_marker_dir(repo_owner, repo_name) + try: + marker_dir.mkdir(parents=True, exist_ok=True) + (marker_dir / version).touch() + except OSError as exc: + logging.warning( + "Could not record auto-managed frontend version %s: %s", + version, exc, ) - if tmp_path.exists(): - try: - tmp_path.unlink() - except OSError: - pass @classmethod def _prune_auto_managed_versions( cls, repo_owner: str, repo_name: str, keep_version: str ) -> None: - tracked = cls._read_auto_managed_versions(repo_owner, repo_name) - if not tracked and keep_version is None: - return - + """Remove previously auto-downloaded versions other than ``keep_version``.""" provider_dir = cls._provider_dir(repo_owner, repo_name) - try: - provider_dir_resolved = provider_dir.resolve() - except OSError as exc: - logging.warning( - "Could not resolve provider directory %s for cleanup: %s", - provider_dir, - exc, - ) - return - - for stale_version in tracked: + for stale_version in cls._read_auto_managed_versions(repo_owner, repo_name): if stale_version == keep_version: continue - # Re-check even though read already filters: keeps rmtree provably - # bounded under provider_dir regardless of caller. - if not cls._is_safe_version_dirname(stale_version): - logging.warning( - "Refusing to clean up suspicious frontend version name: %r", - stale_version, - ) - continue + # stale_version is a single-component marker name, so this is always a + # direct child of provider_dir. stale_path = provider_dir / stale_version - if not stale_path.exists(): - continue - try: - stale_resolved = stale_path.resolve() - except OSError as exc: - logging.warning( - "Could not resolve stale frontend path %s: %s", - stale_path, - exc, - ) - continue - if ( - stale_resolved == provider_dir_resolved - or provider_dir_resolved not in stale_resolved.parents - ): - logging.warning( - "Refusing to remove path outside provider dir: %s (provider=%s)", - stale_resolved, - provider_dir_resolved, - ) - continue - try: - shutil.rmtree(stale_path) - logging.info( - "Removed stale auto-managed frontend version: %s", - stale_path, - ) - except OSError as exc: - logging.warning( - "Failed to remove stale frontend version at %s: %s", - stale_path, - exc, - ) - - new_tracked = [keep_version] if keep_version else [] - cls._write_auto_managed_versions(repo_owner, repo_name, new_tracked) + if stale_path.exists(): + try: + shutil.rmtree(stale_path) + logging.info( + "Removed stale auto-managed frontend version: %s", stale_path + ) + except OSError as exc: + logging.warning( + "Failed to remove stale frontend version at %s: %s", + stale_path, + exc, + ) + continue + cls._untrack_auto_managed_version(repo_owner, repo_name, stale_version) + cls._mark_auto_managed(repo_owner, repo_name, keep_version) @classmethod def _untrack_auto_managed_version( cls, repo_owner: str, repo_name: str, version: str ) -> None: - tracked = cls._read_auto_managed_versions(repo_owner, repo_name) - if version not in tracked: + """Stop auto-managing ``version`` (e.g. the user pinned it); keeps its files.""" + marker = cls._auto_managed_marker_dir(repo_owner, repo_name) / version + try: + marker.unlink() + except FileNotFoundError: return - tracked = [v for v in tracked if v != version] - cls._write_auto_managed_versions(repo_owner, repo_name, tracked) + except OSError as exc: + logging.warning( + "Could not untrack auto-managed frontend version %s: %s", + version, + exc, + ) @classmethod def get_required_frontend_version(cls) -> str: @@ -526,18 +467,13 @@ comfyui-workflow-templates is not installed. is_auto_managed = version in cls.AUTO_MANAGED_VERSION_SPECIFIERS if version.startswith("v"): - expected_path = str( - Path(cls.CUSTOM_FRONTENDS_ROOT) - / f"{repo_owner}_{repo_name}" - / version.lstrip("v") - ) + pinned_version = version.lstrip("v") + expected_path = str(cls._provider_dir(repo_owner, repo_name) / pinned_version) if os.path.exists(expected_path): logging.info( f"Using existing copy of specific frontend version tag: {repo_owner}/{repo_name}@{version}" ) - cls._untrack_auto_managed_version( - repo_owner, repo_name, version.lstrip("v") - ) + cls._untrack_auto_managed_version(repo_owner, repo_name, pinned_version) return expected_path logging.info( @@ -548,40 +484,47 @@ comfyui-workflow-templates is not installed. release = provider.get_release(version) semantic_version = release["tag_name"].lstrip("v") - web_root = str( - Path(cls.CUSTOM_FRONTENDS_ROOT) / provider.folder_name / semantic_version - ) - download_succeeded = os.path.exists(web_root) - if not download_succeeded: - try: - os.makedirs(web_root, exist_ok=True) - logging.info( - "Downloading frontend(%s) version(%s) to (%s)", - provider.folder_name, - semantic_version, - web_root, - ) - logging.debug(release) - download_release_asset_zip(release, destination_path=web_root) - download_succeeded = True - finally: - # Clean up the directory if it is empty, i.e. the download failed - if not os.listdir(web_root): - os.rmdir(web_root) - download_succeeded = False + web_root = str(cls._provider_dir(repo_owner, repo_name) / semantic_version) - if download_succeeded: + if cls._ensure_release_downloaded(provider, semantic_version, web_root, release): if is_auto_managed: - cls._prune_auto_managed_versions( - repo_owner, repo_name, semantic_version - ) + cls._prune_auto_managed_versions(repo_owner, repo_name, semantic_version) else: - cls._untrack_auto_managed_version( - repo_owner, repo_name, semantic_version - ) + cls._untrack_auto_managed_version(repo_owner, repo_name, semantic_version) return web_root + @classmethod + def _ensure_release_downloaded( + cls, + provider: "FrontEndProvider", + semantic_version: str, + web_root: str, + release: Release, + ) -> bool: + """Ensure ``release`` is present at ``web_root``. + + Returns True if the version is available on disk afterwards. A failed + download leaves no empty directory behind. + """ + if os.path.exists(web_root): + return True + try: + os.makedirs(web_root, exist_ok=True) + logging.info( + "Downloading frontend(%s) version(%s) to (%s)", + provider.folder_name, + semantic_version, + web_root, + ) + logging.debug(release) + download_release_asset_zip(release, destination_path=web_root) + finally: + # Clean up the directory if it is empty, i.e. the download failed + if not os.listdir(web_root): + os.rmdir(web_root) + return os.path.isdir(web_root) + @classmethod def init_frontend(cls, version_string: str) -> str: """ diff --git a/tests-unit/app_test/frontend_manager_test.py b/tests-unit/app_test/frontend_manager_test.py index aa2e35c3f..e681af976 100644 --- a/tests-unit/app_test/frontend_manager_test.py +++ b/tests-unit/app_test/frontend_manager_test.py @@ -309,8 +309,10 @@ def _make_version_dir(root, owner, repo, version): return folder -def test_auto_managed_metadata_roundtrip(custom_frontends_root): - FrontendManager._write_auto_managed_versions("o", "r", ["1.0.0", "1.1.0", "1.0.0"]) +def test_auto_managed_markers_roundtrip(custom_frontends_root): + FrontendManager._mark_auto_managed("o", "r", "1.0.0") + FrontendManager._mark_auto_managed("o", "r", "1.1.0") + FrontendManager._mark_auto_managed("o", "r", "1.0.0") # idempotent assert FrontendManager._read_auto_managed_versions("o", "r") == ["1.0.0", "1.1.0"] @@ -318,61 +320,26 @@ def test_read_auto_managed_versions_missing(custom_frontends_root): assert FrontendManager._read_auto_managed_versions("o", "r") == [] -def test_read_auto_managed_versions_corrupt(custom_frontends_root): +def test_mark_auto_managed_does_not_create_version_dir(custom_frontends_root): + FrontendManager._mark_auto_managed("o", "r", "1.0.0") + provider_dir = custom_frontends_root / "o_r" - provider_dir.mkdir() - (provider_dir / FrontendManager.AUTO_MANAGED_METADATA_FILENAME).write_text( - "not json" - ) - assert FrontendManager._read_auto_managed_versions("o", "r") == [] - - -@pytest.mark.parametrize("payload", ["[]", "null", '"oops"', "42"]) -def test_read_auto_managed_versions_non_dict_root(custom_frontends_root, payload): - provider_dir = custom_frontends_root / "o_r" - provider_dir.mkdir() - (provider_dir / FrontendManager.AUTO_MANAGED_METADATA_FILENAME).write_text(payload) - assert FrontendManager._read_auto_managed_versions("o", "r") == [] - - -def test_read_auto_managed_versions_drops_unsafe_names(custom_frontends_root): - provider_dir = custom_frontends_root / "o_r" - provider_dir.mkdir() - (provider_dir / FrontendManager.AUTO_MANAGED_METADATA_FILENAME).write_text( - '{"auto_managed": ["1.0.0", "../escape", "/etc", "..", ".", "ok-1.2"]}' - ) - assert FrontendManager._read_auto_managed_versions("o", "r") == ["1.0.0", "ok-1.2"] - - -def test_write_auto_managed_versions_drops_unsafe_names(custom_frontends_root): - FrontendManager._write_auto_managed_versions( - "o", "r", ["1.0.0", "../escape", "/etc/passwd", "..", "."] - ) + # The marker is recorded, but no version directory is created as a side effect. + assert not (provider_dir / "1.0.0").exists() assert FrontendManager._read_auto_managed_versions("o", "r") == ["1.0.0"] -def test_prune_refuses_to_delete_outside_provider_dir( - custom_frontends_root, monkeypatch -): - sibling = custom_frontends_root / "do-not-touch" - sibling.mkdir() - (sibling / "marker").write_text("keep me") +def test_prune_only_touches_marked_siblings(custom_frontends_root): + # A sibling provider directory must never be affected by pruning. + sibling = _make_version_dir(custom_frontends_root, "other", "repo", "9.9.9") - provider_dir = custom_frontends_root / "o_r" - provider_dir.mkdir() _make_version_dir(custom_frontends_root, "o", "r", "1.0.0") - - monkeypatch.setattr( - FrontendManager, - "_read_auto_managed_versions", - classmethod(lambda cls, owner, repo: ["1.0.0", "../do-not-touch"]), - ) + FrontendManager._mark_auto_managed("o", "r", "1.0.0") FrontendManager._prune_auto_managed_versions("o", "r", keep_version="2.0.0") assert sibling.exists() - assert (sibling / "marker").read_text() == "keep me" - assert not (provider_dir / "1.0.0").exists() + assert not (custom_frontends_root / "o_r" / "1.0.0").exists() def test_prune_auto_managed_versions_removes_stale_and_keeps_pinned( @@ -381,7 +348,8 @@ def test_prune_auto_managed_versions_removes_stale_and_keeps_pinned( _make_version_dir(custom_frontends_root, "o", "r", "1.0.0") _make_version_dir(custom_frontends_root, "o", "r", "1.1.0") pinned = _make_version_dir(custom_frontends_root, "o", "r", "1.2.0") - FrontendManager._write_auto_managed_versions("o", "r", ["1.0.0", "1.1.0"]) + FrontendManager._mark_auto_managed("o", "r", "1.0.0") + FrontendManager._mark_auto_managed("o", "r", "1.1.0") FrontendManager._prune_auto_managed_versions("o", "r", keep_version="1.1.0") @@ -394,7 +362,8 @@ def test_prune_auto_managed_versions_removes_stale_and_keeps_pinned( def test_untrack_auto_managed_version_does_not_delete_folder(custom_frontends_root): version_dir = _make_version_dir(custom_frontends_root, "o", "r", "1.0.0") - FrontendManager._write_auto_managed_versions("o", "r", ["1.0.0", "1.1.0"]) + FrontendManager._mark_auto_managed("o", "r", "1.0.0") + FrontendManager._mark_auto_managed("o", "r", "1.1.0") FrontendManager._untrack_auto_managed_version("o", "r", "1.0.0") @@ -402,6 +371,11 @@ def test_untrack_auto_managed_version_does_not_delete_folder(custom_frontends_ro assert FrontendManager._read_auto_managed_versions("o", "r") == ["1.1.0"] +def test_untrack_auto_managed_version_missing_is_noop(custom_frontends_root): + FrontendManager._untrack_auto_managed_version("o", "r", "1.0.0") + assert FrontendManager._read_auto_managed_versions("o", "r") == [] + + def test_init_frontend_latest_prunes_previous_auto_managed_versions( custom_frontends_root, mock_provider, mock_releases ): @@ -409,9 +383,7 @@ def test_init_frontend_latest_prunes_previous_auto_managed_versions( pinned = _make_version_dir( custom_frontends_root, "test-owner", "test-repo", "1.1.5" ) - FrontendManager._write_auto_managed_versions( - "test-owner", "test-repo", ["1.0.0"] - ) + FrontendManager._mark_auto_managed("test-owner", "test-repo", "1.0.0") def fake_download(release, destination_path): Path(destination_path).mkdir(parents=True, exist_ok=True) @@ -439,9 +411,7 @@ def test_init_frontend_explicit_version_promotes_out_of_auto_managed( custom_frontends_root, mock_provider ): _make_version_dir(custom_frontends_root, "test-owner", "test-repo", "1.0.0") - FrontendManager._write_auto_managed_versions( - "test-owner", "test-repo", ["1.0.0"] - ) + FrontendManager._mark_auto_managed("test-owner", "test-repo", "1.0.0") result = FrontendManager.init_frontend_unsafe( "test-owner/test-repo@v1.0.0", mock_provider @@ -459,9 +429,7 @@ def test_init_frontend_explicit_no_v_prefix_promotes_out_of_auto_managed( custom_frontends_root, mock_provider ): _make_version_dir(custom_frontends_root, "test-owner", "test-repo", "1.0.0") - FrontendManager._write_auto_managed_versions( - "test-owner", "test-repo", ["1.0.0"] - ) + FrontendManager._mark_auto_managed("test-owner", "test-repo", "1.0.0") with patch( "app.frontend_management.download_release_asset_zip"