diff --git a/comfy_extras/nodes_dataset.py b/comfy_extras/nodes_dataset.py index 73fe75b7f..92f263620 100644 --- a/comfy_extras/nodes_dataset.py +++ b/comfy_extras/nodes_dataset.py @@ -42,6 +42,62 @@ def load_and_process_images(image_files, input_dir): return output_images +def secure_subfolder_path(base_dir, folder_name): + """Resolve folder_name inside base_dir, rejecting anything that escapes it. + + Blocks '..', absolute paths, drive letters and symlink escapes using the + same realpath containment check as the core file endpoints. + """ + target = os.path.abspath(os.path.join(base_dir, folder_name)) + if not folder_paths.is_within_directory(base_dir, target): + raise ValueError(f"Invalid folder name {folder_name!r}: resolves outside of {base_dir}") + return target + + +def list_dataset_folders(): + """Relative paths of dataset folders found under all dataset roots. + + Any subfolder containing a metadata.json or *.safetensors shard counts as + a dataset; the walk doesn't descend into matched folders. + """ + found = set() + for root in folder_paths.get_folder_paths("datasets"): + if not os.path.isdir(root): + continue + for dirpath, subdirs, filenames in os.walk(root, followlinks=True): + if dirpath != root and ( + "metadata.json" in filenames + or any(f.endswith(".safetensors") for f in filenames) + ): + found.add(os.path.relpath(dirpath, root).replace(os.sep, "/")) + subdirs[:] = [] + return sorted(found) + + +def get_dataset_save_dir(folder_name): + """Resolve the folder to save a new dataset into, inside the default root. + + The folder is not created here; callers makedirs after validation. + """ + root = folder_paths.get_folder_paths("datasets")[0] + target = secure_subfolder_path(root, folder_name) + if os.path.realpath(target) == os.path.realpath(root): + raise ValueError("folder_name must name a subfolder of the datasets directory, e.g. 'my_dataset'.") + return target + + +def get_dataset_dir(folder_name): + """Find an existing dataset folder by relative name across all dataset roots.""" + roots = folder_paths.get_folder_paths("datasets") + for root in roots: + target = secure_subfolder_path(root, folder_name) + if os.path.realpath(target) == os.path.realpath(root): + raise ValueError("folder_name must name a subfolder of the datasets directory, e.g. 'my_dataset'.") + if os.path.isdir(target): + return target + raise ValueError(f"Dataset folder {folder_name!r} not found in: {', '.join(roots)}") + + class LoadImageDataSetFromFolderNode(io.ComfyNode): @classmethod def define_schema(cls): @@ -252,7 +308,7 @@ class SaveImageDataSetToFolderNode(io.ComfyNode): filename_prefix = filename_prefix[0] mode = mode[0] - output_dir = os.path.join(folder_paths.get_output_directory(), folder_name) + output_dir = secure_subfolder_path(folder_paths.get_output_directory(), folder_name) saved_files = save_images_to_folder(images, output_dir, filename_prefix, mode=='overwrite') logging.info(f"Saved {len(saved_files)} images to {output_dir}.") @@ -306,7 +362,7 @@ class SaveImageTextDataSetToFolderNode(io.ComfyNode): filename_prefix = filename_prefix[0] mode = mode[0] - output_dir = os.path.join(folder_paths.get_output_directory(), folder_name) + output_dir = secure_subfolder_path(folder_paths.get_output_directory(), folder_name) saved_files = save_images_to_folder(images, output_dir, filename_prefix, mode=='overwrite') # Save captions @@ -1443,7 +1499,7 @@ class SaveTrainingDataset(io.ComfyNode): io.String.Input( "folder_name", default="training_dataset", - tooltip="Name of folder to save dataset (inside output directory).", + tooltip="Name of folder to save the dataset into, inside the datasets directory. Subfolders like 'project/run1' are allowed.", ), io.Int.Input( "shard_size", @@ -1473,8 +1529,8 @@ class SaveTrainingDataset(io.ComfyNode): f"Something went wrong in dataset preparation." ) - # Create output directory - output_dir = os.path.join(folder_paths.get_output_directory(), folder_name) + # Create output directory (inside the datasets root, traversal-safe) + output_dir = get_dataset_save_dir(folder_name) os.makedirs(output_dir, exist_ok=True) # Prepare data pairs @@ -1533,10 +1589,10 @@ class LoadTrainingDataset(io.ComfyNode): description="Load encoded training dataset (latents + conditioning) from disk for use in training.", is_experimental=True, inputs=[ - io.String.Input( + io.Combo.Input( "folder_name", - default="training_dataset", - tooltip="Name of folder containing the saved dataset (inside output directory).", + options=list_dataset_folders(), + tooltip="Saved dataset to load, from the datasets directory.", ), ], outputs=[ @@ -1555,11 +1611,8 @@ class LoadTrainingDataset(io.ComfyNode): @classmethod def execute(cls, folder_name): - # Get dataset directory - dataset_dir = os.path.join(folder_paths.get_output_directory(), folder_name) - - if not os.path.exists(dataset_dir): - raise ValueError(f"Dataset directory not found: {dataset_dir}") + # Get dataset directory (searched across all dataset roots, traversal-safe) + dataset_dir = get_dataset_dir(folder_name) # Find all shard files shard_files = sorted(