ausboss
0f689e3bea
Widget values reach the loaders through ComfyUI's /prompt route, and resolve_input_path and local_preview_allowed are what keep those reads inside ComfyUI's input, output and temp folders. No test covered them: removing the outside-the-input-folder check, the network-path check or local path mode's network check passed every test. tests/test_media_paths.py covers climbing out with "..", absolute paths, links pointing outside, "[output]" names, network paths refused before the path is touched (even when a ComfyUI folder itself sits on the share), and local path mode. The guards need neither torch nor PyAV, so the media module's heavy imports are stubbed when missing and the file runs in CI. Each guard's removal fails at least two tests.
241 lines
11 KiB
Python
241 lines
11 KiB
Python
"""File reads stay inside ComfyUI's folders.
|
|
|
|
Widget values reach the pack through ComfyUI's /prompt route, which needs no
|
|
login, and the loaders' VALIDATE_INPUTS hand their source straight to these
|
|
functions. So they are the gate between a crafted prompt and the disk:
|
|
|
|
- resolve_input_path takes only files inside ComfyUI's input folder;
|
|
- a network path ("//host/share" or "\\\\host\\share") is refused from its
|
|
text alone, before anything touches the path, since on Windows merely
|
|
resolving one contacts that host;
|
|
- local path mode (local_preview_allowed) reads only inside the input,
|
|
output and temp folders, and never a network path either.
|
|
|
|
The guards need neither torch nor PyAV, so this file runs in CI: the media
|
|
module's heavy imports are stubbed when they are missing.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import importlib.util
|
|
import os
|
|
from pathlib import Path
|
|
import sys
|
|
import tempfile
|
|
import types
|
|
import unittest
|
|
from unittest.mock import patch
|
|
|
|
ROOT = Path(__file__).resolve().parent.parent
|
|
sys.path.insert(0, str(ROOT))
|
|
if "nodes" in sys.modules and not hasattr(sys.modules["nodes"], "__path__"):
|
|
del sys.modules["nodes"]
|
|
|
|
HEAVY = ("av", "PIL", "numpy", "torch")
|
|
REAL_MEDIA_STACK = all(importlib.util.find_spec(name) for name in HEAVY)
|
|
|
|
# Only the import line needs these; the guards never call into them.
|
|
if importlib.util.find_spec("av") is None:
|
|
sys.modules["av"] = types.ModuleType("av")
|
|
if importlib.util.find_spec("PIL") is None:
|
|
pil = types.ModuleType("PIL")
|
|
for part in ("Image", "ImageOps", "ImageSequence"):
|
|
setattr(pil, part, types.ModuleType(f"PIL.{part}"))
|
|
sys.modules["PIL"] = pil
|
|
|
|
from nodes import _media_helpers # noqa: E402
|
|
from nodes._media_helpers import local_preview_allowed, resolve_input_path, resolve_video_path # noqa: E402
|
|
|
|
NETWORK_REFUSAL = "Network source paths are not allowed."
|
|
OUTSIDE_REFUSAL = "The selected source is outside ComfyUI's input folder."
|
|
|
|
|
|
class Folders:
|
|
"""A stand-in for ComfyUI's folder_paths over three chosen folders.
|
|
|
|
A name may end in " [input]", " [output]" or " [temp]" to pick the
|
|
folder, as the frontend writes for MaskEditor saves; otherwise it is
|
|
joined onto the input folder the way ComfyUI does, so an absolute name
|
|
stays absolute.
|
|
"""
|
|
|
|
def __init__(self, input_dir, output_dir, temp_dir):
|
|
self.dirs = {"input": str(input_dir), "output": str(output_dir), "temp": str(temp_dir)}
|
|
self.looked_up = []
|
|
|
|
def get_input_directory(self):
|
|
return self.dirs["input"]
|
|
|
|
def get_output_directory(self):
|
|
return self.dirs["output"]
|
|
|
|
def get_temp_directory(self):
|
|
return self.dirs["temp"]
|
|
|
|
def get_annotated_filepath(self, name):
|
|
self.looked_up.append(name)
|
|
base = self.dirs["input"]
|
|
for kind in ("input", "output", "temp"):
|
|
tag = f" [{kind}]"
|
|
if name.endswith(tag):
|
|
name, base = name[: -len(tag)], self.dirs[kind]
|
|
return os.path.join(base, name)
|
|
|
|
|
|
class TempFolders(unittest.TestCase):
|
|
def setUp(self):
|
|
self.tmp = tempfile.TemporaryDirectory()
|
|
self.addCleanup(self.tmp.cleanup)
|
|
self.root = Path(self.tmp.name).resolve()
|
|
self.input = self.root / "input"
|
|
self.output = self.root / "output"
|
|
self.temp = self.root / "temp"
|
|
for folder in (self.input, self.output, self.temp, self.input / "sub", self.input / "clipspace"):
|
|
folder.mkdir(parents=True, exist_ok=True)
|
|
self.inside = self.touch(self.input / "pic.png")
|
|
self.nested = self.touch(self.input / "sub" / "pic.png")
|
|
self.masked = self.touch(self.input / "clipspace" / "painted.png")
|
|
self.outside = self.touch(self.root / "secret.png")
|
|
self.in_output = self.touch(self.output / "render.mp4")
|
|
self.folders = Folders(self.input, self.output, self.temp)
|
|
patcher = patch.object(_media_helpers, "folder_paths", self.folders)
|
|
patcher.start()
|
|
self.addCleanup(patcher.stop)
|
|
|
|
@staticmethod
|
|
def touch(path):
|
|
path.write_bytes(b"x")
|
|
return path
|
|
|
|
|
|
class InputFolderTests(TempFolders):
|
|
def test_files_inside_the_input_folder_resolve(self):
|
|
self.assertEqual(resolve_input_path("pic.png"), self.inside)
|
|
self.assertEqual(resolve_input_path("sub/pic.png"), self.nested)
|
|
# A MaskEditor save names its folder; it is still the input folder.
|
|
self.assertEqual(resolve_input_path("clipspace/painted.png [input]"), self.masked)
|
|
|
|
def test_a_path_that_climbs_out_is_refused(self):
|
|
with self.assertRaisesRegex(ValueError, OUTSIDE_REFUSAL):
|
|
resolve_input_path("../secret.png")
|
|
with self.assertRaisesRegex(ValueError, OUTSIDE_REFUSAL):
|
|
resolve_input_path("sub/../../secret.png")
|
|
|
|
def test_an_absolute_path_elsewhere_is_refused(self):
|
|
# os.path.join keeps an absolute name as it is, so this is the path
|
|
# a crafted prompt would reach for.
|
|
with self.assertRaisesRegex(ValueError, OUTSIDE_REFUSAL):
|
|
resolve_input_path(str(self.outside))
|
|
|
|
def test_a_link_inside_that_points_outside_is_refused(self):
|
|
(self.input / "shortcut.png").symlink_to(self.outside)
|
|
with self.assertRaisesRegex(ValueError, OUTSIDE_REFUSAL):
|
|
resolve_input_path("shortcut.png")
|
|
|
|
def test_another_comfy_folder_is_still_outside(self):
|
|
# Core Load Image reads "[output]" names; these loaders read input only.
|
|
with self.assertRaisesRegex(ValueError, OUTSIDE_REFUSAL):
|
|
resolve_input_path("render.mp4 [output]")
|
|
|
|
def test_blank_and_missing_names_say_what_to_do(self):
|
|
with self.assertRaisesRegex(ValueError, "Select or upload a source file first."):
|
|
resolve_input_path("")
|
|
with self.assertRaisesRegex(ValueError, "no longer exists"):
|
|
resolve_input_path("gone.png")
|
|
|
|
|
|
class NetworkPathTests(TempFolders):
|
|
NAMES = ("//fileserver/share/pic.png", "\\\\fileserver\\share\\pic.png", "//fileserver/share/pic.png [input]")
|
|
|
|
def test_a_network_source_is_refused_before_the_path_is_touched(self):
|
|
resolved = []
|
|
real_resolve = Path.resolve
|
|
|
|
def recording_resolve(path, *args, **kwargs):
|
|
resolved.append(str(path))
|
|
return real_resolve(path, *args, **kwargs)
|
|
|
|
for name in self.NAMES:
|
|
with self.subTest(name=name), patch.object(Path, "resolve", recording_resolve):
|
|
with self.assertRaisesRegex(ValueError, NETWORK_REFUSAL):
|
|
resolve_input_path(name)
|
|
self.assertEqual(self.folders.looked_up, [], "the name was looked up before it was refused")
|
|
self.assertEqual(resolved, [], "a network path was resolved before it was refused")
|
|
|
|
def test_the_refusal_holds_when_the_input_folder_itself_is_on_a_share(self):
|
|
# Inside the input folder by text, still a network read.
|
|
self.folders.dirs["input"] = "//fileserver/comfy/input"
|
|
with self.assertRaisesRegex(ValueError, NETWORK_REFUSAL):
|
|
resolve_input_path("//fileserver/comfy/input/pic.png")
|
|
|
|
|
|
class LocalPathTests(TempFolders):
|
|
def test_comfy_folders_are_readable_in_place(self):
|
|
self.assertTrue(local_preview_allowed(str(self.in_output)))
|
|
self.assertTrue(local_preview_allowed(f'"{self.inside}"'))
|
|
self.assertEqual(resolve_video_path("local path", "", str(self.in_output)), self.in_output)
|
|
|
|
def test_anywhere_else_is_refused(self):
|
|
self.assertFalse(local_preview_allowed(str(self.outside)))
|
|
self.assertFalse(local_preview_allowed(str(self.input / ".." / "secret.png")))
|
|
self.assertFalse(local_preview_allowed(""))
|
|
movie = self.touch(self.root / "movie.mp4")
|
|
with self.assertRaisesRegex(ValueError, "Local path mode reads only videos inside"):
|
|
resolve_video_path("local path", "", str(movie))
|
|
|
|
def test_a_network_path_is_refused_from_its_text_alone(self):
|
|
resolved = []
|
|
real_resolve = Path.resolve
|
|
|
|
def recording_resolve(path, *args, **kwargs):
|
|
resolved.append(str(path))
|
|
return real_resolve(path, *args, **kwargs)
|
|
|
|
with patch.object(Path, "resolve", recording_resolve):
|
|
self.assertFalse(local_preview_allowed("//fileserver/share/clip.mp4"))
|
|
self.assertFalse(local_preview_allowed("\\\\fileserver\\share\\clip.mp4"))
|
|
self.assertFalse([path for path in resolved if "fileserver" in path], resolved)
|
|
|
|
def test_a_network_path_is_refused_even_inside_a_comfy_folder_on_a_share(self):
|
|
self.folders.dirs["output"] = "//fileserver/comfy/output"
|
|
self.assertFalse(local_preview_allowed("//fileserver/comfy/output/clip.mp4"))
|
|
with self.assertRaisesRegex(ValueError, "Local path mode reads only videos inside"):
|
|
resolve_video_path("local path", "", "//fileserver/comfy/output/clip.mp4")
|
|
|
|
|
|
@unittest.skipUnless(REAL_MEDIA_STACK, "needs torch, NumPy, Pillow and PyAV (ComfyUI's Python)")
|
|
class LoaderValidationTests(TempFolders):
|
|
"""The loaders refuse the same paths before a run, with one message each."""
|
|
|
|
def test_image_loaders_refuse_paths_outside_the_input_folder(self):
|
|
from nodes import node_image_crop_rotate_pad, node_load_image_pad
|
|
|
|
for module, cls, label in (
|
|
(node_image_crop_rotate_pad, node_image_crop_rotate_pad.AusBossImageCropRotatePad, "Image Crop + Rotate + Pad"),
|
|
(node_load_image_pad, node_load_image_pad.AusBossLoadImagePad, "Load Image + Pad"),
|
|
):
|
|
with self.subTest(node=label):
|
|
self.assertIs(cls.VALIDATE_INPUTS(image="pic.png"), True)
|
|
for name in ("../secret.png", str(self.outside), "//fileserver/share/pic.png"):
|
|
verdict = cls.VALIDATE_INPUTS(image=name)
|
|
self.assertIsInstance(verdict, str, name)
|
|
self.assertTrue(verdict.startswith(f"{label}: "), verdict)
|
|
|
|
def test_video_loaders_refuse_paths_outside_comfy_folders(self):
|
|
from nodes import node_load_video, node_video_crop_rotate_pad
|
|
|
|
self.touch(self.input / "clip.mp4")
|
|
load = node_load_video.AusBossLoadVideo
|
|
self.assertIs(load.VALIDATE_INPUTS("clip.mp4"), True)
|
|
self.assertIn("outside", load.VALIDATE_INPUTS("../secret.png"))
|
|
frame = node_video_crop_rotate_pad.AusBossVideoCropRotatePad
|
|
self.assertIs(frame.VALIDATE_INPUTS("clip.mp4", "input folder", ""), True)
|
|
self.assertIs(frame.VALIDATE_INPUTS("", "local path", str(self.in_output)), True)
|
|
movie = self.touch(self.root / "movie.mp4")
|
|
self.assertIn("Local path mode reads only", frame.VALIDATE_INPUTS("", "local path", str(movie)))
|
|
self.assertIn("Network", frame.VALIDATE_INPUTS("//fileserver/share/clip.mp4", "input folder", ""))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|