civitai.red and civitai.green (and their subdomains) are now allowed redirect targets under the civitai.com rules: HTTPS only, default port, no credentials, matched on the parsed host name, at most 3 hops, never with the key. Their www/api/image hosts are also allowed starting hosts, whose first request carries the key as image.civitai.com does. Each failed download now logs one warning with the URL without query or credentials and a short reason: HTTP status, refused redirect, hop limit, size cap, decode or save error, timeout or network error.
This commit is contained in:
@@ -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