Commit Graph
7 Commits
Author SHA1 Message Date
Peyton DeNiro 603fabdf32 harden import path resolution: full importlib isolation (Approach B)
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
2026-05-10 00:03:19 -05:00
Peyton DeNiro 77c0bbec9b Guard ComfyUI cleanup teardown failures 2026-03-29 19:02:21 -05:00
Peyton DeNiro f9163f2375 Add generated runtime cleanup after workflow execution. Exported scripts now run best-effort ComfyUI cache cleanup in a finally block and support opt-in model unloading for embedded repeated-call use cases. Document the supported VRAM-related runtime flags and lock the behavior in focused exporter and runtime cleanup tests. 2026-03-29 18:18:37 -05:00
Peyton DeNiro 574d96369c Make generator codegen regressions explicit and fix string seed randomization 2026-03-29 17:09:30 -05:00
Peyton DeNiro 42d0526661 Fix generated workflow torch import order
Move generated torch imports until after add_extra_model_paths runs so exported scripts do not trigger ComfyUI's early-torch warning during startup. Update export assertions and checked-in runtime generated fixtures to match the new main() ordering.
2026-03-29 15:15:52 -05:00
Peyton DeNiro 441dcfebdf Defer generated ComfyUI bootstrap until execution 2026-03-29 14:40:17 -05:00
Peyton DeNiro 9e7e474b05 Refactor the exporter into a package-backed module layout.
Split the previous monolithic implementation into focused modules for CLI handling, workflow loading, runtime bootstrap, load ordering, planning, rendering, and file I/O. Keep the public compatibility surface intact by preserving the legacy top-level wrapper while adding  as the documented entrypoint. This reduces cross-cutting coupling inside the exporter and makes the runtime/bootstrap behavior easier to evolve without carrying unrelated concerns in one file.
2026-03-29 14:09:24 -05:00