[Security] Unsafe pickle loading in checkpoint_pickle.py allows RCE (Remote Code Execution)
- Dominant language
- Python
- Stars
- 133k
- Forks
- 15.7k
- Avg merge
- 1d 7h
- Merged PRs (30d)
- 158
Description
### Custom Node Testing
- [x] I have tried disabling custom nodes and the issue persists (see [how to disable custom nodes](https://docs.comfy.org/troubleshooting/custom-node-issues#step-1%3A-test-with-all-custom-nodes-disabled) if you need help)
### Expected Behavior
When loading a PyTorch checkpoint (`.pt`, `.ckpt`, `.bin`), the system should restrict the unpickler to only load safe modules (like `torch`, `numpy`, `collections`).
It should block dangerous globals (like `os.system`, `subprocess.Popen`, `eval`) to prevent malicious models from executing code on the user's machine.
The `Unpickler` class in `comfy/checkpoint_pickle.py` should implement the logic hinted at in the comment `#TODO: safe unpickle`.
### Actual Behavior
The current implementation of `Unpickler.find_class` in `comfy/checkpoint_pickle.py` only filters out `pytorch_lightning`.
For all other modules, it falls back to `super().find_class(module, name)`, which allows loading arbitrary Python objects.
If a user downloads a malicious checkpoint (e.g., from Civitai or a compromised HF repo) that contains a pickle payload with `os.system`, ComfyUI will execute that command immediately upon loading the checkpoint, giving the attacker control over the system.
### Steps to Reproduce
1. Create a standard Python pickle file containing a malicious payload (e.g., printing a message or creating a file).
```python
import pickle, os
class RCE:
def __reduce__(self):
return (os.system, ("echo HACKED",))
torch.save(RCE(), "malicious_model.pt")
1. Place `malicious_model.pt` into the `models/checkpoints/`directory.
2. Start ComfyUI.
3. Load the checkpoint using a "Load Checkpoint" node.
4. Result: The code `echo HACKED` is executed in the terminal.
### Debug Logs
```powershell
N/A - This is a security vulnerability found via static code analysis of `comfy/checkpoint_pickle.py`.
```
### Other
**Code Reference:**
https://github.com/comfyanonymous/ComfyUI/blob/master/comfy/checkpoint_pickle.py#L18
```python
# comfy/checkpoint_pickle.py
class Unpickler(pickle.Unpickler):
def find_class(self, module, name):
#TODO: safe unpickle <-- The vulnerability is here
if module.startswith("pytorch_lightning"):
return Empty
return super().find_class(module, name) # Allows RCE
**Proposed Solution:**
I maintain an open-source security tool called [Veritensor](https://github.com/ArseniiBrazhnyk/Veritensor) (Apache 2.0) that solves this exact problem. It implements a secure Pickle engine with stack emulation and a strict allowlist.
You can fix this by adding a pre-flight check in `comfy/utils.py` before `torch.load` is called, or by updating `checkpoint_pickle.py` to use a strict allowlist similar to Veritensor's engine.
I've attached an example of the full code.
[ComfyUI_scan.py](https://github.com/user-attachments/files/25047323/ComfyUI_scan.py)
Contributor guide
Assessment
This issue has not been assessed yet.