fix(lora): follow CivitAI image redirects without the key (#337) (#338)

image.civitai.com now answers every image request with a redirect to blobs-b2.civitai.com, and the downloader refused all redirects, so every LoRA import download failed. Follow a redirect only over HTTPS to civitai.com or a subdomain of it (parsed host name, default port, no credentials in the URL), at most 3 hops, and never send the API key on a redirected request. Every other redirect is still refused; the size cap, timeout and image checks apply to the final response. Logged URLs drop their query, so signed storage URLs stay out of the log. The System import job's status line now counts failed downloads, and when every planned download fails it leads with a warning, which the Download errors outcome lists too. The legacy admin scan's events are unchanged.
2026-10-25 11:27:57 -07:00
parent 8252d9d226
commit 863d717435
5 changed files with 434 additions and 15 deletions
@@ -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
@@ -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)
@@ -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 "<unparsable URL>"
+    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


@@ -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."""

@@ -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):