v4: Check model installs against the model list, including the url (#3506)

A model install is limited to the entries of the model list. The legacy
server compared save_path, base and filename with the list, but not the
url the file is downloaded from, so a request with a listed name and
another url was queued as given. The glob server lost the list check,
and the check for formats other than .safetensors, when its routes
moved to data models (0be0a2e6), so it queued and downloaded whatever a
request named.

Compare the url in both servers; for the legacy server this is the port
of #3428. In the glob server, run both checks in the worker, which
queue/install_model and queue/task both reach. A default channel list
without models counts as empty there, so the model is refused and the
worker keeps running.

CHANGELOG.md

@@ -40,6 +40,13 @@ and this project adheres to [Semantic Versioning 2.0.0](https://semver.org/spec/
   environment variable at a CA bundle that includes the proxy's CA, or set
   `bypass_ssl = True` in `config.ini`, which turns verification off for this
   fetch too (#3302).
+- Check every model install against the model list, including the download URL:
+  the folder, base, file name, and URL must match an entry. With the legacy UI,
+  a request with a listed name and another URL was accepted. Without it, the
+  list was not checked, and neither was the rule that a model in a format other
+  than `.safetensors` must be in the default channel's list unless the security
+  level is `normal-` or lower on a loopback listener, or `weak` in
+  `personal_cloud` mode (#3428).
 
 ### Fixed
 

comfyui_manager/glob/manager_server.py

@@ -87,6 +87,7 @@ from ..common.security_messages import (
     SECURITY_MESSAGE_MIDDLE,
     SECURITY_MESSAGE_MIDDLE_P,
     SECURITY_MESSAGE_HIGH_P,
+    SECURITY_MESSAGE_NORMAL_MINUS_MODEL,
 )
 from .constants import (
     model_dir_name_map,
@@ -1060,6 +1061,18 @@ async def task_worker():
 
         json_data = params.model_dump()
 
+        # Checked here, not in a route: queue/install_model and queue/task both end up here.
+        if not await model_utils.check_whitelist_for_model(json_data):
+            logging.error(f"[ComfyUI-Manager] Invalid model install request is detected: {json_data}")
+            return "Invalid model install request is detected"
+
+        if not json_data["filename"].endswith(".safetensors") and not security_utils.is_allowed_security_level("high+"):
+            models_json = await core.get_data_by_mode(ManagerDatabaseSource.cache.value, "model-list.json", "default")
+
+            if not any(x.get("url") == json_data["url"] for x in models_json.get("models", [])):
+                logging.error(SECURITY_MESSAGE_NORMAL_MINUS_MODEL)
+                return "A security error has occurred. Please check the terminal logs"
+
         model_path = model_utils.get_model_path(json_data)
         model_url = json_data.get("url")
 

comfyui_manager/glob/utils/model_utils.py

@@ -145,6 +145,7 @@ async def check_whitelist_for_model(item):
             x["save_path"] == item["save_path"]
             and x["base"] == item["base"]
             and x["filename"] == item["filename"]
+            and x.get("url") == item.get("url")
         ):
             return True
 
@@ -155,6 +156,7 @@ async def check_whitelist_for_model(item):
             x["save_path"] == item["save_path"]
             and x["base"] == item["base"]
             and x["filename"] == item["filename"]
+            and x.get("url") == item.get("url")
         ):
             return True
 

comfyui_manager/legacy/manager_server.py

@@ -1713,13 +1713,13 @@ async def check_whitelist_for_model(item):
     json_obj = await core.get_data_by_mode('cache', 'model-list.json')
 
     for x in json_obj.get('models', []):
-        if x['save_path'] == item['save_path'] and x['base'] == item['base'] and x['filename'] == item['filename']:
+        if x['save_path'] == item['save_path'] and x['base'] == item['base'] and x['filename'] == item['filename'] and x.get('url') == item.get('url'):
             return True
 
     json_obj = await core.get_data_by_mode('local', 'model-list.json')
 
     for x in json_obj.get('models', []):
-        if x['save_path'] == item['save_path'] and x['base'] == item['base'] and x['filename'] == item['filename']:
+        if x['save_path'] == item['save_path'] and x['base'] == item['base'] and x['filename'] == item['filename'] and x.get('url') == item.get('url'):
             return True
         
     return False

tests/e2e/test_e2e_task_operations.py

@@ -643,10 +643,10 @@ class TestInstallModel:
                   advances beyond baseline within a bounded timeout)
               (c) optionally, queue/history records a row matching this ui_id
             Download completion to a real model file is NOT an E2E responsibility —
-            it depends on external network + the target URL. The test URL points
-            to example.com, which will not yield a valid safetensors file; the
-            task will still transition through running → history (as failed /
-            trivial-success), which is the E2E-observable we need.
+            it depends on external network + the target URL. The test model is
+            not in the model list, so the worker refuses it without a download;
+            the task still transitions through running → history (as failed),
+            which is the E2E-observable we need.
         """
         _reset_queue()
 

tests/test_glob_model_install_whitelist.py

@@ -0,0 +1,110 @@
+"""The glob server installs only models that are in the model list.
+
+Each case hands a model to the real `do_install_model`, the worker step that
+queue/install_model and queue/task both reach, with the real
+`check_whitelist_for_model`. The downloaders are replaced by a recorder, so
+nothing is fetched.
+"""
+import ast
+import asyncio
+import importlib.util
+import logging
+import sys
+from types import SimpleNamespace
+
+import pytest
+
+from legacy_html_testutil import PACKAGE, load_functions
+
+SAFETENSORS = {"name": "a", "type": "checkpoint", "base": "SDXL", "save_path": "checkpoints",
+               "filename": "a.safetensors", "url": "https://huggingface.co/a/resolve/main/a.safetensors"}
+PICKLE = {"name": "b", "type": "upscale", "base": "upscale", "save_path": "upscale_models",
+          "filename": "b.pth", "url": "https://huggingface.co/b/resolve/main/b.pth"}
+OTHER_URL = "http://other.example/x.bin"
+
+
+def load_data_models(monkeypatch):
+    spec = importlib.util.spec_from_file_location("_generated_models", PACKAGE / "data_models/generated_models.py")
+    module = importlib.util.module_from_spec(spec)
+    monkeypatch.setitem(sys.modules, spec.name, module)
+    spec.loader.exec_module(module)
+    return module
+
+
+def install(monkeypatch, item, channel_models, default_models, local_models=(), high_plus=False):
+    """What the worker answers for `item`, and the downloads it started.
+
+    The three arguments are the model lists of the configured channel, of the default
+    channel, and of the local file; None stands for a list file without `models`.
+    """
+    models = load_data_models(monkeypatch)
+    monkeypatch.setitem(sys.modules, "comfyui_manager.data_models", models)
+
+    async def get_data_by_mode(mode, filename, channel_url=None):
+        listed = {"cache": default_models if channel_url == "default" else channel_models, "local": local_models}[mode]
+        return {} if listed is None else {"models": list(listed)}
+
+    downloads = []
+    core = SimpleNamespace(get_data_by_mode=get_data_by_mode, get_config=lambda: {"model_download_by_agent": True})
+    whitelist = load_functions(PACKAGE / "glob/utils/model_utils.py", {"check_whitelist_for_model"}, {"core": core})
+    namespace = {
+        "logging": logging, "core": core,
+        "ModelMetadata": models.ModelMetadata, "OperationResult": models.OperationResult,
+        "ManagerDatabaseSource": models.ManagerDatabaseSource,
+        "SECURITY_MESSAGE_NORMAL_MINUS_MODEL": "",
+        "security_utils": SimpleNamespace(is_allowed_security_level=lambda level: level != "high+" or high_plus),
+        "model_utils": SimpleNamespace(check_whitelist_for_model=whitelist["check_whitelist_for_model"],
+                                       get_model_path=lambda data: f"{data['save_path']}/{data['filename']}"),
+        "download_url_with_agent": lambda url, path: downloads.append((url, path)) or True,
+    }
+    path = PACKAGE / "glob/manager_server.py"
+    nodes = [node for node in ast.walk(ast.parse(path.read_text(encoding="utf-8")))
+             if isinstance(node, ast.AsyncFunctionDef) and node.name == "do_install_model"]
+    assert len(nodes) == 1
+    exec(compile(ast.Module(nodes, []), str(path), "exec"), namespace)
+
+    result = asyncio.run(namespace["do_install_model"](models.ModelMetadata.model_validate(item)))
+    return result, downloads
+
+
+def test_listed_model_is_downloaded_from_its_listed_url(monkeypatch):
+    result, downloads = install(monkeypatch, PICKLE, [SAFETENSORS, PICKLE], [SAFETENSORS, PICKLE])
+    assert result == "success"
+    assert downloads == [(PICKLE["url"], "upscale_models/b.pth")]
+
+
+@pytest.mark.parametrize("change", [
+    {"url": OTHER_URL},
+    {"save_path": "custom_nodes/SomePack"},
+    {"filename": "__init__.py"},
+], ids=["other url", "other folder", "other filename"])
+def test_unlisted_model_is_not_downloaded(monkeypatch, change):
+    result, downloads = install(monkeypatch, {**SAFETENSORS, **change}, [SAFETENSORS, PICKLE], [SAFETENSORS, PICKLE])
+    assert result != "success"
+    assert downloads == []
+
+
+@pytest.mark.parametrize("high_plus, expected", [(False, []), (True, [(PICKLE["url"], "upscale_models/b.pth")])],
+                         ids=["high+ not allowed", "high+ allowed"])
+def test_pickle_model_outside_the_default_channel_needs_high_plus(monkeypatch, high_plus, expected):
+    result, downloads = install(monkeypatch, PICKLE, [SAFETENSORS, PICKLE], [SAFETENSORS], high_plus=high_plus)
+    assert (result == "success") == bool(expected)
+    assert downloads == expected
+
+
+def test_pickle_model_is_refused_when_the_default_list_has_no_models(monkeypatch):
+    result, downloads = install(monkeypatch, PICKLE, [SAFETENSORS, PICKLE], None)
+    assert isinstance(result, str) and result != "success"
+    assert downloads == []
+
+
+def test_model_listed_only_in_the_local_file_is_downloaded(monkeypatch):
+    result, downloads = install(monkeypatch, SAFETENSORS, [], [], local_models=[SAFETENSORS])
+    assert result == "success"
+    assert downloads == [(SAFETENSORS["url"], "checkpoints/a.safetensors")]
+
+
+def test_safetensors_model_outside_the_default_channel_is_downloaded(monkeypatch):
+    result, downloads = install(monkeypatch, SAFETENSORS, [SAFETENSORS], [])
+    assert result == "success"
+    assert downloads == [(SAFETENSORS["url"], "checkpoints/a.safetensors")]

tests/test_legacy_model_whitelist.py

@@ -0,0 +1,43 @@
+"""The legacy server installs a listed model only from its listed url (#3428).
+
+Each case hands a request to the real `_install_model` with the real
+`check_whitelist_for_model`; the route only queues, so nothing is fetched.
+"""
+import asyncio
+import logging
+from types import SimpleNamespace
+
+import pytest
+from aiohttp import web
+
+from legacy_html_testutil import PACKAGE, load_functions
+
+LISTED = {"name": "a", "type": "checkpoint", "base": "SDXL", "save_path": "checkpoints",
+          "filename": "a.safetensors", "url": "https://huggingface.co/a/resolve/main/a.safetensors"}
+OTHER = {**LISTED, "name": "b", "filename": "b.safetensors", "url": "https://huggingface.co/b/resolve/main/b.safetensors"}
+
+
+def install(item):
+    """The status `_install_model` answers for `item`, and what it queued."""
+    async def get_data_by_mode(mode, filename, channel_url=None):
+        return {"models": [LISTED, OTHER]}
+
+    queued = []
+    functions = load_functions(PACKAGE / "legacy/manager_server.py", {"check_whitelist_for_model", "_install_model"}, {
+        "web": web, "logging": logging, "core": SimpleNamespace(get_data_by_mode=get_data_by_mode),
+        "is_allowed_security_level": lambda level: level == "middle+", "temp_queue_batch": queued,
+    })
+    return asyncio.run(functions["_install_model"](item)).status, queued
+
+
+def test_listed_model_is_queued_with_its_listed_url():
+    status, queued = install({**LISTED, "ui_id": "u"})
+    assert status == 200
+    assert [(kind, item["url"]) for kind, (_, item) in queued] == [("install-model", LISTED["url"])]
+
+
+@pytest.mark.parametrize("url", [{"url": OTHER["url"]}, {}], ids=["url of another entry", "no url"])
+def test_listed_name_with_another_url_is_refused(url):
+    status, queued = install({**{k: v for k, v in LISTED.items() if k != "url"}, **url, "ui_id": "u"})
+    assert status == 400
+    assert queued == []