nodes/ComfyAssets-ComfyUI_PromptManager
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Merge remote-tracking branch 'origin/main' into feat/326-system-s9-backup-ui

Browse Source
This commit is contained in:
Vito Sansevero
2026-10-25 13:03:08 -07:00
parent 447fd4668b f9835aa4d6
commit 367db8d7d6
4 changed files with 415 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+1 -1
View File
@@ -501,7 +501,7 @@ To add your key:
2. Paste it into the **CivitAI API Key** field in the Integrations settings
3. Save and re-import to download any previously skipped images
The key is sent only over HTTPS to CivitAI's own hosts. CivitAI serves images through a redirect to its storage host; PromptManager follows such a redirect only over HTTPS to a `civitai.com` host, at most three hops, and never sends the key with it. When downloads fail, the import's status line says how many; when none could be downloaded, it starts with a warning.
The key is sent only over HTTPS to CivitAI's own hosts. CivitAI serves images through a redirect to its storage host; PromptManager follows such a redirect only over HTTPS to a host on CivitAI's own `civitai.com`, `civitai.red` or `civitai.green` domains, at most three hops, and never sends the key with it. When downloads fail, the import's status line says how many; when none could be downloaded, it starts with a warning. The server log has one warning per failed image with the reason (HTTP status, refused redirect, size cap, unreadable image, timeout or network error).
> **Note:** Re-importing is safe — PromptManager skips LoRAs that were already imported. To do a fresh import, the previous `lora-manager` data is cleared automatically before re-scanning.
TESTING.md
+8 -1
View File
@@ -3217,7 +3217,14 @@ replace binds the category, not one parameter per prompt, so SQLite's variable l
never caps a collection. Cached CivitAI images are written to a `.pm-*.part` file in the
same folder and renamed into place, with the umask's mode as a plain save gives; each import
first removes `.pm-*.part` files older than an hour from the image cache and its LoRA
folders (`TestStaleDownloadSweep`). The legacy `POST /lora/scan` runs the same stages under its SSE
folders (`TestStaleDownloadSweep`). A download starts only at an HTTPS `civitai.com`,
`civitai.red` or `civitai.green` host from the fixed `CIVITAI_HOSTS` list (the key goes
with that first request only); a redirect is followed only over HTTPS to one of those
domains or a subdomain (parsed host name, default port, no credentials), at most three
hops, without the key (`TestCivitaiRedirects`, `TestCivitaiRedAndGreenDomains`). Each
failed download logs one warning naming the URL without its query and a short reason
(HTTP status, refused redirect, hop limit, size cap, decode or save error, timeout or
network error), never the key (`TestDownloadFailureWarnings`). The legacy `POST /lora/scan` runs the same stages under its SSE
event shape; a page that closes its connection stops staging before its next LoRA, and the
scan keeps its records hold until staging has stopped, then ends `cancelled` in Activity
(`test_a_disconnected_legacy_scan_stops_and_holds_until_it_stopped`) (`test_legacy_scan_events_unchanged_on_success`), holds
py/lora_utils.py
+96 -44
View File
@@ -8,6 +8,7 @@ return empty results rather than raising.
"""
import hashlib
import http.client
import json
import os
import re
@@ -420,22 +421,35 @@ def get_preview_image_from_metadata(
_THUMB_MAX_SIZE = 512
# CivitAI's own domains (#340: civitai.red and civitai.green serve the same
# content as civitai.com).
CIVITAI_DOMAINS = ("civitai.com", "civitai.red", "civitai.green")
# Only these hosts may ever receive the CivitAI API key, and only over HTTPS.
CIVITAI_HOSTS = frozenset(
{"civitai.com", "www.civitai.com", "api.civitai.com", "image.civitai.com"}
f"{prefix}{domain}"
for domain in CIVITAI_DOMAINS
for prefix in ("", "www.", "api.", "image.")
)
MAX_CIVITAI_DOWNLOAD_BYTES = 50 * 1024 * 1024
_DOWNLOAD_CHUNK_BYTES = 64 * 1024
# A redirect is followed (#337: image.civitai.com now answers with a redirect
# to blobs-b2.civitai.com) only over HTTPS to civitai.com or a subdomain of
# it, at most this many hops, and never with the API key.
# to blobs-b2.civitai.com) only over HTTPS to one of CIVITAI_DOMAINS or a
# subdomain of one, at most this many hops, and never with the API key.
MAX_CIVITAI_REDIRECTS = 3
_REDIRECT_CODES = frozenset({301, 302, 303, 307, 308})
_USER_AGENT = "ComfyUI-PromptManager/1.0"
class _DownloadFailed(Exception):
"""A download failed for the short, loggable ``reason`` it carries."""
def __init__(self, reason: str):
super().__init__(reason)
self.reason = reason
class _RefuseRedirects(urllib.request.HTTPRedirectHandler):
"""Never follow a redirect inside urllib: it would re-send the Bearer
header. A 3xx surfaces as ``HTTPError``; ``_download_one`` decides."""
@@ -467,10 +481,18 @@ def _loggable(url: str) -> str:
return f"{parts.scheme}://{host}{parts.path}"
def _on_civitai_domain(hostname: str) -> bool:
"""True when ``hostname`` is one of CIVITAI_DOMAINS or a subdomain of one."""
return any(
hostname == domain or hostname.endswith(f".{domain}")
for domain in CIVITAI_DOMAINS
)
def is_civitai_redirect_target(url: str) -> bool:
"""True when a redirect to ``url`` may be followed: HTTPS on civitai.com
or a subdomain of it, matched on the parsed host name, default port and
no credentials in the URL."""
"""True when a redirect to ``url`` may be followed: HTTPS on one of
CIVITAI_DOMAINS or a subdomain of one, matched on the parsed host name,
default port and no credentials in the URL."""
try:
parts = urlsplit(url)
port = parts.port
@@ -479,7 +501,7 @@ def is_civitai_redirect_target(url: str) -> bool:
hostname = (parts.hostname or "").lower()
return (
parts.scheme == "https"
and (hostname == "civitai.com" or hostname.endswith(".civitai.com"))
and _on_civitai_domain(hostname)
and parts.username is None
and parts.password is None
and port in (None, 443)
@@ -500,8 +522,9 @@ def _open_following_civitai_redirects(url: str, api_key: str):
"""Open ``url``, following only allowed CivitAI redirects (keyless).
The first request carries the key; every redirected one carries the
User-Agent alone. Returns the open response, or None when a redirect is
refused or the hop limit is reached. Other HTTP errors propagate.
User-Agent alone. Returns the open response; raises ``_DownloadFailed``
when a redirect is refused or the hop limit is reached. Other HTTP
errors propagate.
"""
headers = {"User-Agent": _USER_AGENT}
if api_key:
@@ -517,20 +540,13 @@ def _open_following_civitai_redirects(url: str, api_key: str):
if target is None:
raise
if hop == MAX_CIVITAI_REDIRECTS:
logger.warning(
f"Refusing redirect past {MAX_CIVITAI_REDIRECTS} hops "
f"from {_loggable(url)}"
)
return None
raise _DownloadFailed(f"more than {MAX_CIVITAI_REDIRECTS} redirects")
if not is_civitai_redirect_target(target):
logger.warning(
f"Refusing redirect from {_loggable(current)} to {_loggable(target)}"
)
return None
raise _DownloadFailed(f"redirect refused to {_loggable(target)}")
logger.debug(f"Following redirect to {_loggable(target)}")
current = target
headers = {"User-Agent": _USER_AGENT} # never the key across a redirect
return None # pragma: no cover - the loop returns on its last hop
raise AssertionError("unreachable") # pragma: no cover - loop ends above
def is_civitai_url(url: str) -> bool:
@@ -573,40 +589,76 @@ def _read_capped(resp, cap: int) -> Optional[bytes]:
return b"".join(chunks)
def _network_reason(error: BaseException) -> str:
"""A short reason for a failed request: a timeout or a network error."""
cause = getattr(error, "reason", error) # URLError wraps the socket error
if isinstance(cause, TimeoutError):
return "timed out"
return f"network error: {cause}"
def _fetch_civitai_body(url: str, api_key: str) -> bytes:
"""The body behind ``url`` (redirects as allowed, size capped), or
``_DownloadFailed`` with the reason it could not be fetched."""
try:
with _open_following_civitai_redirects(url, api_key) as resp:
raw = _read_capped(resp, MAX_CIVITAI_DOWNLOAD_BYTES)
except urllib.error.HTTPError as error:
error.close()
raise _DownloadFailed(f"HTTP {error.code}") from None
except (urllib.error.URLError, OSError, http.client.HTTPException) as error:
raise _DownloadFailed(_network_reason(error)) from None
if raw is None:
cap = MAX_CIVITAI_DOWNLOAD_BYTES
size = f"{cap // (1024 * 1024)} MB" if cap >= 1024 * 1024 else f"{cap} byte"
raise _DownloadFailed(f"larger than the {size} cap")
return raw
def _thumbnail(raw: bytes):
"""``raw`` decoded and shrunk to an RGB thumbnail, or ``_DownloadFailed``."""
from io import BytesIO
from PIL import Image
try:
with Image.open(BytesIO(raw)) as img:
img.thumbnail((_THUMB_MAX_SIZE, _THUMB_MAX_SIZE), Image.LANCZOS)
return img.convert("RGB")
except Exception as error: # PIL raises several types for bad data
raise _DownloadFailed(
f"image could not be decoded ({type(error).__name__})"
) from None
def _download_one(url: str, local_path: Path, api_key: str) -> Optional[str]:
"""Download a single CivitAI image, resize to thumbnail, save as JPEG.
Refuses any URL that is not HTTPS on an allow-listed CivitAI host so the
API key is never sent elsewhere, follows a redirect only as
``_open_following_civitai_redirects`` allows (without the key), and
refuses bodies larger than MAX_CIVITAI_DOWNLOAD_BYTES.
refuses bodies larger than MAX_CIVITAI_DOWNLOAD_BYTES. A failure logs
one warning with a short reason and the URL without query or
credentials (#340), and returns None.
"""
if not is_civitai_url(url):
logger.warning(f"Refusing non-CivitAI download URL: {_loggable(url)}")
return None
try:
resp = _open_following_civitai_redirects(url, api_key)
if resp is None:
return None
with resp:
raw = _read_capped(resp, MAX_CIVITAI_DOWNLOAD_BYTES)
if raw is None:
logger.warning(f"Refusing oversized download: {_loggable(url)}")
return None
# Resize to thumbnail to save disk space
from io import BytesIO
from PIL import Image
with Image.open(BytesIO(raw)) as img:
img.thumbnail((_THUMB_MAX_SIZE, _THUMB_MAX_SIZE), Image.LANCZOS)
rgb = img.convert("RGB")
_save_atomically(rgb, local_path)
if not is_civitai_url(url):
raise _DownloadFailed("not an allowed CivitAI URL")
image = _thumbnail(_fetch_civitai_body(url, api_key))
try:
_save_atomically(image, local_path)
except OSError as error:
raise _DownloadFailed(
f"image could not be saved ({error.strerror or type(error).__name__})"
) from None
except _DownloadFailed as failure:
reason = failure.reason
except Exception as error: # best effort: an import never fails on one image
reason = f"unexpected {type(error).__name__}"
else:
return str(local_path.resolve())
except Exception as e:
logger.debug(f"Failed to download {_loggable(url)}: {e}")
return None
logger.warning(f"CivitAI image download failed: {_loggable(url)}: {reason}")
return None
_PART_PREFIX = ".pm-"
tests/test_lora_utils.py
+310 -6
View File
@@ -296,12 +296,9 @@ class TestCivitaiUrlPolicy(unittest.TestCase):
self.assertEqual(
CIVITAI_HOSTS,
frozenset(
{
"civitai.com",
"www.civitai.com",
"api.civitai.com",
"image.civitai.com",
}
f"{prefix}{domain}"
for domain in ("civitai.com", "civitai.red", "civitai.green")
for prefix in ("", "www.", "api.", "image.")
),
)
self.assertEqual(MAX_CIVITAI_DOWNLOAD_BYTES, 50 * 1024 * 1024)
@@ -334,6 +331,36 @@ class TestCivitaiUrlPolicy(unittest.TestCase):
def test_explicit_default_port_is_allowed(self):
self.assertTrue(is_civitai_url("https://civitai.com:443/x"))
def test_red_and_green_starting_hosts_are_allowed(self):
# #340: CivitAI's own .red and .green domains serve the same content.
for url in (
"https://image.civitai.red/xG1/abc/width=450/1.jpeg",
"https://civitai.red/api/download/1",
"https://civitai.green/api/download/1",
"https://image.civitai.green/xG1/abc/1.jpeg",
"https://IMAGE.CIVITAI.RED/x",
"https://civitai.green:443/x",
):
self.assertTrue(is_civitai_url(url), url)
def test_red_and_green_lookalikes_are_refused(self):
for url in (
"http://image.civitai.red/x",
"http://civitai.green/x",
"https://civitai.red.evil.example/x",
"https://civitai.green.evil.example/x",
"https://evilcivitai.red/x",
"https://evilcivitai.green/x",
"https://civitai.redx/x",
"https://civitai.greenx/x",
"https://civitai.red@evil.example/x",
"https://user:pw@civitai.red/x",
"https://civitai.green:8443/x",
"https://blobs-b2.civitai.red/x", # starting hosts stay a fixed list
"https://evil.example/civitai.red",
):
self.assertFalse(is_civitai_url(url), url)
class TestLoraManagerPathPolicy(unittest.TestCase):
"""Only a direct child of a ComfyUI custom_nodes dir named *lora* qualifies."""
@@ -891,6 +918,283 @@ class TestCivitaiRedirects(unittest.TestCase):
self.assertNotIn("sig123", text)
RED_BLOB = "https://blobs-b2.civitai.red/file/civitai-media-cache/1/original"
GREEN_BLOB = "https://media.civitai.green/file/1/original"
class TestCivitaiRedAndGreenDomains(unittest.TestCase):
"""#340: civitai.red and civitai.green are CivitAI's own domains.
They are redirect targets under the civitai.com rules (HTTPS, default
port, no credentials, parsed host name, 3 hops, never the key) and
starting hosts whose first request may carry the key.
"""
def setUp(self):
self.tmpdir = Path(tempfile.mkdtemp())
self.addCleanup(shutil.rmtree, self.tmpdir, True)
self.target = self.tmpdir / "x.jpg"
def _download(self, routes, url=START, api_key="secret-key"):
handler = _ChainHTTPSHandler(routes)
opener = lora_utils._build_opener(handler)
with patch.object(lora_utils, "_opener", opener):
result = lora_utils._download_one(url, self.target, api_key)
return result, handler.requests
def test_redirect_targets_on_red_and_green_are_allowed(self):
for url in (
"https://civitai.red/x",
"https://civitai.green/x",
RED_BLOB,
GREEN_BLOB,
"https://a.b.civitai.red:443/x",
"https://CIVITAI.GREEN/x",
):
self.assertTrue(lora_utils.is_civitai_redirect_target(url), url)
def test_redirect_target_lookalikes_are_refused(self):
for url in (
"https://civitai.red.evil.example/x",
"https://civitai.green.evil.example/x",
"https://evilcivitai.red/x",
"https://evilcivitai.green/x",
"https://civitai.redx/x",
"https://civitai.greenx/x",
"https://civitai.red@evil.example/x",
"https://user:pw@blobs-b2.civitai.red/x",
"https://user@civitai.green/x",
"https://blobs-b2.civitai.red:8443/x",
"http://blobs-b2.civitai.red/x",
"http://civitai.green/x",
"ftp://civitai.red/x",
"https://evil.example/civitai.red",
"https://evil.example/?h=civitai.green",
"https://red/x",
"https://.civitai.red.evil/x",
):
self.assertFalse(lora_utils.is_civitai_redirect_target(url), url)
def test_a_redirect_to_red_or_green_is_followed_without_the_key(self):
for blob in (RED_BLOB, GREEN_BLOB):
with self.subTest(blob=blob):
if self.target.exists():
os.remove(self.target)
result, requests = self._download(
{START: (302, blob), blob: (200, None)}
)
self.assertIsNotNone(result)
self.assertEqual([r.full_url for r in requests], [START, blob])
self.assertEqual(
requests[0].get_header("Authorization"), "Bearer secret-key"
)
self.assertIsNone(requests[1].get_header("Authorization"))
def test_lookalike_and_plain_http_redirects_are_refused(self):
for location in (
"https://civitai.red.evil.example/x",
"https://evilcivitai.red/x",
"https://civitai.redx/x",
"https://civitai.green.evil.example/x",
"https://evilcivitai.green/x",
"https://civitai.greenx/x",
"http://blobs-b2.civitai.red/x",
"http://civitai.green/x",
):
with self.subTest(location=location):
result, requests = self._download({START: (302, location)})
self.assertIsNone(result)
self.assertEqual([r.full_url for r in requests], [START])
self.assertFalse(self.target.exists())
def test_a_red_starting_host_gets_the_key_and_its_redirect_does_not(self):
start = "https://image.civitai.red/xG1/abc/width=450/1.jpeg"
result, requests = self._download(
{start: (302, BLOB), BLOB: (200, None)}, url=start
)
self.assertIsNotNone(result)
self.assertEqual([r.full_url for r in requests], [start, BLOB])
self.assertEqual(requests[0].get_header("Authorization"), "Bearer secret-key")
self.assertIsNone(requests[1].get_header("Authorization"))
def test_a_green_starting_host_downloads_directly(self):
start = "https://civitai.green/api/download/1"
result, requests = self._download({start: (200, None)}, url=start)
self.assertIsNotNone(result)
self.assertEqual(requests[0].get_header("Authorization"), "Bearer secret-key")
def test_the_hop_limit_spans_domains(self):
hops = [RED_BLOB, GREEN_BLOB, BLOB, "https://civitai.red/4"]
routes = {START: (302, hops[0])}
for here, there in zip(hops, hops[1:]):
routes[here] = (302, there)
routes[hops[-1]] = (200, None)
result, requests = self._download(routes)
self.assertIsNone(result)
self.assertEqual(len(requests), lora_utils.MAX_CIVITAI_REDIRECTS + 1)
self.assertNotIn(hops[-1], [r.full_url for r in requests])
for req in requests[1:]:
self.assertIsNone(req.get_header("Authorization"))
class _RaisingOpener:
"""An ``_open_url`` stand-in that raises ``error`` for every request."""
def __init__(self, error):
self.error = error
self.requests = []
def __call__(self, req, timeout=None):
self.requests.append(req)
raise self.error
class TestDownloadFailureWarnings(unittest.TestCase):
"""#340: each failed download logs one warning line with a short reason,
the URL without its query or credentials, and never the key."""
LOGGER = "prompt_manager.lora_utils"
SIGNED = START + "?token=t0k3n"
def setUp(self):
self.tmpdir = Path(tempfile.mkdtemp())
self.addCleanup(shutil.rmtree, self.tmpdir, True)
self.target = self.tmpdir / "x.jpg"
def _warning(self, run):
"""Run ``run`` and return its one WARNING-or-above record's text."""
with self.assertLogs(self.LOGGER, level="WARNING") as logs:
result = run()
self.assertIsNone(result)
self.assertFalse(self.target.exists())
self.assertEqual(len(logs.records), 1, logs.output)
record = logs.records[0]
self.assertEqual(record.levelname, "WARNING")
text = record.getMessage()
self.assertNotIn("secret-key", text)
self.assertNotIn("t0k3n", text)
self.assertNotIn("sig123", text)
self.assertIn(START.split("?")[0], text)
return text
def _chain(self, routes, url=SIGNED):
handler = _ChainHTTPSHandler(routes)
opener = lora_utils._build_opener(handler)
def run():
with patch.object(lora_utils, "_opener", opener):
return lora_utils._download_one(url, self.target, "secret-key")
return run
def _fake(self, opener, url=SIGNED):
def run():
with patch.object(lora_utils, "_open_url", opener):
return lora_utils._download_one(url, self.target, "secret-key")
return run
def test_http_status(self):
text = self._warning(self._chain({})) # unknown URL answers 404
self.assertIn("HTTP 404", text)
def test_http_status_after_a_redirect(self):
text = self._warning(self._chain({self.SIGNED: (302, BLOB)}))
self.assertIn("HTTP 404", text)
def test_refused_redirect(self):
text = self._warning(
self._chain({self.SIGNED: (302, "https://evil.example/x?sig123=1")})
)
self.assertIn("redirect refused", text)
self.assertIn("https://evil.example/x", text)
def test_hop_limit(self):
text = self._warning(self._chain({self.SIGNED: (302, self.SIGNED)}))
self.assertIn("more than 3 redirects", text)
def test_size_cap(self):
opener = _FakeOpener(
_FakeResponse(_png_bytes(), content_length=MAX_CIVITAI_DOWNLOAD_BYTES + 1)
)
text = self._warning(self._fake(opener))
self.assertIn("larger than the 50 MB cap", text)
def test_decode_error(self):
text = self._warning(self._fake(_FakeOpener(_FakeResponse(b"not an image"))))
self.assertIn("image could not be decoded", text)
def test_timeout(self):
for error in (
TimeoutError("timed out"),
urllib.error.URLError(TimeoutError("timed out")),
):
with self.subTest(error=type(error).__name__):
text = self._warning(self._fake(_RaisingOpener(error)))
self.assertIn("timed out", text)
def test_network_error(self):
for error in (
urllib.error.URLError("[Errno -2] Name or service not known"),
ConnectionResetError(104, "Connection reset by peer"),
):
with self.subTest(error=type(error).__name__):
text = self._warning(self._fake(_RaisingOpener(error)))
self.assertIn("network error", text)
def test_a_timeout_while_reading_the_body(self):
class SlowResponse(_FakeResponse):
def read(self, amt=-1):
raise TimeoutError("timed out")
text = self._warning(self._fake(_FakeOpener(SlowResponse(b""))))
self.assertIn("timed out", text)
def test_a_failed_save(self):
opener = _FakeOpener(_FakeResponse(_png_bytes()))
def run():
with patch.object(
lora_utils.os, "replace", side_effect=OSError(28, "No space left")
):
return self._fake(opener)()
text = self._warning(run)
self.assertIn("could not be saved", text)
def test_a_refused_starting_url(self):
def run():
return lora_utils._download_one(
"http://image.civitai.com/x?token=t0k3n", self.target, "secret-key"
)
with self.assertLogs(self.LOGGER, level="WARNING") as logs:
self.assertIsNone(run())
self.assertEqual(len(logs.records), 1)
text = logs.records[0].getMessage()
self.assertIn("not an allowed CivitAI URL", text)
self.assertNotIn("t0k3n", text)
def test_an_unexpected_error_names_only_its_type(self):
error = ValueError("unknown url type: https://x?token=t0k3n")
text = self._warning(self._fake(_RaisingOpener(error)))
self.assertIn("unexpected ValueError", text)
def test_the_line_names_the_download(self):
text = self._warning(self._chain({}))
self.assertTrue(text.startswith("CivitAI image download failed"), text)
def test_a_successful_download_logs_no_warning(self):
with self.assertNoLogs(self.LOGGER, level="WARNING"):
result = self._fake(_FakeOpener(_FakeResponse(_png_bytes())))()
self.assertIsNotNone(result)
class TestDownloadCivitaiImages(unittest.TestCase):
def setUp(self):