diff --git a/CHANGELOG.md b/CHANGELOG.md index 9c80e6d5..32f44306 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/comfyui_manager/glob/manager_server.py b/comfyui_manager/glob/manager_server.py index 37d87095..0c2e0a5a 100644 --- a/comfyui_manager/glob/manager_server.py +++ b/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") diff --git a/comfyui_manager/glob/utils/model_utils.py b/comfyui_manager/glob/utils/model_utils.py index 2225ec07..6b784f95 100644 --- a/comfyui_manager/glob/utils/model_utils.py +++ b/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 diff --git a/comfyui_manager/legacy/manager_server.py b/comfyui_manager/legacy/manager_server.py index 7da62e73..625f15f5 100644 --- a/comfyui_manager/legacy/manager_server.py +++ b/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 diff --git a/tests/e2e/test_e2e_task_operations.py b/tests/e2e/test_e2e_task_operations.py index 41161f5f..c6c3ebf8 100644 --- a/tests/e2e/test_e2e_task_operations.py +++ b/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() diff --git a/tests/test_glob_model_install_whitelist.py b/tests/test_glob_model_install_whitelist.py new file mode 100644 index 00000000..8e7566c9 --- /dev/null +++ b/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")] diff --git a/tests/test_legacy_model_whitelist.py b/tests/test_legacy_model_whitelist.py new file mode 100644 index 00000000..ec97672f --- /dev/null +++ b/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 == []