diff --git a/py/lora_utils.py b/py/lora_utils.py index f23cc07..0a20330 100644 --- a/py/lora_utils.py +++ b/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}") diff --git a/tests/test_lora_utils.py b/tests/test_lora_utils.py index ed22377..c713a04 100644 --- a/tests/test_lora_utils.py +++ b/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, "")