fix: CivitAI key is never returned or sent off-host

py/api/lora_integration.py

@@ -76,7 +76,7 @@ class LoraIntegrationMixin:
                     "enabled": config["enabled"],
                     "path": config["path"],
                     "trigger_words_enabled": config["trigger_words_enabled"],
-                    "civitai_api_key": config.get("civitai_api_key", ""),
+                    "has_civitai_api_key": bool(config.get("civitai_api_key")),
                     "detected": detected_path is not None,
                     "detected_path": detected_path or "",
                     "trigger_cache_loaded": cache.is_loaded,

py/lora_utils.py

@@ -15,6 +15,7 @@ import threading
 import urllib.request
 from pathlib import Path
 from typing import Dict, List, Optional, Tuple
+from urllib.parse import urlsplit
 
 try:
     from ..utils.logging_config import get_logger
@@ -298,16 +299,67 @@ def get_preview_image_from_metadata(
 
 _THUMB_MAX_SIZE = 512
 
+# 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"}
+)
+MAX_CIVITAI_DOWNLOAD_BYTES = 50 * 1024 * 1024
+_DOWNLOAD_CHUNK_BYTES = 64 * 1024
+
+
+def is_civitai_url(url: str) -> bool:
+    """True when ``url`` is an HTTPS URL on an allow-listed CivitAI host."""
+    try:
+        parts = urlsplit(url)
+    except ValueError:
+        return False
+    hostname = (parts.hostname or "").lower()
+    return parts.scheme == "https" and hostname in CIVITAI_HOSTS
+
+
+def _read_capped(resp, cap: int) -> Optional[bytes]:
+    """Read a response body, refusing anything larger than ``cap`` bytes."""
+    declared = resp.headers.get("Content-Length")
+    if declared is not None:
+        try:
+            if int(declared) > cap:
+                return None
+        except ValueError:
+            pass
+
+    chunks = []
+    total = 0
+    while True:
+        chunk = resp.read(_DOWNLOAD_CHUNK_BYTES)
+        if not chunk:
+            break
+        total += len(chunk)
+        if total > cap:
+            return None
+        chunks.append(chunk)
+    return b"".join(chunks)
+
 
 def _download_one(url: str, local_path: Path, api_key: str) -> Optional[str]:
