Comfy-Org / Comfy-Org/ComfyUI

[Security] Unsafe pickle loading in checkpoint_pickle.py allows RCE (Remote Code Execution)

Open
#12,245 0 comments 0 reactions 0 assignees View on GitHub
Potential Bug
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.