diff --git a/README.md b/README.md index 7662730..5ff1860 100644 --- a/README.md +++ b/README.md @@ -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. diff --git a/TESTING.md b/TESTING.md index 898cba8..8ea2d7f 100644 --- a/TESTING.md +++ b/TESTING.md @@ -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 diff --git a/py/lora_utils.py b/py/lora_utils.py index ef77b86..5965af4 100644 --- a/py/lora_utils.py +++ b/py/lora_utils.py @@ -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-" diff --git a/tests/test_lora_utils.py b/tests/test_lora_utils.py index 5de87d2..7348c83 100644 --- a/tests/test_lora_utils.py +++ b/tests/test_lora_utils.py @@ -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):