# ComfyUI — Remote Code Execution via Unsafe Deserialization in LoadTrainingDataset

## meta
platform: huntr
program: ComfyUI
asset: https://github.com/comfyanonymous/ComfyUI
date: 2026-02-13
status: DRAFT

````
Repository URL: https://github.com/comfyanonymous/ComfyUI
Package Manager: pip
Version Affected: 0.13.0 (latest)
Vulnerability Type: Deserialization of Untrusted Data
CVSS:
  - Attack Vector: Network
  - Attack Complexity: Low
  - Privileges Required: None
  - User Interaction: None
  - Scope: Unchanged
  - Confidentiality: High
  - Integrity: High
  - Availability: High
Title: Remote Code Execution via unsafe torch.load() in LoadTrainingDataset node
Description:

# Description

The `LoadTrainingDataset` node in `comfy_extras/nodes_dataset.py` uses `torch.load(f)` (line 1472) **without** `weights_only=True` to deserialize `.pkl` shard files. This allows arbitrary Python code execution through crafted serialized payloads.

This is an inconsistency in the codebase's security posture. Every other `torch.load` call in ComfyUI uses safe deserialization:
- `comfy/utils.py:144` uses `weights_only=True`
- `comfy/utils.py:147` uses `pickle_module=comfy.checkpoint_pickle` (restricted unpickler)
- `comfy/sd1_clip.py:439` uses `weights_only=True`
- `nodes.py:546` and `sd1_clip.py:436` use `safetensors.torch.load_file()` (safe format)

The `LoadTrainingDataset` node was the only one missed.

Additionally, the `folder_name` parameter has **zero path traversal validation**. It is joined directly with the output directory via `os.path.join(folder_paths.get_output_directory(), folder_name)` (line 1445). An attacker can use `../` sequences to load `.pkl` files from any directory on the filesystem.

The companion `SaveTrainingDataset` node (line 1364) has the same path traversal issue, allowing an attacker to write arbitrary serialized data to any writable directory.

# Proof of Concept

1. Generate a malicious serialized payload that matches the expected shard format, embedding an OS command in the deserialized object's `__reduce__` method.

2. Place the payload at a known path (e.g., `/tmp/evil_dataset/shard_0000.pkl`), or use the `SaveTrainingDataset` node's path traversal to write it there.

3. Submit a workflow prompt that loads the malicious dataset:

```bash
curl -X POST http://localhost:8188/prompt -H "Content-Type: application/json" -d '{
  "prompt": {
    "1": {
      "class_type": "LoadTrainingDataset",
      "inputs": {
        "folder_name": "../../../tmp/evil_dataset"
      }
    }
  }
}'
```

4. ComfyUI processes the workflow, calls `torch.load(f)` on the malicious `.pkl` file, executing the embedded Python code.

5. Verify code execution occurred (e.g., check for created files, reverse shell, etc.).

**Impact**: This vulnerability is capable of achieving full Remote Code Execution on the ComfyUI server. An attacker who can submit workflow prompts (default: any user with network access, no authentication required) can execute arbitrary commands with the privileges of the ComfyUI process. Combined with the path traversal in `SaveTrainingDataset`, an attacker can also write malicious files to arbitrary locations before loading them.

**Occurrences**:
- Permalink: https://github.com/comfyanonymous/ComfyUI/blob/e1add563f9e89026e8c4e8825a2b279fbd67d23a/comfy_extras/nodes_dataset.py#L1472
  - Description: `torch.load(f)` without `weights_only=True` — unsafe deserialization of user-controlled `.pkl` files. Every other `torch.load` in the codebase uses safe loading.
- Permalink: https://github.com/comfyanonymous/ComfyUI/blob/e1add563f9e89026e8c4e8825a2b279fbd67d23a/comfy_extras/nodes_dataset.py#L1445
  - Description: `os.path.join(folder_paths.get_output_directory(), folder_name)` — no path traversal validation on `folder_name`, allows loading from arbitrary directories
- Permalink: https://github.com/comfyanonymous/ComfyUI/blob/e1add563f9e89026e8c4e8825a2b279fbd67d23a/comfy_extras/nodes_dataset.py#L1364
  - Description: `SaveTrainingDataset` has same path traversal — `os.path.join(folder_paths.get_output_directory(), folder_name)` allows writing malicious `.pkl` files to arbitrary directories

**References**:
- https://cwe.mitre.org/data/definitions/502.html — CWE-502: Deserialization of Untrusted Data
- https://pytorch.org/docs/stable/generated/torch.load.html — PyTorch docs warning about unsafe torch.load
- https://cwe.mitre.org/data/definitions/22.html — CWE-22: Path Traversal
- https://owasp.org/www-community/vulnerabilities/Deserialization_of_untrusted_data — OWASP Deserialization
````