Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
157 changes: 134 additions & 23 deletions app/frontend_management.py

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Way too many long comments.

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fair — trimmed in 0a2b9f8. Net -76 lines of prose across the two files:

  • Killed all narrative/scenario comments inside test bodies and helper functions (they restated the test setup).
  • Removed the multi-line block comments on AUTO_MANAGED_VERSION_SPECIFIERS and the regex constant — kept tight one-liners.
  • Collapsed the init_frontend_unsafe "promote/prune" comments into nothing — the method names already say it.
  • Shortened docstrings on _read/_prune/_untrack_auto_managed_versions to one line each.

Kept just five short comments that aren't restating code:

  • AUTO_MANAGED_VERSION_SPECIFIERS constant (says why this tuple exists vs explicit pins)
  • The regex (security: explains the allowlist purpose)
  • The atomic-write step (non-obvious tmp+rename pattern)
  • The defense-in-depth _is_safe_version_dirname re-check in _prune_auto_managed_versions (would otherwise look redundant and get "simplified" away)

33/33 tests still pass.

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Cut further in 273314a — only one comment remains in the new code:

# 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):
    ...

This one I'd push back on removing: the _is_safe_version_dirname re-check inside _prune_auto_managed_versions looks redundant (the read path already filters), and without a comment a future maintainer will simplify it away. That would re-open the path-traversal vector Oracle flagged in the first review pass. The test test_prune_refuses_to_delete_outside_provider_dir would still pass without the re-check (because it relies on the resolve-and-check below it), so the comment is the only thing protecting the layered defense.

Everything else — class-constant explanations, regex purpose, atomic-write idiom, all docstrings on the new private helpers — is gone. Happy to delete this last one too if you'd rather rename the helper or merge the check into _is_safe_version_dirname's docstring instead.

33/33 tests still pass.

Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
import logging
import os
import re
import shutil
import sys
import tempfile
import zipfile
Expand Down Expand Up @@ -205,6 +206,97 @@ def download_release_asset_zip(release: Release, destination_path: str) -> None:
class FrontendManager:
CUSTOM_FRONTENDS_ROOT = str(Path(__file__).parents[1] / "web_custom_versions")

AUTO_MANAGED_VERSION_SPECIFIERS = ("latest", "prerelease")
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_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]:
"""Return versions ComfyUI auto-downloaded for @latest / @prerelease.
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:
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 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,
)

@classmethod
def _prune_auto_managed_versions(
cls, repo_owner: str, repo_name: str, keep_version: str
) -> None:
"""Remove previously auto-downloaded versions other than ``keep_version``."""
provider_dir = cls._provider_dir(repo_owner, repo_name)
for stale_version in cls._read_auto_managed_versions(repo_owner, repo_name):
if stale_version == keep_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 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:
"""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
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:
"""Get the required frontend package version."""
Expand Down Expand Up @@ -372,17 +464,16 @@ def init_frontend_unsafe(
return cls.default_frontend_path()

repo_owner, repo_name, version = cls.parse_version_string(version_string)
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, pinned_version)
return expected_path

logging.info(
Expand All @@ -393,27 +484,47 @@ def init_frontend_unsafe(
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
)
if not os.path.exists(web_root):
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)
web_root = str(cls._provider_dir(repo_owner, repo_name) / semantic_version)

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)
else:
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:
"""
Expand Down
155 changes: 155 additions & 0 deletions tests-unit/app_test/frontend_manager_test.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,6 @@
import argparse
from pathlib import Path

import pytest
from requests.exceptions import HTTPError
from unittest.mock import patch, mock_open
Expand Down Expand Up @@ -287,3 +289,156 @@ def test_get_installed_templates_version_not_installed():

# Assert
assert version is None


# Auto-managed @latest / @prerelease cleanup (CORE-285)


@pytest.fixture
def custom_frontends_root(tmp_path, monkeypatch):
root = tmp_path / "web_custom_versions"
root.mkdir()
monkeypatch.setattr(FrontendManager, "CUSTOM_FRONTENDS_ROOT", str(root))
return root


def _make_version_dir(root, owner, repo, version):
folder = root / f"{owner}_{repo}" / version
folder.mkdir(parents=True, exist_ok=True)
(folder / "index.html").write_text("<html></html>")
return folder


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"]


def test_read_auto_managed_versions_missing(custom_frontends_root):
assert FrontendManager._read_auto_managed_versions("o", "r") == []


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"
# 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_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")

_make_version_dir(custom_frontends_root, "o", "r", "1.0.0")
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 not (custom_frontends_root / "o_r" / "1.0.0").exists()


def test_prune_auto_managed_versions_removes_stale_and_keeps_pinned(
custom_frontends_root,
):
_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._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")

provider_dir = custom_frontends_root / "o_r"
assert not (provider_dir / "1.0.0").exists()
assert (provider_dir / "1.1.0").exists()
assert pinned.exists()
assert FrontendManager._read_auto_managed_versions("o", "r") == ["1.1.0"]


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._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")

assert version_dir.exists()
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
):
_make_version_dir(custom_frontends_root, "test-owner", "test-repo", "1.0.0")
pinned = _make_version_dir(
custom_frontends_root, "test-owner", "test-repo", "1.1.5"
)
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)
(Path(destination_path) / "index.html").write_text("<html></html>")

with patch(
"app.frontend_management.download_release_asset_zip",
side_effect=fake_download,
):
result = FrontendManager.init_frontend_unsafe(
"test-owner/test-repo@latest", mock_provider
)

provider_dir = custom_frontends_root / "test-owner_test-repo"
assert Path(result) == provider_dir / "2.0.0"
assert (provider_dir / "2.0.0").exists()
assert FrontendManager._read_auto_managed_versions(
"test-owner", "test-repo"
) == ["2.0.0"]
assert not (provider_dir / "1.0.0").exists()
assert pinned.exists()


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._mark_auto_managed("test-owner", "test-repo", "1.0.0")

result = FrontendManager.init_frontend_unsafe(
"test-owner/test-repo@v1.0.0", mock_provider
)

provider_dir = custom_frontends_root / "test-owner_test-repo"
assert Path(result) == provider_dir / "1.0.0"
assert FrontendManager._read_auto_managed_versions(
"test-owner", "test-repo"
) == []
assert (provider_dir / "1.0.0").exists()


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._mark_auto_managed("test-owner", "test-repo", "1.0.0")

with patch(
"app.frontend_management.download_release_asset_zip"
) as mock_download_zip:
FrontendManager.init_frontend_unsafe(
"test-owner/test-repo@1.0.0", mock_provider
)
mock_download_zip.assert_not_called()

assert FrontendManager._read_auto_managed_versions(
"test-owner", "test-repo"
) == []
Loading