diff --git a/.gitignore b/.gitignore index 0f192fe..72f0fb7 100644 --- a/.gitignore +++ b/.gitignore @@ -149,4 +149,4 @@ cython_debug/ docs/ .agents/ .codex/ -AGENTS.md \ No newline at end of file +AGENTS.md diff --git a/comfyui_to_python/app.py b/comfyui_to_python/app.py index acee636..c05f27f 100644 --- a/comfyui_to_python/app.py +++ b/comfyui_to_python/app.py @@ -65,7 +65,9 @@ class ExportApplication: if nodes_mod is not None: fresh_mappings = getattr(nodes_mod, "NODE_CLASS_MAPPINGS", {}) self.node_class_mappings = fresh_mappings - self.base_node_class_mappings = copy.deepcopy(self.node_class_mappings) + # Leave self.base_node_class_mappings unchanged — it represents the + # pre-custom-node baseline used by WorkflowPlanner to decide whether + # a node gets a direct import or a NODE_CLASS_MAPPINGS dict lookup. load_order = LoadOrderDeterminer( data, self.node_class_mappings diff --git a/comfyui_to_python/generator/render.py b/comfyui_to_python/generator/render.py index ec33277..0c4f5d0 100644 --- a/comfyui_to_python/generator/render.py +++ b/comfyui_to_python/generator/render.py @@ -1,9 +1,12 @@ import inspect +import logging from pprint import pformat from typing import Any import black +log = logging.getLogger(__name__) + from ..node_runtime import import_custom_nodes from . import generated_helpers from .model import GenerationPlan @@ -25,7 +28,12 @@ class WorkflowRenderer: # are picked up without manually updating this list. func_strings = [] for name in generated_helpers.__all__: - func = getattr(generated_helpers, name) + func = getattr(generated_helpers, name, None) + if func is None: + log.warning( + "Helper '%s' missing from generated_helpers — skipping", name + ) + continue 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 06c532d..59dd728 100644 --- a/comfyui_to_python/node_runtime.py +++ b/comfyui_to_python/node_runtime.py @@ -9,10 +9,18 @@ log = logging.getLogger(__name__) def _is_comfyui_directory(path: str) -> bool: - """Verify a directory has ComfyUI structural markers (nodes.py).""" + """Verify a directory has ComfyUI structural markers. + + Checks for nodes.py, main.py, and the comfy/ subdirectory to raise + the bar against spoofing via a directory with only a single marker file. + """ if not os.path.isdir(path): return False - return os.path.isfile(os.path.join(path, "nodes.py")) + return ( + os.path.isfile(os.path.join(path, "nodes.py")) + and os.path.isfile(os.path.join(path, "main.py")) + and os.path.isdir(os.path.join(path, "comfy")) + ) def _load_module(module_name: str, filepath: str) -> Any: @@ -43,8 +51,9 @@ def _load_module(module_name: str, filepath: str) -> Any: sys.modules.pop(module_name, None) raise return mod - except Exception as e: + except BaseException as e: log.debug("Failed to load %s from %s: %s", module_name, filepath, e) + sys.modules.pop(module_name, None) # Also clean up on edge-case exceptions return None diff --git a/tests/runtime/generated/text-to-image.py b/tests/runtime/generated/text-to-image.py index cb109f6..cdf0821 100644 --- a/tests/runtime/generated/text-to-image.py +++ b/tests/runtime/generated/text-to-image.py @@ -10,47 +10,6 @@ from typing import Sequence, Mapping, Any, Union log = logging.getLogger(__name__) -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 - from verified file paths instead of relying on sys.path resolution. - """ - if not os.path.isfile(filepath): - log.debug("Module file not found: %s (%s)", module_name, filepath) - return None - try: - spec = importlib.util.spec_from_file_location(module_name, filepath) - if spec is None or spec.loader is None: - log.debug("Could not create spec for %s at %s", module_name, filepath) - return None - mod = importlib.util.module_from_spec(spec) - sys.modules[module_name] = mod - spec.loader.exec_module(mod) - return mod - except Exception as e: - log.debug("Failed to load %s from %s: %s", module_name, filepath, e) - return None - - -def _load_module_temp(module_name: str, filepath: str) -> Any: - """Load a module via _load_module() then remove it from sys.modules. - - Used during bootstrap for modules that ComfyUI's import chain also loads - normally — prevents the cached copy from conflicting with later imports. - """ - mod = _load_module(module_name, filepath) - sys.modules.pop(module_name, None) - return mod - - -def _is_comfyui_directory(path: str) -> bool: - """Verify a directory has ComfyUI structural markers (nodes.py).""" - if not os.path.isdir(path): - return False - return os.path.isfile(os.path.join(path, "nodes.py")) - - def _find_from_extension_location() -> str | None: """Walk up from this file's location to find ComfyUI root.""" ext_dir = os.path.dirname(os.path.realpath(__file__)) @@ -66,46 +25,64 @@ def _find_from_extension_location() -> str | None: return None -def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any: - """Return a sequence or mapping result item by index.""" +def _is_comfyui_directory(path: str) -> bool: + """Verify a directory has ComfyUI structural markers. + + Checks for nodes.py, main.py, and the comfy/ subdirectory to raise + the bar against spoofing via a directory with only a single marker file. + """ + if not os.path.isdir(path): + return False + return ( + os.path.isfile(os.path.join(path, "nodes.py")) + and os.path.isfile(os.path.join(path, "main.py")) + and os.path.isdir(os.path.join(path, "comfy")) + ) + + +def _load_module(module_name: str, filepath: str) -> Any: + """Load a Python module from an explicit file path, bypassing sys.path. + + 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 try: - return obj[index] - except KeyError: - return obj["result"][index] + spec = importlib.util.spec_from_file_location(module_name, filepath) + if spec is None or spec.loader is None: + log.debug("Could not create spec for %s at %s", module_name, filepath) + return None + mod = importlib.util.module_from_spec(spec) + sys.modules[module_name] = mod + try: + spec.loader.exec_module(mod) + except BaseException: + sys.modules.pop(module_name, None) + raise + return mod + except BaseException as e: + log.debug("Failed to load %s from %s: %s", module_name, filepath, e) + sys.modules.pop(module_name, None) # Also clean up on edge-case exceptions + return None -def get_comfyui_path() -> str | None: - """Resolve ComfyUI path via prioritized multi-strategy fallback. +def _load_module_temp(module_name: str, filepath: str) -> Any: + """Load a module via _load_module() then remove it from sys.modules. - Strategy order: - 1. COMFYUI_PATH env var (verified with _is_comfyui_directory) - 2. Relative walk from extension location (realpath + verified) - 3. CWD walk (legacy fallback, depth-limited) + Used during bootstrap for modules that ComfyUI's import chain also loads + normally — prevents the cached copy from conflicting with later imports. """ - p = os.environ.get("COMFYUI_PATH") - if p and _is_comfyui_directory(p): - return p - p = _find_from_extension_location() - if p: - return p - return find_path("ComfyUI", max_depth=20) - - -def find_path(name: str, max_depth: int = 20) -> str | None: - """Iteratively walk up from CWD to find a directory by name. - - Depth-limited to prevent slow startup on deep trees. - Each candidate verified by _is_comfyui_directory() at the caller. - """ - candidate = os.getcwd() - for _ in range(max_depth): - parent = os.path.dirname(candidate) - if parent == candidate: - break - candidate = parent - if os.path.basename(candidate) == name: - return candidate - return None + mod = _load_module(module_name, filepath) + sys.modules.pop(module_name, None) + return mod def add_comfyui_directory_to_sys_path() -> None: @@ -247,6 +224,51 @@ def cleanup_comfyui_runtime(unload_models: bool | None = None) -> None: gc.collect() +def find_path(name: str, max_depth: int = 20) -> str | None: + """Iteratively walk up from CWD to find a directory by name. + + Depth-limited to prevent slow startup on deep trees. + Each candidate verified by _is_comfyui_directory() at the caller. + """ + candidate = os.getcwd() + for _ in range(max_depth): + parent = os.path.dirname(candidate) + if parent == candidate: + break + candidate = parent + if os.path.basename(candidate) == name: + return candidate + return None + + +def get_comfyui_path() -> str | None: + """Resolve ComfyUI path via prioritized multi-strategy fallback. + + Strategy order: + 1. COMFYUI_PATH env var (verified with _is_comfyui_directory) + 2. Relative walk from extension location (realpath + verified) + 3. CWD walk (legacy fallback, depth-limited) + """ + p = os.environ.get("COMFYUI_PATH") + if p and _is_comfyui_directory(p): + return p + p = _find_from_extension_location() + if p: + return p + p = find_path("ComfyUI", max_depth=20) + if p and _is_comfyui_directory(p): + return p + return None + + +def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any: + """Return a sequence or mapping result item by index.""" + try: + return obj[index] + except KeyError: + return obj["result"][index] + + # Workflow data def build_workflow() -> dict[str, Any]: return { diff --git a/tests/runtime/generated/upscale-model-loader.py b/tests/runtime/generated/upscale-model-loader.py index 642d7aa..3b76a26 100644 --- a/tests/runtime/generated/upscale-model-loader.py +++ b/tests/runtime/generated/upscale-model-loader.py @@ -10,47 +10,6 @@ from typing import Sequence, Mapping, Any, Union log = logging.getLogger(__name__) -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 - from verified file paths instead of relying on sys.path resolution. - """ - if not os.path.isfile(filepath): - log.debug("Module file not found: %s (%s)", module_name, filepath) - return None - try: - spec = importlib.util.spec_from_file_location(module_name, filepath) - if spec is None or spec.loader is None: - log.debug("Could not create spec for %s at %s", module_name, filepath) - return None - mod = importlib.util.module_from_spec(spec) - sys.modules[module_name] = mod - spec.loader.exec_module(mod) - return mod - except Exception as e: - log.debug("Failed to load %s from %s: %s", module_name, filepath, e) - return None - - -def _load_module_temp(module_name: str, filepath: str) -> Any: - """Load a module via _load_module() then remove it from sys.modules. - - Used during bootstrap for modules that ComfyUI's import chain also loads - normally — prevents the cached copy from conflicting with later imports. - """ - mod = _load_module(module_name, filepath) - sys.modules.pop(module_name, None) - return mod - - -def _is_comfyui_directory(path: str) -> bool: - """Verify a directory has ComfyUI structural markers (nodes.py).""" - if not os.path.isdir(path): - return False - return os.path.isfile(os.path.join(path, "nodes.py")) - - def _find_from_extension_location() -> str | None: """Walk up from this file's location to find ComfyUI root.""" ext_dir = os.path.dirname(os.path.realpath(__file__)) @@ -66,46 +25,64 @@ def _find_from_extension_location() -> str | None: return None -def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any: - """Return a sequence or mapping result item by index.""" +def _is_comfyui_directory(path: str) -> bool: + """Verify a directory has ComfyUI structural markers. + + Checks for nodes.py, main.py, and the comfy/ subdirectory to raise + the bar against spoofing via a directory with only a single marker file. + """ + if not os.path.isdir(path): + return False + return ( + os.path.isfile(os.path.join(path, "nodes.py")) + and os.path.isfile(os.path.join(path, "main.py")) + and os.path.isdir(os.path.join(path, "comfy")) + ) + + +def _load_module(module_name: str, filepath: str) -> Any: + """Load a Python module from an explicit file path, bypassing sys.path. + + 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 try: - return obj[index] - except KeyError: - return obj["result"][index] + spec = importlib.util.spec_from_file_location(module_name, filepath) + if spec is None or spec.loader is None: + log.debug("Could not create spec for %s at %s", module_name, filepath) + return None + mod = importlib.util.module_from_spec(spec) + sys.modules[module_name] = mod + try: + spec.loader.exec_module(mod) + except BaseException: + sys.modules.pop(module_name, None) + raise + return mod + except BaseException as e: + log.debug("Failed to load %s from %s: %s", module_name, filepath, e) + sys.modules.pop(module_name, None) # Also clean up on edge-case exceptions + return None -def get_comfyui_path() -> str | None: - """Resolve ComfyUI path via prioritized multi-strategy fallback. +def _load_module_temp(module_name: str, filepath: str) -> Any: + """Load a module via _load_module() then remove it from sys.modules. - Strategy order: - 1. COMFYUI_PATH env var (verified with _is_comfyui_directory) - 2. Relative walk from extension location (realpath + verified) - 3. CWD walk (legacy fallback, depth-limited) + Used during bootstrap for modules that ComfyUI's import chain also loads + normally — prevents the cached copy from conflicting with later imports. """ - p = os.environ.get("COMFYUI_PATH") - if p and _is_comfyui_directory(p): - return p - p = _find_from_extension_location() - if p: - return p - return find_path("ComfyUI", max_depth=20) - - -def find_path(name: str, max_depth: int = 20) -> str | None: - """Iteratively walk up from CWD to find a directory by name. - - Depth-limited to prevent slow startup on deep trees. - Each candidate verified by _is_comfyui_directory() at the caller. - """ - candidate = os.getcwd() - for _ in range(max_depth): - parent = os.path.dirname(candidate) - if parent == candidate: - break - candidate = parent - if os.path.basename(candidate) == name: - return candidate - return None + mod = _load_module(module_name, filepath) + sys.modules.pop(module_name, None) + return mod def add_comfyui_directory_to_sys_path() -> None: @@ -247,6 +224,54 @@ def cleanup_comfyui_runtime(unload_models: bool | None = None) -> None: gc.collect() +def find_path(name: str, max_depth: int = 20) -> str | None: + """Iteratively walk up from CWD to find a directory by name. + + Depth-limited to prevent slow startup on deep trees. + Each candidate verified by _is_comfyui_directory() at the caller. + """ + candidate = os.getcwd() + for _ in range(max_depth): + parent = os.path.dirname(candidate) + if parent == candidate: + break + candidate = parent + if os.path.basename(candidate) == name: + return candidate + return None + + +def get_comfyui_path() -> str | None: + """Resolve ComfyUI path via prioritized multi-strategy fallback. + + Strategy order: + 1. COMFYUI_PATH env var (verified with _is_comfyui_directory) + 2. Relative walk from extension location (realpath + verified) + 3. CWD walk (legacy fallback, depth-limited) + """ + p = os.environ.get("COMFYUI_PATH") + if p and _is_comfyui_directory(p): + return p + p = _find_from_extension_location() + if p: + return p + p = find_path("ComfyUI", max_depth=20) + if p and _is_comfyui_directory(p): + return p + return None + + +def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any: + """Return a sequence or mapping result item by index.""" + try: + return obj[index] + except KeyError: + return obj["result"][index] + + +import_custom_nodes() + + def import_custom_nodes() -> None: """Initialize ComfyUI custom nodes in the exporter runtime. @@ -277,10 +302,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( diff --git a/tests/test_app_base_mappings.py b/tests/test_app_base_mappings.py new file mode 100644 index 0000000..b5d8633 --- /dev/null +++ b/tests/test_app_base_mappings.py @@ -0,0 +1,165 @@ +"""Tests for ExportApplication base_node_class_mappings behavior. + +Verifies that after custom node import, base_node_class_mappings retains the +pre-custom-node baseline so WorkflowPlanner correctly distinguishes between +built-in nodes (direct import) and custom/extras nodes (NODE_CLASS_MAPPINGS lookup). +""" + +import sys +import unittest +from io import StringIO + + +def _make_workflow(): + """Minimal workflow JSON string for testing.""" + return '{"1": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": "test.safetensors"}}}' + + +class TestBaseNodeClassMappings(unittest.TestCase): + """Tests for base_node_class_mappings stability across custom node reloads.""" + + def tearDown(self): + # Clean up any sys.modules pollution to avoid cross-test contamination + if "nodes" in sys.modules: + del sys.modules["nodes"] + + def test_base_mappings_unchanged_after_custom_node_import(self): + """Given custom nodes are loaded, base_node_class_mappings stays at original baseline.""" + from comfyui_to_python.app import ExportApplication + + class FakeNode: + CATEGORY = "loaders" + FUNCTION = "execute" + + @staticmethod + def INPUT_TYPES(): + return {"required": {"ckpt_name": ("STRING",)}} + + def execute(self, ckpt_name: str): + return ({},) + + class FakeCustomNode: + CATEGORY = "custom" + FUNCTION = "execute" + + @staticmethod + def INPUT_TYPES(): + return {"required": {"ckpt_name": ("STRING",)}} + + def execute(self, ckpt_name: str): + return ({},) + + initial_mappings = {"CheckpointLoaderSimple": FakeNode} + after_import_mappings = dict(initial_mappings) + # Simulate custom nodes adding a new class + after_import_mappings["CustomNode"] = FakeCustomNode + + call_count = [0] + + def mock_node_loader(): + return initial_mappings + + def mock_custom_node_importer(): + call_count[0] += 1 + # Simulate populating sys.modules["nodes"].NODE_CLASS_MAPPINGS + nodes_mod_type = type("NodesModule", (), {})() + nodes_mod_type.NODE_CLASS_MAPPINGS = after_import_mappings + sys.modules["nodes"] = nodes_mod_type + + output = StringIO() + app = ExportApplication( + workflow=_make_workflow(), + output_file=output, + node_mapping_loader=mock_node_loader, + custom_node_importer=mock_custom_node_importer, + needs_init_custom_nodes=True, + ) + + # base_node_class_mappings should be a deep copy of initial (pre-custom-node) state + self.assertEqual(app.base_node_class_mappings, initial_mappings) + app.execute() + + # After execute(), base_node_class_mappings must still be the original baseline + # (not overwritten with after_import_mappings which includes custom nodes) + self.assertEqual( + app.base_node_class_mappings, + initial_mappings, + "base_node_class_mappings should remain at pre-custom-node baseline", + ) + + def test_base_mappings_is_deep_copy_not_shared_reference(self): + """Given node mappings change in-place, base_node_class_mappings is independent.""" + from comfyui_to_python.app import ExportApplication + + class FakeNode2: + CATEGORY = "loaders" + FUNCTION = "execute" + + @staticmethod + def INPUT_TYPES(): + return {"required": {"ckpt_name": ("STRING",)}} + + def execute(self, ckpt_name: str): + return ({},) + + initial_mappings = {"CheckpointLoaderSimple": FakeNode2} + + def mock_node_loader(): + return dict(initial_mappings) + + output = StringIO() + app = ExportApplication( + workflow=_make_workflow(), + output_file=output, + node_mapping_loader=mock_node_loader, + custom_node_importer=lambda: None, + ) + + # Mutate the live mappings + app.node_class_mappings["NewKey"] = "value" + + # base_node_class_mappings must not reflect the mutation (deep copy) + self.assertNotIn( + "NewKey", + app.base_node_class_mappings, + "base_node_class_mappings must be a deep copy, not shared reference", + ) + + +class TestRenderHardenedImports(unittest.TestCase): + """Integration test: render output contains hardened _load_module pattern.""" + + def test_generated_output_contains_exec_failure_cleanup(self): + """Given a GenerationPlan, rendered output includes try/except BaseException around exec_module().""" + from comfyui_to_python.generator.render import WorkflowRenderer + from comfyui_to_python.generator.model import GenerationPlan + + plan = GenerationPlan( + workflow_data={"1": {"class_type": "CheckpointLoaderSimple"}}, + metadata_workflow_data=None, + custom_nodes=False, + import_statements={}, + special_functions_code=[], + loop_code=["result = some_node()"], + queue_size=1, + ) + renderer = WorkflowRenderer() + generated = renderer.render(plan) + + # The _load_module in rendered output must have the exec failure cleanup pattern + self.assertIn("try:", generated) + self.assertIn("spec.loader.exec_module(mod)", generated) + self.assertIn( + "except BaseException:", + generated, + "Must catch BaseException for edge-case safety", + ) + self.assertIn( + "sys.modules.pop(module_name, None)", + generated, + "Must clean up sys.modules on exec failure", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_import_path_resolution.py b/tests/test_import_path_resolution.py index ff162bc..64d77f4 100644 --- a/tests/test_import_path_resolution.py +++ b/tests/test_import_path_resolution.py @@ -441,20 +441,40 @@ 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.""" + """Given a crash during import_custom_nodes(), sys.path is still restored. + + The crash fires on the "server" module load (after comfy/ filtering) + so that the try/finally sys.path restoration path is actually exercised. + """ from comfyui_to_python.node_runtime import import_custom_nodes original_path = sys.path[:] + call_order = [] + + def fake_load(name, path): + call_order.append(name) + if name == "server": + raise RuntimeError("simulated crash during server load") + return None # Let execution/nodes return None (falls through gracefully) with patch( "comfyui_to_python.node_runtime._load_module", - side_effect=RuntimeError("simulated crash during server load"), + side_effect=fake_load, + ), patch( + "comfyui_to_python.node_runtime.get_comfyui_path", + return_value=_REAL_COMFYUI if os.path.isdir(_REAL_COMFYUI) else None, ): try: import_custom_nodes() except RuntimeError: pass # Expected crash + # When ComfyUI path is found, "execution" and "nodes" should be called + # before reaching the server module load inside try/finally. + self.assertIn("execution", call_order, "Expected 'execution' to be loaded") + self.assertIn("nodes", call_order, "Expected 'nodes' to be loaded") + self.assertIn("server", call_order, "Expected 'server' to trigger the crash") + # sys.path should still contain all original entries (no dangling state) for entry in original_path: self.assertIn(entry, sys.path)