From 8088b347d0aeff7ac0319996d351eec62450a7ba Mon Sep 17 00:00:00 2001 From: jakelodwick Date: Sun, 5 Apr 2026 19:13:08 -0600 Subject: [PATCH 1/2] Fix torch.load missing weights_only in LoadTrainingDataset The rest of the codebase uses weights_only=True. Also adds a path check on folder_name. --- comfy_extras/nodes_dataset.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/comfy_extras/nodes_dataset.py b/comfy_extras/nodes_dataset.py index 98ed25d7e..f7d9afbff 100644 --- a/comfy_extras/nodes_dataset.py +++ b/comfy_extras/nodes_dataset.py @@ -1450,7 +1450,11 @@ 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) + output_dir = folder_paths.get_output_directory() + dataset_dir = os.path.join(output_dir, folder_name) + # Prevent path traversal (e.g. folder_name="../../etc") + if not os.path.realpath(dataset_dir).startswith(os.path.realpath(output_dir)): + raise ValueError(f"Invalid folder_name: path traversal detected") if not os.path.exists(dataset_dir): raise ValueError(f"Dataset directory not found: {dataset_dir}") @@ -1477,7 +1481,7 @@ class LoadTrainingDataset(io.ComfyNode): shard_path = os.path.join(dataset_dir, shard_file) with open(shard_path, "rb") as f: - shard_data = torch.load(f) + shard_data = torch.load(f, weights_only=True) all_latents.extend(shard_data["latents"]) all_conditioning.extend(shard_data["conditioning"]) From cf43a3a63e633c9e77bde0a6e1d8ef6fc0ed4f9d Mon Sep 17 00:00:00 2001 From: jakelodwick Date: Sun, 5 Apr 2026 19:16:52 -0600 Subject: [PATCH 2/2] Use commonpath for path traversal check Matches the pattern used in folder_paths.py and server.py. The startswith approach is vulnerable to sibling directory bypasses. --- comfy_extras/nodes_dataset.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/comfy_extras/nodes_dataset.py b/comfy_extras/nodes_dataset.py index f7d9afbff..6a64b2166 100644 --- a/comfy_extras/nodes_dataset.py +++ b/comfy_extras/nodes_dataset.py @@ -1453,7 +1453,9 @@ class LoadTrainingDataset(io.ComfyNode): output_dir = folder_paths.get_output_directory() dataset_dir = os.path.join(output_dir, folder_name) # Prevent path traversal (e.g. folder_name="../../etc") - if not os.path.realpath(dataset_dir).startswith(os.path.realpath(output_dir)): + real_output_dir = os.path.realpath(output_dir) + real_dataset_dir = os.path.realpath(dataset_dir) + if os.path.commonpath((real_output_dir, real_dataset_dir)) != real_output_dir: raise ValueError(f"Invalid folder_name: path traversal detected") if not os.path.exists(dataset_dir):