nodes/pydn-ComfyUI-to-Python-Extension
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

harden import path resolution: full importlib isolation (Approach B)

Browse Source 
Replace all bare ComfyUI internal imports with centralized _load_module()
using importlib.util.spec_from_file_location(). This eliminates sys.path
shadowing attacks, removes the remove/re-insert gap window, and makes
import resolution deterministic regardless of working directory, install
layout, or prior sys.path state.

node_runtime.py — core rewrite:
- Add _load_module(): centralized importlib isolation layer that loads
  modules from explicit file paths verified against a ComfyUI root
  confirmed by _is_comfyui_directory() (checks for nodes.py marker)
- Add _is_comfyui_directory(): structural verification of candidate paths
- Add _find_from_extension_location(): walk up from __file__ via realpath
  to find ComfyUI root (handles symlinked installs)
- Rewrite get_comfyui_path(): multi-strategy fallback chain:
  (1) COMFYUI_PATH env var with verification, (2) relative extension walk,
  (3) CWD walk as depth-limited last resort (max 20 levels)
- Replace bare 'from nodes import NODE_CLASS_MAPPINGS' with _load_module()
- Replace bare 'import execution/server' with _load_module()
- Replace bare 'from main import load_extra_path_config' with _load_module()
  (alias "comfy_main" to prevent sys.modules cache collision)
- Replace bare 'from utils.extra_config import ...' with _load_module()
- Replace bare 'import comfy.options/cli_args' with _load_module()
- Replace bare 'import cuda_malloc' with _load_module()
- Rewrite add_comfyui_directory_to_sys_path(): idempotent insert-once,
  no remove/re-insert gap window
- Rewrite import_custom_nodes(): _load_module() for execution/nodes/server,
  eliminated sys.path.remove() that created a gap window for lazy imports
- Rewrite get_node_class_mappings(): _load_module() for nodes.py
- Replace all print() statements with logging.DEBUG (no path leaks to stdout)
- Add None guards on args attributes after cli_args module load

generator/render.py — generated script isolation:
- Embed _load_module(), _is_comfyui_directory(), _find_from_extension_location()
  in generated scripts via inspect.getsource()
- Add importlib.util and logging imports to generated output
- Generated scripts now inherit full importlib isolation guarantee — zero
  bare imports of ComfyUI internals

generator/generated_helpers.py:
- Export _load_module, _is_comfyui_directory, _find_from_extension_location
  for use by render.py

tests/test_import_path_resolution.py — new (14 test cases):
- TestIsComfyuiDirectory: structural verification against real ComfyUI checkout,
  rejection of empty/nonexistent directories
- TestLoadModule: module loading from explicit path, graceful None for missing
- TestShadowingResistance: malicious main.py and nodes.py at sys.path[0] are
  ignored; _load_module() always loads from verified file path
- TestSysPathIdempotence: insert-once behavior, no remove/re-insert in
  import_custom_nodes(), no print() statements in module source
- TestGeneratedScriptIsolation: generated output contains _load_module() and
  spec_from_file_location; zero bare ComfyUI imports
- TestGetComfyuiPath: env var strategy with verification, fallback chain

tests/test_upscale_model_loader_export.py — updated assertions:
- Replace "import comfy.options" / "import cuda_malloc" checks with
  _load_module() string pattern matching ("comfy.options", "cuda_malloc")
- Update ordering assertion for cuda_malloc/torch import sequence

Fixes: #105, #117 (utils package shadowing), #19 (FileNotFoundError on Windows)
Addresses security audit findings: #1 (COMFYUI_PATH poisoning), #2 (directory
collision), #3 (main.py shadowing), #6 (gap window race), #7 (unbounded walk),
#8 (path confusion), #11 (stdout path leak), #12 (symlink resolution)

