diff --git a/AGENTS.md b/AGENTS.md index bfe0976fd..6682f71a4 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -305,6 +305,12 @@ - Follow existing node conventions: `INPUT_TYPES`, `RETURN_TYPES`, `FUNCTION`, `CATEGORY`, and registration through the local mapping used by that file. +- Treat legacy combo inputs, `io.Combo`, and `io.DynamicCombo` values as + untrusted when they affect filesystem access. Any value used as a file or + folder name, path component, format, or extension must be validated again at + the load/save boundary using an existing `folder_paths` resolver or + containment helper, or a fixed allowlist/mapping. Do not rely only on the + advertised combo options or prompt validation. - Keep node changes backward compatible by default. Add inputs with sensible defaults and avoid changing output types unless the request requires it. - Model implementations should add the minimal number of ComfyUI nodes required diff --git a/comfy_api/latest/_ui.py b/comfy_api/latest/_ui.py index b48713d41..f1f359e6a 100644 --- a/comfy_api/latest/_ui.py +++ b/comfy_api/latest/_ui.py @@ -260,6 +260,7 @@ class ImageSaveHelper: class AudioSaveHelper: """A helper class with static methods to handle audio saving and metadata.""" _OPUS_RATES = [8000, 12000, 16000, 24000, 48000] + _FORMATS = {"flac", "mp3", "opus"} @staticmethod def save_audio( @@ -270,6 +271,9 @@ class AudioSaveHelper: format: str = "flac", quality: str = "128k", ) -> list[SavedResult]: + if format not in AudioSaveHelper._FORMATS: + raise ValueError(f"Unsupported audio format: {format!r}") + full_output_folder, filename, counter, subfolder, _ = folder_paths.get_save_image_path( filename_prefix, _get_directory_by_folder_type(folder_type) ) diff --git a/comfy_extras/nodes_dataset.py b/comfy_extras/nodes_dataset.py index 5e0454d8b..71e5ee368 100644 --- a/comfy_extras/nodes_dataset.py +++ b/comfy_extras/nodes_dataset.py @@ -195,7 +195,7 @@ class LoadImageDataSetFromFolderNode(io.ComfyNode): @classmethod def execute(cls, folder): - sub_input_dir = os.path.join(folder_paths.get_input_directory(), folder) + sub_input_dir = secure_subfolder_path(folder_paths.get_input_directory(), folder) valid_extensions = [".png", ".jpg", ".jpeg", ".webp"] image_files = [ f @@ -241,7 +241,7 @@ class LoadImageTextDataSetFromFolderNode(io.ComfyNode): def execute(cls, folder): logging.info(f"Loading images from folder: {folder}") - sub_input_dir = os.path.join(folder_paths.get_input_directory(), folder) + sub_input_dir = secure_subfolder_path(folder_paths.get_input_directory(), folder) valid_extensions = [".png", ".jpg", ".jpeg", ".webp"] image_files = [] @@ -310,7 +310,7 @@ class LoadVideoDataSetFromFolderNode(io.ComfyNode): @classmethod def execute(cls, folder): - sub_input_dir = os.path.join(folder_paths.get_input_directory(), folder) + sub_input_dir = secure_subfolder_path(folder_paths.get_input_directory(), folder) video_files = sorted([ f for f in os.listdir(sub_input_dir) if any(f.lower().endswith(ext) for ext in VALID_VIDEO_EXTENSIONS) @@ -357,7 +357,7 @@ class LoadVideoTextDataSetFromFolderNode(io.ComfyNode): @classmethod def execute(cls, folder): - sub_input_dir = os.path.join(folder_paths.get_input_directory(), folder) + sub_input_dir = secure_subfolder_path(folder_paths.get_input_directory(), folder) video_files = [] for item in sorted(os.listdir(sub_input_dir)): diff --git a/comfy_extras/nodes_gaussian_splat.py b/comfy_extras/nodes_gaussian_splat.py index 116c14fde..a53317b4e 100644 --- a/comfy_extras/nodes_gaussian_splat.py +++ b/comfy_extras/nodes_gaussian_splat.py @@ -471,6 +471,12 @@ def _mat_to_quat(m): class SplatToFile3D(IO.ComfyNode): + FORMAT_WRITERS = { + "ply": _gaussian_ply_bytes, + "ksplat": _gaussian_ksplat_bytes, + "spz": _gaussian_spz_bytes, + } + @classmethod def define_schema(cls): return IO.Schema( @@ -482,7 +488,7 @@ class SplatToFile3D(IO.ComfyNode): "Supports one item per batch only.", inputs=[ IO.Splat.Input("splat"), - IO.Combo.Input("format", options=["ply", "ksplat", "spz"], # TODO: add "splat" when we have a writer for it + IO.Combo.Input("format", options=list(cls.FORMAT_WRITERS), # TODO: add "splat" when we have a writer for it tooltip="ply: standard 3D Gaussian Splat with full spherical harmonics. " "ksplat: mkkellogg SplatBuffer (level 0, uncompressed), base color only " "spz: Niantic gzip-compressed (~10x smaller), base color only " @@ -493,10 +499,13 @@ class SplatToFile3D(IO.ComfyNode): @classmethod def execute(cls, splat, format="ply") -> IO.NodeOutput: + writer = cls.FORMAT_WRITERS.get(format) + if writer is None: + raise ValueError(f"Unsupported splat format: {format!r}") + if splat.positions.shape[0] > 1: logging.warning("SplatToFile3D supports one item per batch only. Got %d; using first.", splat.positions.shape[0]) end = _real_len(splat, 0) - writer = {"ksplat": _gaussian_ksplat_bytes, "spz": _gaussian_spz_bytes}.get(format, _gaussian_ply_bytes) data = writer(splat.positions[0, :end], splat.scales[0, :end], splat.rotations[0, :end], splat.opacities[0, :end], splat.sh[0, :end]) return IO.NodeOutput(Types.File3D(BytesIO(data), file_format=format)) diff --git a/nodes.py b/nodes.py index de86e93fe..08529d8dc 100644 --- a/nodes.py +++ b/nodes.py @@ -633,15 +633,18 @@ class DiffusersLoader: SEARCH_ALIASES = ["load diffusers model"] @classmethod - def INPUT_TYPES(cls): + def _model_paths(cls): paths = [] for search_path in folder_paths.get_folder_paths("diffusers"): if os.path.exists(search_path): for root, subdir, files in os.walk(search_path, followlinks=True): if "model_index.json" in files: paths.append(os.path.relpath(root, start=search_path)) + return paths - return {"required": {"model_path": (paths,), }} + @classmethod + def INPUT_TYPES(cls): + return {"required": {"model_path": (cls._model_paths(),), }} RETURN_TYPES = ("MODEL", "CLIP", "VAE") FUNCTION = "load_checkpoint" DEPRECATED = True @@ -649,14 +652,20 @@ class DiffusersLoader: CATEGORY = "model/loaders" def load_checkpoint(self, model_path, output_vae=True, output_clip=True): + if model_path not in self._model_paths(): + raise ValueError(f"Invalid diffusers model path: {model_path!r}") + + resolved_model_path = None for search_path in folder_paths.get_folder_paths("diffusers"): if os.path.exists(search_path): path = os.path.join(search_path, model_path) - if os.path.exists(path): - model_path = path + if os.path.isfile(os.path.join(path, "model_index.json")): + resolved_model_path = path break + if resolved_model_path is None: + raise FileNotFoundError(f"Diffusers model {model_path!r} not found.") - return comfy.diffusers_load.load_diffusers(model_path, output_vae=output_vae, output_clip=output_clip, embedding_directory=folder_paths.get_folder_paths("embeddings")) + return comfy.diffusers_load.load_diffusers(resolved_model_path, output_vae=output_vae, output_clip=output_clip, embedding_directory=folder_paths.get_folder_paths("embeddings")) class unCLIPCheckpointLoader: