Compare commits
1
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
0aab817a6a |
@@ -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.
|
||||
|
||||
|
||||
+8
-1
@@ -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
|
||||
|
||||
+96
-44
@@ -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-"
|
||||
|
||||
+310
-6
@@ -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):
|
||||
|
||||
Reference in New Issue
Block a user