diff --git a/README.md b/README.md index 1f50bda..7662730 100644 --- a/README.md +++ b/README.md @@ -501,6 +501,8 @@ 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. + > **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. ### 🌐 Web Interface Features diff --git a/py/api/lora_import_routes.py b/py/api/lora_import_routes.py index fa63234..2d45de3 100644 --- a/py/api/lora_import_routes.py +++ b/py/api/lora_import_routes.py @@ -100,6 +100,25 @@ def _import_line(result, skipped): return " · ".join(parts) +def _all_downloads_failed(errors): + """The warning when every planned download failed (#337).""" + if errors == 1: + return "Warning: the example image could not be downloaded" + return f"Warning: none of the {errors:,} example images could be downloaded" + + +def _with_downloads(line, downloaded, errors): + """``line`` with the failed downloads: counted when some failed, led by + a warning when every planned download failed (#337). The job still + finishes: the import itself succeeded and the downloads are best effort.""" + if not errors: + return line + if not downloaded: + return f"{_all_downloads_failed(errors)} · {line}" + failed = f"{_count(errors, 'example image')} could not be downloaded" + return f"{line} · {failed}" + + class LoraImportRoutesMixin: """Mixin providing the reviewed LoRA import (preview, apply, job).""" @@ -415,7 +434,12 @@ class LoraImportRoutesMixin: downloaded = staged.stats.get("downloaded", 0) errors = staged.stats.get("download_errors", 0) job.record_counted("downloaded", downloaded) - job.record_counted("download_errors", errors) + rows = [] + if errors and not downloaded: + warning = _all_downloads_failed(errors) + rows = [{"path": "", "detail": f"{warning}; the log has the details"}] + job.log_line(warning) + job.record_counted("download_errors", errors, rows) if downloaded or errors: job.log_line( f"Downloaded {_count(downloaded, 'example image')}" @@ -451,4 +475,9 @@ class LoraImportRoutesMixin: if result["removed"]: job.log_line(f"Replaced {_count(result['removed'], 'previous import')}") job.log_line(line) + line = _with_downloads( + line, + staged.stats.get("downloaded", 0), + staged.stats.get("download_errors", 0), + ) job.update(progress=100, processed=job.total, status_line=line) diff --git a/py/lora_utils.py b/py/lora_utils.py index cde7cc3..ef77b86 100644 --- a/py/lora_utils.py +++ b/py/lora_utils.py @@ -12,10 +12,11 @@ import json import os import re import threading +import urllib.error import urllib.request from pathlib import Path from typing import Dict, List, Optional, Tuple -from urllib.parse import urlsplit +from urllib.parse import urljoin, urlsplit try: from ..utils.logging_config import get_logger @@ -427,11 +428,19 @@ 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. +MAX_CIVITAI_REDIRECTS = 3 +_REDIRECT_CODES = frozenset({301, 302, 303, 307, 308}) +_USER_AGENT = "ComfyUI-PromptManager/1.0" + + class _RefuseRedirects(urllib.request.HTTPRedirectHandler): - """Never follow a redirect: urllib would re-send the Bearer header to it.""" + """Never follow a redirect inside urllib: it would re-send the Bearer + header. A 3xx surfaces as ``HTTPError``; ``_download_one`` decides.""" def redirect_request(self, req, fp, code, msg, headers, newurl): - logger.warning(f"Refusing redirect from {req.full_url} to {newurl}") return None @@ -448,6 +457,82 @@ def _open_url(req: urllib.request.Request, timeout: float = 10): return _opener.open(req, timeout=timeout) +def _loggable(url: str) -> str: + """``url`` without its query, fragment or credentials (signed URLs).""" + try: + parts = urlsplit(url) + host = parts.hostname or "" + except ValueError: + return "" + return f"{parts.scheme}://{host}{parts.path}" + + +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.""" + try: + parts = urlsplit(url) + port = parts.port + except ValueError: + return False + hostname = (parts.hostname or "").lower() + return ( + parts.scheme == "https" + and (hostname == "civitai.com" or hostname.endswith(".civitai.com")) + and parts.username is None + and parts.password is None + and port in (None, 443) + ) + + +def _redirect_target(url: str, error) -> Optional[str]: + """The URL a 3xx ``error`` for ``url`` points to, or None when absent.""" + if error.code not in _REDIRECT_CODES: + return None + location = (error.headers.get("Location") or "").strip() + if not location: + return None + return urljoin(url, location) + + +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. + """ + headers = {"User-Agent": _USER_AGENT} + if api_key: + headers["Authorization"] = f"Bearer {api_key}" + current = url + for hop in range(MAX_CIVITAI_REDIRECTS + 1): + req = urllib.request.Request(current, headers=headers) + try: + return _open_url(req, timeout=10) + except urllib.error.HTTPError as error: + with error: + target = _redirect_target(current, error) + 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 + if not is_civitai_redirect_target(target): + logger.warning( + f"Refusing redirect from {_loggable(current)} to {_loggable(target)}" + ) + return None + 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 + + def is_civitai_url(url: str) -> bool: """True when ``url`` is an HTTPS URL on an allow-listed CivitAI host.""" try: @@ -492,21 +577,21 @@ 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, and refuses bodies larger than - MAX_CIVITAI_DOWNLOAD_BYTES. + 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. """ if not is_civitai_url(url): - logger.warning(f"Refusing non-CivitAI download URL: {url}") + logger.warning(f"Refusing non-CivitAI download URL: {_loggable(url)}") return None try: - headers = {"User-Agent": "ComfyUI-PromptManager/1.0"} - if api_key: - headers["Authorization"] = f"Bearer {api_key}" - req = urllib.request.Request(url, headers=headers) - with _open_url(req, timeout=10) as resp: + 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: {url}") + logger.warning(f"Refusing oversized download: {_loggable(url)}") return None # Resize to thumbnail to save disk space @@ -520,7 +605,7 @@ def _download_one(url: str, local_path: Path, api_key: str) -> Optional[str]: _save_atomically(rgb, local_path) return str(local_path.resolve()) except Exception as e: - logger.debug(f"Failed to download {url}: {e}") + logger.debug(f"Failed to download {_loggable(url)}: {e}") return None diff --git a/tests/test_lora_import.py b/tests/test_lora_import.py index 54885d7..b50b671 100644 --- a/tests/test_lora_import.py +++ b/tests/test_lora_import.py @@ -239,10 +239,13 @@ class LoraImportAPITestCase(LoraAPITestCase): self.requests = [] self.gate = None # a threading.Event a download waits on self.downloading = threading.Event() + self.failing = set() # URLs whose download fails def fake_open(req, timeout=None): self.requests.append(req) self.downloading.set() + if req.full_url in self.failing: + raise OSError("connection reset") if self.gate is not None: self.gate.wait(10) return _FakeResponse(_png_bytes()) @@ -597,6 +600,114 @@ class LoraImportRouteTests(LoraImportAPITestCase): self.assertEqual(sorted(p["text"] for p in imported), ["neon street", "plain"]) +class LoraImportDownloadOutcomeTests(LoraImportAPITestCase): + """#337: failed CivitAI downloads show in the job's status line.""" + + URLS = ["https://civitai.com/a.png", "https://civitai.com/b.png"] + + def _lora_with_images(self): + _write_metadata( + self.loras, + "neon", + words=["glow"], + images=list(self.URLS), + prompt="neon street", + ) + _write_metadata(self.loras, "plain") + + async def _run(self): + _plan, job = await self._apply_reviewed() + await job.task + return job + + async def test_status_line_unchanged_when_every_download_succeeds(self): + self._lora_with_images() + + job = await self._run() + + self.assertEqual(job.state, "finished", job.log) + self.assertEqual(job.status_line, "Imported 2 LoRAs") + self.assertEqual(job.outcomes["download_errors"], []) + + async def test_status_line_counts_failed_downloads(self): + self._lora_with_images() + self.failing = {self.URLS[1]} + + job = await self._run() + + self.assertEqual(job.state, "finished", job.log) + self.assertEqual( + (job.counters["downloaded"], job.counters["download_errors"]), (1, 1) + ) + self.assertEqual( + job.status_line, + "Imported 2 LoRAs · 1 example image could not be downloaded", + ) + self.assertNotIn("Warning", job.status_line) + self.assertEqual(job.outcomes["download_errors"], []) + + async def test_every_download_failing_is_a_warning(self): + self._lora_with_images() + self.failing = set(self.URLS) + + job = await self._run() + + # The import itself succeeded, so the job still finishes. + self.assertEqual(job.state, "finished", job.log) + self.assertEqual( + (job.counters["downloaded"], job.counters["download_errors"]), (0, 2) + ) + warning = "Warning: none of the 2 example images could be downloaded" + self.assertEqual(job.status_line, f"{warning} · Imported 2 LoRAs") + self.assertEqual( + job.outcomes["download_errors"], + [ + { + "path": "", + "root": "", + "root_index": None, + "detail": f"{warning}; the log has the details", + } + ], + ) + self.assertTrue(any(warning in line for line in job.log), job.log) + self.assertEqual(job.summary()["status_line"], job.status_line) + self.assertNotIn(SECRET_KEY, json.dumps(job.report())) + + async def test_one_planned_download_failing_reads_in_the_singular(self): + _write_metadata(self.loras, "neon", images=[self.URLS[0]], prompt="neon street") + self.failing = {self.URLS[0]} + + job = await self._run() + + self.assertEqual( + job.status_line, + "Warning: the example image could not be downloaded · Imported 1 LoRA", + ) + + async def test_legacy_scan_events_unchanged_when_every_download_fails(self): + self._lora_with_images() + self.failing = set(self.URLS) + + resp = await self.client.request("POST", SCAN) + events = self._sse(await resp.text()) + + self.assertEqual(resp.status, 200) + self.assertEqual( + events[-1], + { + "type": "complete", + "progress": 100, + "total": 2, + "imported": 2, + "skipped": 0, + }, + ) + for event in events[:-1]: + self.assertEqual(event["type"], "progress") + self.assertNotIn("Warning", event.get("status", "")) + + class LoraImportRefusalTests(LoraImportAPITestCase): """Every refusal leaves the previous imports as they were.""" diff --git a/tests/test_lora_utils.py b/tests/test_lora_utils.py index 81f3357..5de87d2 100644 --- a/tests/test_lora_utils.py +++ b/tests/test_lora_utils.py @@ -15,6 +15,7 @@ import threading import time import types import unittest +import urllib.error import urllib.parse import urllib.request import urllib.response @@ -686,7 +687,7 @@ class TestRedirectsAreRefused(unittest.TestCase): if host not in CIVITAI_HOSTS: self.assertIsNone(req.get_header("Authorization")) - def test_redirect_to_civitai_is_not_followed_either(self): + def test_redirect_to_civitai_is_followed_without_the_key(self): handler = _FakeHTTPSHandler("https://image.civitai.com/other.png") opener = lora_utils._build_opener(handler) @@ -695,9 +696,200 @@ class TestRedirectsAreRefused(unittest.TestCase): "https://civitai.com/api/download/1", self.tmpdir / "x.jpg", "key" ) + # image.civitai.com redirects again (the fake handler always does), + # so the hop limit ends this chain; the first request alone had the key. self.assertIsNone(result) + self.assertEqual(len(handler.requests), lora_utils.MAX_CIVITAI_REDIRECTS + 1) + self.assertEqual(handler.requests[0].get_header("Authorization"), "Bearer key") + for req in handler.requests[1:]: + self.assertIsNone(req.get_header("Authorization")) + + +class _ChainHTTPSHandler(urllib.request.HTTPSHandler): + """Serves ``routes`` (url -> (status, Location or None)) through a real + urllib opener, recording every request; an unknown URL answers 404.""" + + def __init__(self, routes): + super().__init__() + self.routes = routes + self.requests = [] + + def https_open(self, req): + self.requests.append(req) + status, location = self.routes.get(req.full_url, (404, None)) + if location is not None: + headers = email.message_from_string(f"Location: {location}\n") + resp = urllib.response.addinfourl( + BytesIO(b""), headers, req.full_url, status + ) + resp.msg = "Redirect" + return resp + body = _png_bytes() if status == 200 else b"" + headers = email.message_from_string("Content-Type: image/png\n") + resp = urllib.response.addinfourl(BytesIO(body), headers, req.full_url, status) + resp.msg = "OK" if status == 200 else "Not Found" + return resp + + def http_open(self, req): # a plain-HTTP hop must never be requested + self.requests.append(req) + raise AssertionError(f"plain HTTP requested: {req.full_url}") + + +START = "https://image.civitai.com/xG1/abc/width=450/1.jpeg" +BLOB = "https://blobs-b2.civitai.com/file/civitai-media-cache/1/original" + + +class TestCivitaiRedirects(unittest.TestCase): + """#337: image.civitai.com now redirects to blobs-b2.civitai.com. + + A redirect is followed only over HTTPS to civitai.com or a subdomain of + it, at most MAX_CIVITAI_REDIRECTS hops, and never with the API 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 _refused(self, location, status=302): + result, requests = self._download({START: (status, location)}) + self.assertIsNone(result) + self.assertEqual([r.full_url for r in requests], [START]) + self.assertFalse(self.target.exists()) + + def test_redirect_to_an_allowed_host_succeeds_without_the_key(self): + for status in (301, 302, 303, 307, 308): + with self.subTest(status=status): + if self.target.exists(): + os.remove(self.target) + result, requests = self._download( + {START: (status, BLOB), BLOB: (200, None)} + ) + + self.assertIsNotNone(result) + self.assertTrue(self.target.exists()) + 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")) + # urllib adds Host itself; nothing else rides along. + self.assertEqual( + [h for h in requests[1].header_items() if h[0] != "Host"], + [("User-agent", "ComfyUI-PromptManager/1.0")], + ) + + def test_a_relative_location_resolves_on_the_same_host(self): + follow = "https://image.civitai.com/elsewhere/1.jpeg" + result, requests = self._download( + {START: (302, "/elsewhere/1.jpeg"), follow: (200, None)} + ) + + self.assertIsNotNone(result) + self.assertEqual([r.full_url for r in requests], [START, follow]) + self.assertIsNone(requests[1].get_header("Authorization")) + + def test_a_plain_http_redirect_is_refused(self): + self._refused("http://blobs-b2.civitai.com/file/1") + + def test_a_redirect_to_another_host_is_refused(self): + self._refused("https://evil.example/steal") + + def test_lookalike_hosts_are_refused(self): + for location in ( + "https://civitai.com.evil.example/x", + "https://evilcivitai.com/x", + "https://civitai.com@evil.example/x", + "https://user:pw@blobs-b2.civitai.com/x", + "https://blobs-b2.civitai.com:8443/x", + "https://civitai.co/x", + "//evil.example/x", + "ftp://blobs-b2.civitai.com/x", + "https://[::1]/x", + ): + with self.subTest(location=location): + self._refused(location) + + def test_a_redirect_without_a_location_is_refused(self): + result, requests = self._download({START: (302, "")}) + self.assertIsNone(result) + self.assertEqual(len(requests), 1) + + def test_the_hop_limit_is_enforced(self): + limit = lora_utils.MAX_CIVITAI_REDIRECTS + self.assertEqual(limit, 3) + hops = [f"https://blobs-b2.civitai.com/hop/{i}" for i in range(limit + 1)] + + # Exactly ``limit`` redirects: followed. + routes = {START: (302, hops[0])} + for i in range(limit - 1): + routes[hops[i]] = (302, hops[i + 1]) + routes[hops[limit - 1]] = (200, None) + result, requests = self._download(routes) + self.assertIsNotNone(result) + self.assertEqual(len(requests), limit + 1) + self.assertTrue( + all(r.get_header("Authorization") is None for r in requests[1:]) + ) + + # One more redirect: refused, and the next hop is never requested. + os.remove(self.target) + routes[hops[limit - 1]] = (302, hops[limit]) + routes[hops[limit]] = (200, None) + result, requests = self._download(routes) + self.assertIsNone(result) + self.assertEqual(len(requests), limit + 1) + self.assertNotIn(hops[limit], [r.full_url for r in requests]) + self.assertFalse(self.target.exists()) + + def test_a_redirect_loop_ends_at_the_hop_limit(self): + result, requests = self._download({START: (302, START)}) + self.assertIsNone(result) + self.assertEqual(len(requests), lora_utils.MAX_CIVITAI_REDIRECTS + 1) + + def test_size_cap_applies_after_a_redirect(self): + with patch.object(lora_utils, "MAX_CIVITAI_DOWNLOAD_BYTES", 10): + result, requests = self._download({START: (302, BLOB), BLOB: (200, None)}) + self.assertIsNone(result) + self.assertEqual(len(requests), 2) + self.assertFalse(self.target.exists()) + + def test_an_error_after_a_redirect_fails_cleanly(self): + result, requests = self._download({START: (302, BLOB)}) # BLOB is 404 + self.assertIsNone(result) + self.assertEqual(len(requests), 2) + self.assertFalse(self.target.exists()) + + def test_the_opener_itself_still_refuses_every_redirect(self): + handler = _ChainHTTPSHandler({START: (302, BLOB), BLOB: (200, None)}) + opener = lora_utils._build_opener(handler) + req = urllib.request.Request(START, headers={"Authorization": "Bearer k"}) + + with self.assertRaises(urllib.error.HTTPError) as caught: + opener.open(req, timeout=5) + caught.exception.close() self.assertEqual(len(handler.requests), 1) + def test_redirect_logs_never_carry_the_key_or_the_query(self): + signed = BLOB + "?X-Amz-Signature=sig123" + with self.assertLogs("prompt_manager.lora_utils", level="DEBUG") as logs: + self._download( + {START: (302, "https://evil.example/x?token=t0k3n")}, + api_key="secret-key", + ) + self._download({START: (302, signed), signed: (200, None)}) + text = "\n".join(logs.output) + self.assertNotIn("secret-key", text) + self.assertNotIn("t0k3n", text) + self.assertNotIn("sig123", text) + class TestDownloadCivitaiImages(unittest.TestCase):