diff --git a/comfyui_to_python/generator/render.py b/comfyui_to_python/generator/render.py index 707bed7..ec33277 100644 --- a/comfyui_to_python/generator/render.py +++ b/comfyui_to_python/generator/render.py @@ -5,19 +5,7 @@ from typing import Any import black from ..node_runtime import import_custom_nodes -from .generated_helpers import ( - _find_from_extension_location, - _is_comfyui_directory, - _load_module, - _load_module_temp, - add_comfyui_directory_to_sys_path, - add_extra_model_paths, - bootstrap_comfyui_runtime, - cleanup_comfyui_runtime, - find_path, - get_comfyui_path, - get_value_at_index, -) +from . import generated_helpers from .model import GenerationPlan @@ -33,20 +21,11 @@ class WorkflowRenderer: {"workflow": plan.metadata_workflow_data} ) + # Auto-discover helpers from generated_helpers.__all__ so new helpers + # are picked up without manually updating this list. func_strings = [] - for func in [ - _load_module, - _load_module_temp, - _is_comfyui_directory, - _find_from_extension_location, - get_value_at_index, - get_comfyui_path, - find_path, - add_comfyui_directory_to_sys_path, - add_extra_model_paths, - bootstrap_comfyui_runtime, - cleanup_comfyui_runtime, - ]: + for name in generated_helpers.__all__: + func = getattr(generated_helpers, name) func_strings.append(f"\n{inspect.getsource(func)}") static_imports = [ diff --git a/comfyui_to_python/node_runtime.py b/comfyui_to_python/node_runtime.py index e9127ff..06c532d 100644 --- a/comfyui_to_python/node_runtime.py +++ b/comfyui_to_python/node_runtime.py @@ -18,9 +18,15 @@ def _is_comfyui_directory(path: str) -> bool: def _load_module(module_name: str, filepath: str) -> Any: """Load a Python module from an explicit file path, bypassing sys.path. - This eliminates ALL bare import shadowing attacks by loading modules + Significantly reduces bare import shadowing risk by loading modules from verified file paths instead of relying on sys.path resolution. + If exec_module() raises, the partially-loaded module is removed from + sys.modules so subsequent calls start fresh. """ + # Return cached module if already loaded — prevents re-execution + # that would reset state (e.g. NODE_CLASS_MAPPINGS after init_extra_nodes). + if module_name in sys.modules: + return sys.modules[module_name] if not os.path.isfile(filepath): log.debug("Module file not found: %s (%s)", module_name, filepath) return None @@ -31,7 +37,11 @@ def _load_module(module_name: str, filepath: str) -> Any: return None mod = importlib.util.module_from_spec(spec) sys.modules[module_name] = mod - spec.loader.exec_module(mod) + try: + spec.loader.exec_module(mod) + except BaseException: + sys.modules.pop(module_name, None) + raise return mod except Exception as e: log.debug("Failed to load %s from %s: %s", module_name, filepath, e) @@ -84,7 +94,10 @@ def get_comfyui_path() -> str | None: p = _find_from_extension_location() if p: return p - return find_path("ComfyUI", max_depth=20) + p = find_path("ComfyUI", max_depth=20) + if p and _is_comfyui_directory(p): + return p + return None def add_comfyui_directory_to_sys_path() -> None: @@ -267,10 +280,12 @@ def import_custom_nodes() -> None: original_sys_path = list(sys.path) sys.path[:] = [p for p in sys.path if p != comfy_subdir] - server_mod = _load_module("server", os.path.join(comfyui_path, "server.py")) - - # Restore sys.path so nodes and other modules that need comfy/ still work - sys.path[:] = original_sys_path + try: + server_mod = _load_module("server", os.path.join(comfyui_path, "server.py")) + finally: + # Restore sys.path so nodes and other modules that need comfy/ still work. + # Guaranteed even if _load_module raises mid-execution. + sys.path[:] = original_sys_path if execution_mod is None or server_mod is None: log.debug( @@ -292,14 +307,21 @@ def import_custom_nodes() -> None: def get_node_class_mappings() -> dict: - """Load ComfyUI node mappings on demand via _load_module().""" + """Load ComfyUI node mappings on demand via _load_module(). + + Reuses the cached "nodes" module from sys.modules if already loaded + (e.g. by import_custom_nodes) to avoid resetting NODE_CLASS_MAPPINGS. + """ add_comfyui_directory_to_sys_path() comfyui_path = get_comfyui_path() if comfyui_path is None: log.debug("get_node_class_mappings: ComfyUI path not found") return {} - nodes_mod = _load_module("nodes", os.path.join(comfyui_path, "nodes.py")) + # Reuse cached module if already loaded to avoid resetting mappings. + nodes_mod = sys.modules.get("nodes") + if nodes_mod is None: + nodes_mod = _load_module("nodes", os.path.join(comfyui_path, "nodes.py")) if nodes_mod is None: return {} return getattr(nodes_mod, "NODE_CLASS_MAPPINGS", {}) diff --git a/pyproject.toml b/pyproject.toml index d30b11f..ef0b50f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -9,9 +9,6 @@ dependencies = ["black"] [project.optional-dependencies] dev = ["ruff"] -[tool.uv] -dev-dependencies = ["ruff"] - [project.urls] Repository = "https://github.com/pydn/ComfyUI-to-Python-Extension" # Used by Comfy Registry https://comfyregistry.org @@ -20,5 +17,3 @@ Repository = "https://github.com/pydn/ComfyUI-to-Python-Extension" PublisherId = "pydn" DisplayName = "ComfyUI-to-Python-Extension" Icon = "" - - diff --git a/tests/test_import_path_resolution.py b/tests/test_import_path_resolution.py index aa06469..ff162bc 100644 --- a/tests/test_import_path_resolution.py +++ b/tests/test_import_path_resolution.py @@ -12,9 +12,7 @@ Covers Approach B from docs/specs/harden-import-path-resolution.html: import os import sys import tempfile -import types import unittest -from io import StringIO from unittest.mock import patch # Real ComfyUI checkout for integration-style unit tests (parent of this repo) @@ -192,7 +190,9 @@ class TestGeneratedScriptIsolation(unittest.TestCase): renderer = WorkflowRenderer() generated = renderer.render(plan) - self.assertIn("_load_module", generated, "Generated script must embed _load_module()") + self.assertIn( + "_load_module", generated, "Generated script must embed _load_module()" + ) self.assertIn( "spec_from_file_location", generated, @@ -283,5 +283,182 @@ class TestIsComfyuiDirectory(unittest.TestCase): self.assertFalse(_is_comfyui_directory("/nonexistent/path/to/comfyui")) +class TestLoadModuleFailureCleanup(unittest.TestCase): + """Tests for _load_module() behavior when exec_module() raises.""" + + def tearDown(self): + mod_names_to_clean = [ + name for name in sys.modules if name.startswith("_test_fail_") + ] + for name in mod_names_to_clean: + del sys.modules[name] + + def test_broken_module_removed_from_sys_modules_on_failure(self): + """Given a module that raises during exec, it is removed from sys.modules.""" + from comfyui_to_python.node_runtime import _load_module + + with tempfile.TemporaryDirectory() as tmpdir: + mod_file = os.path.join(tmpdir, "badmodule.py") + with open(mod_file, "w") as f: + f.write("raise RuntimeError('mid-execution failure')\n") + + result = _load_module("_test_fail_bad", mod_file) + self.assertIsNone(result) + self.assertNotIn( + "_test_fail_bad", + sys.modules, + "Failed module must not remain cached in sys.modules", + ) + + def test_broken_module_retried_cleanly(self): + """Given a failed load, a subsequent load attempt starts fresh (not from cache).""" + from comfyui_to_python.node_runtime import _load_module + + with tempfile.TemporaryDirectory() as tmpdir: + bad_file = os.path.join(tmpdir, "bad.py") + good_file = os.path.join(tmpdir, "good.py") + + # First load: bad module that raises + with open(bad_file, "w") as f: + f.write("raise RuntimeError('fail')\n") + result_bad = _load_module("_test_fail_retry", bad_file) + self.assertIsNone(result_bad) + + # Overwrite with good module + os.remove(bad_file) + with open(good_file, "w") as f: + f.write("VALUE = 99\n") + + # Reload same name from a different (working) file + result_good = _load_module("_test_fail_retry", good_file) + self.assertIsNotNone(result_good) + self.assertEqual(result_good.VALUE, 99) + + +class TestLoadModuleTemp(unittest.TestCase): + """Tests for _load_module_temp() removal from sys.modules.""" + + def tearDown(self): + mod_names_to_clean = [ + name for name in sys.modules if name.startswith("_test_temp_") + ] + for name in mod_names_to_clean: + del sys.modules[name] + + def test_module_removed_from_sys_modules_after_temp_load(self): + """After _load_module_temp(), the module key is absent from sys.modules.""" + from comfyui_to_python.node_runtime import _load_module_temp + + with tempfile.TemporaryDirectory() as tmpdir: + mod_file = os.path.join(tmpdir, "tempmod.py") + with open(mod_file, "w") as f: + f.write("VALUE = 123\n") + + mod = _load_module_temp("_test_temp_mod", mod_file) + self.assertIsNotNone(mod) + self.assertEqual(mod.VALUE, 123) + self.assertNotIn( + "_test_temp_mod", + sys.modules, + "Temp module must be removed from sys.modules after load", + ) + + def test_module_removed_on_temp_load_failure(self): + """After _load_module_temp() fails, the module key is absent from sys.modules.""" + from comfyui_to_python.node_runtime import _load_module_temp + + with tempfile.TemporaryDirectory() as tmpdir: + mod_file = os.path.join(tmpdir, "badtemp.py") + with open(mod_file, "w") as f: + f.write("raise RuntimeError('fail')\n") + + result = _load_module_temp("_test_temp_bad", mod_file) + self.assertIsNone(result) + self.assertNotIn( + "_test_temp_bad", + sys.modules, + "Failed temp module must be removed from sys.modules", + ) + + +class TestRepeatedNodeLoad(unittest.TestCase): + """Tests that repeated _load_module('nodes', ...) calls respect the cache.""" + + def tearDown(self): + if "_test_nodes" in sys.modules: + del sys.modules["_test_nodes"] + + def test_repeated_load_preserves_cached_attributes(self): + """Second _load_module() for same name returns cached module, not a fresh copy.""" + from comfyui_to_python.node_runtime import _load_module + + with tempfile.TemporaryDirectory() as tmpdir: + mod_file = os.path.join(tmpdir, "mynodes.py") + with open(mod_file, "w") as f: + f.write("NODE_CLASS_MAPPINGS = {'A': 1}\n") + + first = _load_module("_test_nodes", mod_file) + self.assertIsNotNone(first) + + # Mutate the cached module to simulate init_extra_nodes() + first.NODE_CLASS_MAPPINGS["B"] = 2 + + # Second load with same name must return cached copy, preserving mutation + second = _load_module("_test_nodes", mod_file) + self.assertIs(first, second) + self.assertEqual(second.NODE_CLASS_MAPPINGS.get("B"), 2) + + +class TestGetComfyuiPathThirdStrategy(unittest.TestCase): + """Tests for get_comfyui_path() third-strategy verification.""" + + def test_fallback_rejects_non_comfyui_directory(self): + """Given a directory named ComfyUI without nodes.py, third strategy rejects it.""" + from comfyui_to_python.node_runtime import get_comfyui_path + + with tempfile.TemporaryDirectory() as tmpdir: + # Create a fake "ComfyUI" directory (no nodes.py) + fake_dir = os.path.join(tmpdir, "ComfyUI") + os.makedirs(fake_dir) + + original_cwd = os.getcwd() + try: + os.chdir(tmpdir) + with patch.dict("os.environ", {}, clear=False): + # Strip COMFYUI_PATH so fallback strategies activate + env_copy = dict(os.environ) + env_copy.pop("COMFYUI_PATH", None) + with patch.dict("os.environ", env_copy, clear=True): + result = get_comfyui_path() + finally: + os.chdir(original_cwd) + + # Must NOT return the fake directory + self.assertNotEqual(result, fake_dir) + + +class TestImportCustomNodesSysPathRestoration(unittest.TestCase): + """Tests for sys.path restoration in import_custom_nodes().""" + + def test_sys_path_restored_on_crash(self): + """Given a crash during import_custom_nodes(), sys.path is still restored.""" + from comfyui_to_python.node_runtime import import_custom_nodes + + original_path = sys.path[:] + + with patch( + "comfyui_to_python.node_runtime._load_module", + side_effect=RuntimeError("simulated crash during server load"), + ): + try: + import_custom_nodes() + except RuntimeError: + pass # Expected crash + + # sys.path should still contain all original entries (no dangling state) + for entry in original_path: + self.assertIn(entry, sys.path) + + if __name__ == "__main__": unittest.main() diff --git a/uv.lock b/uv.lock index 5e153fc..a9c9abf 100644 --- a/uv.lock +++ b/uv.lock @@ -68,11 +68,6 @@ dev = [ { name = "ruff" }, ] -[package.dev-dependencies] -dev = [ - { name = "ruff" }, -] - [package.metadata] requires-dist = [ { name = "black" }, @@ -80,9 +75,6 @@ requires-dist = [ ] provides-extras = ["dev"] -[package.metadata.requires-dev] -dev = [{ name = "ruff" }] - [[package]] name = "mypy-extensions" version = "1.1.0"