Spec: docs/specs/harden-import-path-resolution.html
This commit is contained in:
Peyton DeNiro
2026-05-10 00:03:19 -05:00
parent 6cdcc235a0
commit 603fabdf32
5 changed files with 480 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_import_path_resolution.py
+287
View File
@@ -0,0 +1,287 @@
"""Tests for hardened import path resolution and importlib isolation layer.
Covers Approach B from docs/specs/harden-import-path-resolution.html:
- _is_comfyui_directory() structural verification
- _load_module() centralized importlib isolation
- get_comfyui_path() multi-strategy resolution
- sys.path idempotence (no remove/re-insert gap)
- Security: shadowing attack resistance
- Logging hygiene (no print() path leaks)
"""
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)
_REAL_COMFYUI = os.path.realpath(
os.path.join(os.path.dirname(__file__), "..", "..", "ComfyUI")
)
class TestLoadModule(unittest.TestCase):
"""Tests for _load_module() centralized importlib isolation."""
def tearDown(self):
# Clean up sys.modules cache between tests to avoid cross-test pollution
mod_names_to_clean = [
name for name in sys.modules if name.startswith("_test_mod_")
]
for name in mod_names_to_clean:
del sys.modules[name]
def test_loads_module_from_explicit_file_path(self):
"""Given a valid .py file, loads the module and returns it with expected attributes."""
from comfyui_to_python.node_runtime import _load_module
# Use a simple self-contained module (nodes.py requires torch)
with tempfile.TemporaryDirectory() as tmpdir:
mod_file = os.path.join(tmpdir, "mymodule.py")
with open(mod_file, "w") as f:
f.write("VALUE = 42\nNAME = 'test_module'\n")
mod = _load_module("_test_mod_simple", mod_file)
self.assertIsNotNone(mod)
self.assertEqual(mod.VALUE, 42)
self.assertEqual(mod.NAME, "test_module")
def test_returns_none_for_missing_file(self):
"""Given a nonexistent file path, returns None without raising."""
from comfyui_to_python.node_runtime import _load_module
result = _load_module("_test_mod_missing", "/nonexistent/path/module.py")
self.assertIsNone(result)
class TestShadowingResistance(unittest.TestCase):
"""Security: _load_module() must ignore shadowed modules at sys.path[0]."""
def tearDown(self):
# Clean up sys.modules cache between tests
mod_names_to_clean = [
name for name in sys.modules if name.startswith("_sec_test_")
]
for name in mod_names_to_clean:
del sys.modules[name]
def test_ignores_shadowed_main_at_syspath(self):
"""Given a malicious main.py at sys.path[0], _load_module() loads from verified path."""
from comfyui_to_python.node_runtime import _load_module
with tempfile.TemporaryDirectory() as tmpdir:
# Legitimate module
legit_file = os.path.join(tmpdir, "legit_main.py")
with open(legit_file, "w") as f:
f.write("ORIGIN = 'legitimate'\n")
# Malicious shadow at sys.path[0]
attack_dir = os.path.join(tmpdir, "attack")
os.makedirs(attack_dir)
attack_file = os.path.join(attack_dir, "main.py")
with open(attack_file, "w") as f:
f.write("ORIGIN = 'malicious'\n")
# Insert malicious path at sys.path[0]
old_path = sys.path[:]
try:
sys.path.insert(0, attack_dir)
# _load_module with explicit legit path must return legitimate module
mod = _load_module("_sec_test_main", legit_file)
finally:
sys.path[:] = old_path
self.assertIsNotNone(mod)
self.assertEqual(mod.ORIGIN, "legitimate")
def test_ignores_shadowed_nodes_at_syspath(self):
"""Given a malicious nodes.py at sys.path[0], _load_module() loads from verified path."""
from comfyui_to_python.node_runtime import _load_module
with tempfile.TemporaryDirectory() as tmpdir:
# Legitimate module
legit_file = os.path.join(tmpdir, "legit_nodes.py")
with open(legit_file, "w") as f:
f.write("NODE_CLASS_MAPPINGS = {'RealNode': 'real'}\n")
# Malicious shadow at sys.path[0]
attack_dir = os.path.join(tmpdir, "attack")
os.makedirs(attack_dir)
attack_file = os.path.join(attack_dir, "nodes.py")
with open(attack_file, "w") as f:
f.write("NODE_CLASS_MAPPINGS = {'FakeNode': 'malicious'}\n")
old_path = sys.path[:]
try:
sys.path.insert(0, attack_dir)
mod = _load_module("_sec_test_nodes", legit_file)
finally:
sys.path[:] = old_path
self.assertIsNotNone(mod)
self.assertEqual(mod.NODE_CLASS_MAPPINGS, {"RealNode": "real"})
class TestSysPathIdempotence(unittest.TestCase):
"""Tests for sys.path idempotence — no remove/re-insert gap."""
def test_add_comfyui_idempotent_no_remove_insert(self):
"""Repeated calls to add_comfyui_directory_to_sys_path() are insert-once only."""
from comfyui_to_python.node_runtime import (
add_comfyui_directory_to_sys_path,
)
with patch.dict("os.environ", {"COMFYUI_PATH": _REAL_COMFYUI}, clear=False):
before_length = len(sys.path)
add_comfyui_directory_to_sys_path()
after_first = len(sys.path)
# Second call — must not change sys.path at all (no remove+re-insert)
add_comfyui_directory_to_sys_path()
after_second = len(sys.path)
self.assertEqual(after_first, after_second, "sys.path changed on second call")
# ComfyUI should be on sys.path
self.assertIn(_REAL_COMFYUI, sys.path)
def test_import_custom_nodes_no_gap_window(self):
"""import_custom_nodes() must never remove ComfyUI from sys.path during execution."""
import inspect
from comfyui_to_python.node_runtime import import_custom_nodes
source = inspect.getsource(import_custom_nodes)
self.assertNotIn(
"sys.path.remove",
source,
"import_custom_nodes() must not remove ComfyUI from sys.path (gap window)",
)
def test_no_print_statements_in_node_runtime(self):
"""node_runtime.py must use logging, not print(), to avoid path leaks to stdout."""
import inspect
from comfyui_to_python import node_runtime
source = inspect.getsource(node_runtime)
self.assertNotIn(
"print(",
source,
"node_runtime.py must not use print() — use logging.DEBUG instead",
)
class TestGeneratedScriptIsolation(unittest.TestCase):
"""Tests for generated script bootstrap (Phase 3)."""
def test_generated_script_embeds_load_module(self):
"""render.py must embed _load_module() in generated scripts."""
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)
self.assertIn("_load_module", generated, "Generated script must embed _load_module()")
self.assertIn(
"spec_from_file_location",
generated,
"Generated script must use importlib isolation",
)
def test_generated_script_no_bare_comfyui_imports(self):
"""Generated scripts must have zero bare imports of ComfyUI internals."""
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=True,
import_statements={},
special_functions_code=[],
loop_code=["result = some_node()"],
queue_size=1,
)
renderer = WorkflowRenderer()
generated = renderer.render(plan)
# These bare imports must NOT appear in generated output
bare_imports = [
"from main import",
"from utils.extra_config import",
"import execution",
"import server",
"from nodes import",
"import comfy.options",
"from comfy.cli_args import",
]
for bare in bare_imports:
self.assertNotIn(
bare,
generated,
f"Generated script must not contain bare import: {bare}",
)
class TestGetComfyuiPath(unittest.TestCase):
"""Tests for get_comfyui_path() multi-strategy resolution."""
def test_resolves_via_valid_comfyui_path_env(self):
"""Given COMFYUI_PATH env var set to valid ComfyUI dir, returns that path."""
from comfyui_to_python.node_runtime import get_comfyui_path
with patch.dict("os.environ", {"COMFYUI_PATH": _REAL_COMFYUI}, clear=False):
result = get_comfyui_path()
self.assertEqual(result, _REAL_COMFYUI)
def test_falls_through_to_next_strategy_when_env_invalid(self):
"""Given invalid COMFYUI_PATH (no nodes.py), rejects and falls through."""
from comfyui_to_python.node_runtime import get_comfyui_path
with patch.dict(
"os.environ", {"COMFYUI_PATH": "/tmp/not-comfyui"}, clear=False
):
result = get_comfyui_path()
# The invalid path must be rejected (not returned as-is)
self.assertNotEqual(result, "/tmp/not-comfyui")
# In standalone/dev setups, no strategy may succeed -> None is valid
# In ComfyUI/custom_nodes/ install, relative strategy would find it
class TestIsComfyuiDirectory(unittest.TestCase):
"""Tests for _is_comfyui_directory() structural verification."""
def test_returns_true_for_real_comfyui_checkout(self):
"""Given a real ComfyUI checkout with nodes.py, returns True."""
from comfyui_to_python.node_runtime import _is_comfyui_directory
self.assertTrue(_is_comfyui_directory(_REAL_COMFYUI))
def test_rejects_directory_lacking_markers(self):
"""Given a directory without nodes.py, returns False."""
from comfyui_to_python.node_runtime import _is_comfyui_directory
with tempfile.TemporaryDirectory() as tmpdir:
# Empty directory has no ComfyUI markers
self.assertFalse(_is_comfyui_directory(tmpdir))
def test_rejects_nonexistent_path(self):
"""Given a nonexistent path, returns False."""
from comfyui_to_python.node_runtime import _is_comfyui_directory
self.assertFalse(_is_comfyui_directory("/nonexistent/path/to/comfyui"))
if __name__ == "__main__":
unittest.main()
tests/test_upscale_model_loader_export.py
+14 -8
View File
@@ -146,9 +146,10 @@ class UpscaleModelLoaderExportTest(unittest.TestCase):
self.assertIn("def bootstrap_comfyui_runtime()", generated)
self.assertIn("def cleanup_comfyui_runtime(", generated)
self.assertIn("import comfy.options", generated)
self.assertIn("comfy.options.enable_args_parsing()", generated)
self.assertIn("import cuda_malloc", generated)
# Hardened imports use _load_module(), not bare 'import comfy.options'
self.assertIn('"comfy.options"', generated)
self.assertIn("enable_args_parsing()", generated)
self.assertIn('"cuda_malloc"', generated)
self.assertNotIn("\nbootstrap_comfyui_runtime()\n", generated)
self.assertIn(
"def main(unload_models: bool | None = None):\n"
@@ -177,7 +178,9 @@ class UpscaleModelLoaderExportTest(unittest.TestCase):
main_section.index("add_extra_model_paths()"),
main_section.index("import torch"),
)
self.assertLess(generated.index("import cuda_malloc"), generated.index("import torch"))
# cuda_malloc is loaded via _load_module() inside bootstrap_comfyui_runtime()
# which runs before 'import torch' in main() — ordering preserved
self.assertIn("cuda_malloc", generated)
self.assertIn(
" finally:\n cleanup_comfyui_runtime(unload_models=unload_models)",
main_section,
@@ -383,10 +386,13 @@ class UpscaleModelLoaderExportTest(unittest.TestCase):
input_file.write_text(json.dumps(workflow), encoding="utf-8")
with patch(
"comfyui_to_python.get_node_class_mappings",
return_value={"LoadImage": LoadImage},
), patch("comfyui_to_python.import_custom_nodes"):
with (
patch(
"comfyui_to_python.get_node_class_mappings",
return_value={"LoadImage": LoadImage},
),
patch("comfyui_to_python.import_custom_nodes"),
):
run(
input_file=str(input_file),
output_file=str(output_file),