fix: CivitAI downloads never follow redirects

py/lora_utils.py

@@ -307,6 +307,27 @@ MAX_CIVITAI_DOWNLOAD_BYTES = 50 * 1024 * 1024
 _DOWNLOAD_CHUNK_BYTES = 64 * 1024
 
 
+class _RefuseRedirects(urllib.request.HTTPRedirectHandler):
+    """Never follow a redirect: urllib would re-send the Bearer header to it."""
+
+    def redirect_request(self, req, fp, code, msg, headers, newurl):
+        logger.warning(f"Refusing redirect from {req.full_url} to {newurl}")
+        return None
+
+
+def _build_opener(*handlers) -> urllib.request.OpenerDirector:
+    """Build the outbound opener; extra handlers let tests serve canned responses."""
+    return urllib.request.build_opener(_RefuseRedirects, *handlers)
+
+
+_opener = _build_opener()
+
+
+def _open_url(req: urllib.request.Request, timeout: float = 10):
+    """Open ``req`` through the redirect-refusing opener."""
+    return _opener.open(req, timeout=timeout)
+
+
 def is_civitai_url(url: str) -> bool:
     """True when ``url`` is an HTTPS URL on an allow-listed CivitAI host."""
     try:
@@ -355,7 +376,7 @@ def _download_one(url: str, local_path: Path, api_key: str) -> Optional[str]:
         if api_key:
             headers["Authorization"] = f"Bearer {api_key}"
         req = urllib.request.Request(url, headers=headers)
-        with urllib.request.urlopen(req, timeout=10) as resp:
+        with _open_url(req, timeout=10) as resp:
             raw = _read_capped(resp, MAX_CIVITAI_DOWNLOAD_BYTES)
         if raw is None:
             logger.warning(f"Refusing oversized download: {url}")

tests/test_lora_utils.py

@@ -5,6 +5,7 @@ Tests metadata parsing, trigger word extraction, image URL extraction,
 directory detection, TriggerWordCache, and image download logic.
 """
 
+import email
 import json
 import os
 import shutil
@@ -12,6 +13,10 @@ import sys
 import tempfile
 import threading
 import unittest
+import urllib.parse
+import urllib.request
+import urllib.response
+from io import BytesIO
 from pathlib import Path
 from unittest.mock import patch
 
@@ -323,7 +328,7 @@ class TestDownloadOne(unittest.TestCase):
         self.target = self.tmpdir / "out.jpg"
 
     def _run(self, url, opener, api_key="secret-key"):
-        with patch("urllib.request.urlopen", opener):
+        with patch.object(lora_utils, "_open_url", opener):
             return lora_utils._download_one(url, self.target, api_key)
 
     def test_off_host_url_is_refused_without_a_request(self):
@@ -384,6 +389,68 @@ class TestDownloadOne(unittest.TestCase):
         self.assertFalse(self.target.exists())
 
 
+class _FakeHTTPSHandler(urllib.request.HTTPSHandler):
+    """Serves canned responses through a real urllib opener, recording requests.
+
+    CivitAI hosts answer with a 302 to ``redirect_to``; every other host
+    answers 200 with a PNG body.
+    """
+
+    def __init__(self, redirect_to):
+        super().__init__()
+        self.redirect_to = redirect_to
+        self.requests = []
+
+    def https_open(self, req):
+        self.requests.append(req)
+        host = (urllib.parse.urlsplit(req.full_url).hostname or "").lower()
+        if host in CIVITAI_HOSTS:
+            headers = email.message_from_string(f"Location: {self.redirect_to}\n")
+            return urllib.response.addinfourl(BytesIO(b""), headers, req.full_url, 302)
+        headers = email.message_from_string("Content-Type: image/png\n")
+        return urllib.response.addinfourl(
+            BytesIO(_png_bytes()), headers, req.full_url, 200
+        )
+
+
+class TestRedirectsAreRefused(unittest.TestCase):
+
+    def setUp(self):
+        self.tmpdir = Path(tempfile.mkdtemp())
+        self.addCleanup(shutil.rmtree, self.tmpdir, True)
+
+    def test_redirect_to_other_host_is_not_followed(self):
+        handler = _FakeHTTPSHandler("https://evil.example/steal")
+        opener = lora_utils._build_opener(handler)
+
+        with patch.object(lora_utils, "_opener", opener):
+            result = lora_utils._download_one(
+                "https://civitai.com/api/download/1", self.tmpdir / "x.jpg", "key"
+            )
+
+        self.assertIsNone(result)
+        self.assertEqual(
+            [r.full_url for r in handler.requests],
+            ["https://civitai.com/api/download/1"],
+        )
+        for req in handler.requests:
+            host = urllib.parse.urlsplit(req.full_url).hostname
+            if host not in CIVITAI_HOSTS:
+                self.assertIsNone(req.get_header("Authorization"))
+
+    def test_redirect_to_civitai_is_not_followed_either(self):
+        handler = _FakeHTTPSHandler("https://image.civitai.com/other.png")
+        opener = lora_utils._build_opener(handler)
+
+        with patch.object(lora_utils, "_opener", opener):
+            result = lora_utils._download_one(
+                "https://civitai.com/api/download/1", self.tmpdir / "x.jpg", "key"
+            )
+
+        self.assertIsNone(result)
+        self.assertEqual(len(handler.requests), 1)
+
+
 class TestDownloadCivitaiImages(unittest.TestCase):
 
     def setUp(self):
@@ -402,7 +469,7 @@ class TestDownloadCivitaiImages(unittest.TestCase):
             ]
         )
 
-        with patch("urllib.request.urlopen", opener):
+        with patch.object(lora_utils, "_open_url", opener):
             paths = download_civitai_images(
                 metadata, self.meta_path, self.tmpdir / "cache", "key"
             )
@@ -418,7 +485,7 @@ class TestDownloadCivitaiImages(unittest.TestCase):
         metadata = _make_metadata(images=[{"url": "https://civitai.com/a.png"}])
         cache = self.tmpdir / "cache"
 
-        with patch("urllib.request.urlopen", opener):
+        with patch.object(lora_utils, "_open_url", opener):
             first = download_civitai_images(metadata, self.meta_path, cache, "")
             second = download_civitai_images(metadata, self.meta_path, cache, "")