fix(lora): allow civitai.red/.green and log why a CivitAI download failed (#340) (#341)

civitai.red and civitai.green (and their subdomains) are now allowed
redirect targets under the civitai.com rules: HTTPS only, default port,
no credentials, matched on the parsed host name, at most 3 hops, never
with the key. Their www/api/image hosts are also allowed starting hosts,
whose first request carries the key as image.civitai.com does.

Each failed download now logs one warning with the URL without query or
credentials and a short reason: HTTP status, refused redirect, hop
limit, size cap, decode or save error, timeout or network error.
This commit is contained in:
Vito
2026-10-25 12:53:26 -07:00
committed by GitHub
parent 863d717435
commit f9835aa4d6
4 changed files with 415 additions and 52 deletions
+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.
+8 -1
View File
@@ -3151,7 +3151,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
+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-"
+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):