-    """Download a single image, resize to thumbnail, save as JPEG."""
+    """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.
+    """
+    if not is_civitai_url(url):
+        logger.warning(f"Refusing non-CivitAI download URL: {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 urllib.request.urlopen(req, timeout=10) as resp:
-            raw = resp.read()
+            raw = _read_capped(resp, MAX_CIVITAI_DOWNLOAD_BYTES)
+        if raw is None:
+            logger.warning(f"Refusing oversized download: {url}")
+            return None
 
         # Resize to thumbnail to save disk space
         from io import BytesIO
@@ -361,7 +413,7 @@ def download_civitai_images(
         if not isinstance(img, dict):
             continue
         url = img.get("url", "")
-        if not isinstance(url, str) or not url.startswith("http"):
+        if not isinstance(url, str) or not is_civitai_url(url):
             continue
 
         url_hash = hashlib.md5(url.encode()).hexdigest()[:12]

tests/test_api_lora.py

@@ -0,0 +1,119 @@
+"""Route tests for the LoraManager integration endpoints.
+
+Filesystem detection and network downloads are patched out so the tests run
+anywhere without a LoraManager install or a CivitAI connection.
+"""
+
+import json
+import os
+import shutil
+import sys
+import tempfile
+import types
+import unittest
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+if "folder_paths" not in sys.modules:
+    _stub = types.ModuleType("folder_paths")
+    _stub.models_dir = tempfile.gettempdir()
+    _stub.base_path = tempfile.gettempdir()
+    _stub.get_output_directory = lambda: tempfile.gettempdir()
+    _stub.get_folder_paths = lambda name: []
+    sys.modules["folder_paths"] = _stub
+
+if "server" not in sys.modules:
+    _server = MagicMock()
+    _server.PromptServer.instance.routes = MagicMock()
+    sys.modules["server"] = _server
+
+from aiohttp import web
+from aiohttp.test_utils import AioHTTPTestCase
+
+from database.operations import PromptDatabase
+from py.api import PromptManagerAPI
+from py.config import IntegrationConfig
+from py.lora_utils import get_trigger_cache
+
+SECRET_KEY = "sk-civitai-super-secret-0123456789"
+
+
+class LoraAPITestCase(AioHTTPTestCase):
+    """Stands up the PromptManager routes with a temp DB and neutral LoRA config."""
+
+    async def get_application(self):
+        self._temp_db = tempfile.NamedTemporaryFile(delete=False, suffix=".db")
+        self._temp_db.close()
+        self.tmp_root = Path(tempfile.mkdtemp())
+
+        self._config_patch = patch.multiple(
+            IntegrationConfig,
+            LORA_MANAGER_ENABLED=False,
+            LORA_MANAGER_PATH="",
+            LORA_TRIGGER_WORDS_ENABLED=False,
+            CIVITAI_API_KEY="",
+        )
+        self._config_patch.start()
+        self._detect_patch = patch(
+            "py.lora_utils.detect_lora_manager", return_value=None
+        )
+        self.detect_mock = self._detect_patch.start()
+        get_trigger_cache().clear()
+
+        app = web.Application()
+        routes = web.RouteTableDef()
+        self.api = PromptManagerAPI()
+        self.api.db = PromptDatabase(self._temp_db.name)
+        self.api.add_routes(routes)
+        app.router.add_routes(routes)
+        return app
+
+    async def tearDownAsync(self):
+        self._detect_patch.stop()
+        self._config_patch.stop()
+        get_trigger_cache().clear()
+        for path in (
+            self._temp_db.name,
+            self._temp_db.name + "-wal",
+            self._temp_db.name + "-shm",
+        ):
+            if os.path.exists(path):
+                os.unlink(path)
+        shutil.rmtree(self.tmp_root, ignore_errors=True)
+
+    async def _post_json(self, url, body):
+        return await self.client.request(
+            "POST",
+            url,
+            data=json.dumps(body),
+            headers={"Content-Type": "application/json"},
+        )
+
+
+class TestLoraStatusNeverReturnsKey(LoraAPITestCase):
+
+    async def test_status_reports_flag_without_the_key(self):
+        IntegrationConfig.CIVITAI_API_KEY = SECRET_KEY
+
+        resp = await self.client.request("GET", "/prompt_manager/lora/status")
+
+        self.assertEqual(resp.status, 200)
+        body = await resp.text()
+        data = json.loads(body)
+        self.assertTrue(data["success"])
+        self.assertIs(data["has_civitai_api_key"], True)
+        self.assertNotIn("civitai_api_key", data)
+        self.assertNotIn(SECRET_KEY, body)
+
+    async def test_status_flag_is_false_without_a_key(self):
+        resp = await self.client.request("GET", "/prompt_manager/lora/status")
+
+        data = await resp.json()
+        self.assertIs(data["has_civitai_api_key"], False)
+        self.assertNotIn("civitai_api_key", data)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_lora_utils.py

@@ -7,6 +7,7 @@ directory detection, TriggerWordCache, and image download logic.
 
 import json
 import os
+import shutil
 import sys
 import tempfile
 import threading
@@ -16,12 +17,17 @@ from unittest.mock import patch
 
 sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 
+from py import lora_utils
 from py.lora_utils import (
+    CIVITAI_HOSTS,
+    MAX_CIVITAI_DOWNLOAD_BYTES,
     TriggerWordCache,
+    download_civitai_images,
     get_civitai_image_urls,
     get_example_prompt_from_metadata,
     get_lora_image_cache_dir,
     get_trigger_words_from_metadata,
+    is_civitai_url,
     read_lora_metadata,
 )
 
@@ -224,6 +230,207 @@ class TestGetLoraImageCacheDir(unittest.TestCase):
         self.assertTrue(cache_dir.is_dir())
 
 
+# ── CivitAI download tests (no network) ───────────────────────────────
+
+
+def _png_bytes():
+    """A valid 2x2 PNG so PIL can decode the fake download."""
+    from io import BytesIO
+
+    from PIL import Image
+
+    buf = BytesIO()
+    Image.new("RGB", (2, 2), (10, 20, 30)).save(buf, "PNG")
+    return buf.getvalue()
+
+
+class _FakeResponse:
+    """Minimal stand-in for the object returned by urllib.request.urlopen."""
+
+    def __init__(self, body, content_length=None):
+        self._body = body
+        self._pos = 0
+        self.headers = {}
+        if content_length is not None:
+            self.headers["Content-Length"] = str(content_length)
+
+    def read(self, amt=-1):
+        if amt is None or amt < 0:
+            chunk = self._body[self._pos :]
+            self._pos = len(self._body)
+            return chunk
+        chunk = self._body[self._pos : self._pos + amt]
+        self._pos += len(chunk)
+        return chunk
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *exc):
+        return False
+
+
+class _FakeOpener:
+    """Records every request and serves a canned response."""
+
+    def __init__(self, response):
+        self.response = response
+        self.requests = []
+
+    def __call__(self, req, timeout=None):
+        self.requests.append(req)
+        return self.response
+
+
+class TestCivitaiUrlPolicy(unittest.TestCase):
+
+    def test_allow_list_contains_expected_hosts(self):
+        self.assertEqual(
+            CIVITAI_HOSTS,
+            frozenset(
+                {
+                    "civitai.com",
+                    "www.civitai.com",
+                    "api.civitai.com",
+                    "image.civitai.com",
+                }
+            ),
+        )
+        self.assertEqual(MAX_CIVITAI_DOWNLOAD_BYTES, 50 * 1024 * 1024)
+
+    def test_https_civitai_hosts_are_allowed(self):
+        for host in CIVITAI_HOSTS:
+            self.assertTrue(is_civitai_url(f"https://{host}/api/download/1"), host)
+
+    def test_other_hosts_and_schemes_are_refused(self):
+        for url in (
+            "https://evil.example/x",
+            "http://civitai.com/api/download/1",
+            "https://civitai.com.evil.example/x",
+            "https://notcivitai.com/x",
+            "ftp://civitai.com/x",
+            "not a url",
+            "",
+        ):
+            self.assertFalse(is_civitai_url(url), url)
+
+
+class TestDownloadOne(unittest.TestCase):
+
+    def setUp(self):
+        self.tmpdir = Path(tempfile.mkdtemp())
+        self.addCleanup(shutil.rmtree, self.tmpdir, True)
+        self.target = self.tmpdir / "out.jpg"
+
+    def _run(self, url, opener, api_key="secret-key"):
+        with patch("urllib.request.urlopen", opener):
+            return lora_utils._download_one(url, self.target, api_key)
+
+    def test_off_host_url_is_refused_without_a_request(self):
+        opener = _FakeOpener(_FakeResponse(_png_bytes()))
+
+        result = self._run("https://evil.example/x", opener)
+
+        self.assertIsNone(result)
+        self.assertEqual(opener.requests, [])
+        self.assertFalse(self.target.exists())
+
+    def test_plain_http_civitai_is_refused(self):
+        opener = _FakeOpener(_FakeResponse(_png_bytes()))
+
+        result = self._run("http://civitai.com/api/download/1", opener)
+
+        self.assertIsNone(result)
+        self.assertEqual(opener.requests, [])
+
+    def test_civitai_download_sends_bearer_header(self):
+        opener = _FakeOpener(_FakeResponse(_png_bytes()))
+
+        result = self._run("https://civitai.com/api/download/1", opener)
+
+        self.assertEqual(os.path.realpath(result), os.path.realpath(str(self.target)))
+        self.assertEqual(len(opener.requests), 1)
+        req = opener.requests[0]
+        self.assertEqual(req.get_header("Authorization"), "Bearer secret-key")
+        self.assertTrue(self.target.exists())
+
+    def test_no_bearer_header_without_a_key(self):
+        opener = _FakeOpener(_FakeResponse(_png_bytes()))
+
+        self._run("https://image.civitai.com/a.png", opener, api_key="")
+
+        self.assertIsNone(opener.requests[0].get_header("Authorization"))
+
+    def test_oversized_content_length_is_refused(self):
+        opener = _FakeOpener(
+            _FakeResponse(_png_bytes(), content_length=MAX_CIVITAI_DOWNLOAD_BYTES + 1)
+        )
+
+        result = self._run("https://civitai.com/api/download/1", opener)
+
+        self.assertIsNone(result)
+        self.assertFalse(self.target.exists())
+
+    def test_reading_stops_past_cap_when_length_is_absent(self):
+        body = b"x" * (MAX_CIVITAI_DOWNLOAD_BYTES + 1)
+        response = _FakeResponse(body)
+        opener = _FakeOpener(response)
+
+        with patch.object(lora_utils, "MAX_CIVITAI_DOWNLOAD_BYTES", 1024):
+            result = self._run("https://civitai.com/api/download/1", opener)
+
+        self.assertIsNone(result)
+        self.assertLessEqual(response._pos, 1024 + lora_utils._DOWNLOAD_CHUNK_BYTES)
+        self.assertFalse(self.target.exists())
+
+
+class TestDownloadCivitaiImages(unittest.TestCase):
+
+    def setUp(self):
+        self.tmpdir = Path(tempfile.mkdtemp())
+        self.addCleanup(shutil.rmtree, self.tmpdir, True)
+        self.meta_path = self.tmpdir / "lora.safetensors.metadata.json"
+        self.meta_path.write_text("{}")
+
+    def test_only_civitai_urls_are_fetched(self):
+        opener = _FakeOpener(_FakeResponse(_png_bytes()))
+        metadata = _make_metadata(
+            images=[
+                {"url": "https://evil.example/steal"},
+                {"url": "https://civitai.com/api/download/1"},
+                {"url": "not-a-url"},
+            ]
+        )
+
+        with patch("urllib.request.urlopen", opener):
+            paths = download_civitai_images(
+                metadata, self.meta_path, self.tmpdir / "cache", "key"
+            )
+
+        self.assertEqual(len(paths), 1)
+        self.assertEqual(
+            [r.full_url for r in opener.requests],
+            ["https://civitai.com/api/download/1"],
+        )
+
+    def test_cached_files_are_not_refetched(self):
+        opener = _FakeOpener(_FakeResponse(_png_bytes()))
+        metadata = _make_metadata(images=[{"url": "https://civitai.com/a.png"}])
+        cache = self.tmpdir / "cache"
+
+        with patch("urllib.request.urlopen", opener):
+            first = download_civitai_images(metadata, self.meta_path, cache, "")
+            second = download_civitai_images(metadata, self.meta_path, cache, "")
+
+        self.assertEqual(first, second)
+        self.assertEqual(len(opener.requests), 1)
+
+    def test_no_images_returns_empty(self):
+        self.assertEqual(
+            download_civitai_images({}, self.meta_path, self.tmpdir / "cache"), []
+        )
+
+
 # ── TriggerWordCache tests ────────────────────────────────────────────