# Pickle 扫描

Pickle 是机器学习领域中广泛使用的序列化格式。尤其值得注意的是，它是 PyTorch 模型权重的默认格式。

加载 pickle 文件时存在危险的任意代码执行攻击风险。我们建议从您信任的用户和组织加载模型，依赖已签名的提交，和/或使用 `from_tf=True` 自动转换机制从 TF 或 Jax 格式加载模型。此外，我们通过在 Hub 上直接显示/审查任何 pickle 文件中的导入列表来缓解此问题。最后，我们正在试验一种名为 [`safetensors`](https://github.com/huggingface/safetensors) 的新的、简单的权重序列化格式。

## 什么是 pickle？

来自[官方文档](https://docs.python.org/3/library/pickle.html)：

> `pickle` 模块实现了用于序列化和反序列化 Python 对象结构的二进制协议。

这意味着 pickle 是一种序列化协议，用于在各方之间高效地共享数据。

我们将 pickle 定义为在序列化过程中生成的二进制文件。

从本质上讲，pickle 基本上是一系列指令或操作码（opcodes）的堆栈。正如你可能已经猜到的那样，它并不是人类可读的。操作码在序列化（pickling）时生成，并在反序列化（unpickling）时按顺序读取。根据操作码，会执行相应的操作。

下面是一个简单的示例：

```python
import pickle
import pickletools

var = "data I want to share with a friend"

# store the pickle data in a file named 'payload.pkl'
with open('payload.pkl', 'wb') as f:
    pickle.dump(var, f)

# disassemble the pickle
# and print the instructions to the command line
with open('payload.pkl', 'rb') as f:
    pickletools.dis(f)
```

运行此操作时，将创建一个 pickle 文件，并在终端中打印以下说明：

```python
    0: \x80 PROTO      4
    2: \x95 FRAME      48
   11: \x8c SHORT_BINUNICODE 'data I want to share with a friend'
   57: \x94 MEMOIZE    (as 0)
   58: .    STOP
highest protocol among opcodes = 4
```

现在不必过于担心这些指令，只需知道 [pickletools](https://docs.python.org/3/library/pickletools.html) 模块对于分析 pickle 非常有用。它允许你读取文件中的指令，而**无需**执行任何代码。

Pickle 不仅仅是一种序列化协议，它通过允许用户在反序列化时执行 Python 代码，提供了更高的灵活性。听起来不太安全，对吧？

## 为什么这很危险？

如上所述，反序列化 pickle 意味着可以执行代码。但这存在某些限制：你只能引用顶层模块中的函数和类，无法将它们嵌入到 pickle 文件本身中。

回到绘图板：

```python
import pickle
import pickletools

class Data:
    def __init__(self, important_stuff: str):
        self.important_stuff = important_stuff

d = Data("42")

with open('payload.pkl', 'wb') as f:
    pickle.dump(d, f)
```

运行此脚本后，我们再次得到 `payload.pkl`。检查文件内容时：

```bash

# cat payload.pkl
__main__Data)}important_stuff42sb.%

# hexyl payload.pkl
┌────────┬─────────────────────────┬─────────────────────────┬────────┬────────┐
│00000000│ 80 04 95 33 00 00 00 00 ┊ 00 00 00 8c 08 5f 5f 6d │ו×30000┊000ו__m│
│00000010│ 61 69 6e 5f 5f 94 8c 04 ┊ 44 61 74 61 94 93 94 29 │ain__×ו┊Data×××)│
│00000020│ 81 94 7d 94 8c 0f 69 6d ┊ 70 6f 72 74 61 6e 74 5f │××}×וim┊portant_│
│00000030│ 73 74 75 66 66 94 8c 02 ┊ 34 32 94 73 62 2e       │stuff×ו┊42×sb.  │
└────────┴─────────────────────────┴─────────────────────────┴────────┴────────┘
```

我们可以看到，其中并没有太多内容，只有几个操作码（opcodes）以及与之关联的数据。你可能会想，pickle 有什么问题呢？

让我们尝试其他方法：

```python
from fickling.pickle import Pickled
import pickle

# Create a malicious pickle
data = "my friend needs to know this"

pickle_bin = pickle.dumps(data)

p = Pickled.load(pickle_bin)

p.insert_python_exec('print("you\'ve been pwned !")')

with open('payload.pkl', 'wb') as f:
    p.dump(f)

# innocently unpickle and get your friend's data
with open('payload.pkl', 'rb') as f:
    data = pickle.load(f)
    print(data)
```

这里我们使用 [fickling](https://github.com/trailofbits/fickling) 库以简化操作。它允许我们通过 `exec` 函数添加 pickle 指令，以执行包含在字符串中的代码。这就是你绕过 pickle 中无法定义函数或类这一限制的方法：对保存为字符串的 Python 代码执行 exec。

运行此代码时，会创建一个 `payload.pkl` 并输出以下内容：

```
you've been pwned !
my friend needs to know this
```

如果我们检查 pickle 文件的内容，会得到：

```bash
# cat payload.pkl
c__builtin__
exec
(Vprint("you've been pwned !")
tR my friend needs to know this.%

# hexyl payload.pkl
┌────────┬─────────────────────────┬─────────────────────────┬────────┬────────┐
│00000000│ 63 5f 5f 62 75 69 6c 74 ┊ 69 6e 5f 5f 0a 65 78 65 │c__built┊in___exe│
│00000010│ 63 0a 28 56 70 72 69 6e ┊ 74 28 22 79 6f 75 27 76 │c_(Vprin┊t("you'v│
│00000020│ 65 20 62 65 65 6e 20 70 ┊ 77 6e 65 64 20 21 22 29 │e been p┊wned !")│
│00000030│ 0a 74 52 80 04 95 20 00 ┊ 00 00 00 00 00 00 8c 1c │_tR×•× 0┊000000ו│
│00000040│ 6d 79 20 66 72 69 65 6e ┊ 64 20 6e 65 65 64 73 20 │my frien┊d needs │
│00000050│ 74 6f 20 6b 6e 6f 77 20 ┊ 74 68 69 73 94 2e       │to know ┊this×.  │
└────────┴─────────────────────────┴─────────────────────────┴────────┴────────┘
```

基本上，当你执行反序列化（unpickle）时，发生的就是以下过程：

```python
# ...
opcodes_stack = [exec_func, "malicious argument", "REDUCE"]
opcode = stack.pop()
if opcode == "REDUCE":
    arg = opcodes_stack.pop()
    callable = opcodes_stack.pop()
    opcodes_stack.append(callable(arg))
# ...
```

构成威胁的指令是 `STACK_GLOBAL`、`GLOBAL` 和 `REDUCE`。

`REDUCE` 指示反序列化器使用提供的参数执行函数，而 `*GLOBAL` 指令则指示反序列化器 `import` 相关对象。

总而言之，pickle 之所以危险，是因为：

- 在导入 Python 模块时，可能会执行任意代码
- 您可以导入内置函数，例如 `eval` 或 `exec`，这些函数可用于执行任意代码
- 在实例化对象时，可能会调用构造函数

这就是为什么大多数文档在使用 pickle 时都会强调，不要反序列化来自不可信来源的数据。

## 缓解策略

***请勿使用 pickle***

Luc 的建议很中肯，但 pickle 的使用非常广泛，短期内不会消失：找到一个让所有人都满意的新格式并推动这一变更，需要一些时间。

那么，我们现在能做什么呢？

### 从您信任的用户和组织加载文件

在 Hub 上，你可以[使用 GPG 密钥对提交进行签名](./security-gpg)。这**并不**能保证你的文件是安全的，但它可以保证文件的来源。

如果您了解并信任用户 A，并且 Hub 上包含该文件的提交是由用户 A 的 GPG 密钥签名的，那么可以相当安全地认为您可以信任该文件。

### 从 TF 或 Flax 加载模型权重

TensorFlow 和 Flax 检查点不受影响，可以使用 `from_pretrained` 方法的 `from_tf` 和 `from_flax` 关键字参数在 PyTorch 架构中加载它们，以规避此问题。

E.g.:

```python
from transformers import AutoModel

model = AutoModel.from_pretrained("google-bert/bert-base-cased", from_flax=True)
```

### 使用您自己的序列化格式

- [MsgPack](https://msgpack.org/index.html)
- [Protobuf](https://developers.google.com/protocol-buffers)
- [Cap'n'proto](https://capnproto.org/)
- [Avro](https://avro.apache.org/)
- [safetensors](https://github.com/huggingface/safetensors)

最后这种格式 `safetensors` 是一种简单的序列化格式，我们目前正在对其进行开发和实验！如果你能帮忙或贡献代码，请不吝赐教 🔥。

### 改进 `torch.load/save`

PyTorch 上正在进行一项公开讨论，主题是[默认情况下安全地仅从 *.pt 文件加载权重](https://github.com/pytorch/pytorch/issues/52181)——欢迎前往参与讨论！

### Hub 的安全扫描器

#### 我们目前拥有的

我们创建了一个安全扫描器，它会扫描推送到 Hub 的每个文件并执行安全检查。在撰写本文时，它运行两种类型的扫描：

- ClamAV 扫描
- Pickle 导入扫描

对于 ClamAV 扫描，文件会通过开源杀毒软件 [ClamAV](https://www.clamav.net) 进行处理。虽然这能覆盖大量危险文件，但无法覆盖 pickle 漏洞利用。

我们实现了一种 Pickle 导入扫描，用于提取 pickle 文件中引用的导入列表。每次您上传 `pytorch_model.bin` 或任何其他 pickle 文件时，都会运行此扫描。

在 Hub 上，包含导入语句的每个文件旁边都会显示导入列表。如果某个导入看起来可疑，它将被高亮显示。

我们之所以能获取这些数据，要归功于 [`pickletools.genops`](https://docs.python.org/3/library/pickletools.html#pickletools.genops)，它允许我们在不执行潜在危险代码的情况下读取文件。

请注意，正是这一点使得我们能够判断，在反序列化文件时，是否会针对由 `*GLOBAL` 导入的潜在危险函数执行 `REDUCE` 操作。

***免责声明***：此方法并非 100% 可靠。作为用户，您有责任自行检查内容是否安全。我们不会主动对 Python 包进行安全审计，现有的安全/不安全导入列表仅以尽力而为的方式维护。如果您认为某些内容不安全，请通过发送邮件至 website at huggingface.co 联系我们，我们会将其标记为不安全。

#### 潜在解决方案

人们可能会想到创建一个自定义的 [Unpickler](https://docs.python.org/3/library/pickle.html#pickle.Unpickler)，类似于 [这个](https://github.com/facebookresearch/CrypTen/blob/main/crypten/common/serial.py)。但正如我们在 [这个复杂的利用](https://ctftime.org/writeup/16723) 中所看到的，这样做是行不通的。

幸运的是，`eval` 导入总会留下痕迹，因此直接读取操作码应该能够捕获恶意使用。

我目前提出的解决方案是创建一个类似于 `.gitignore` 的文件，但用于管理导入。

此文件将作为导入白名单，如果存在未包含在白名单中的导入，则会将 `pytorch_model.bin` 文件标记为危险。

可以设想一种类似正则表达式的格式，例如通过一行简单的 `numpy.*` 来允许所有 numpy 子模块。

## 延伸阅读

[pickle - Python 对象序列化 - Python 3.10.6 文档](https://docs.python.org/3/library/pickle.html#what-can-be-pickled-and-unpickled)

[危险的 Pickles - 恶意的 Python 序列化](https://intoli.com/blog/dangerous-pickles/)

[GitHub - trailofbits/fickling：一个 Python 反序列化反编译器和静态分析工具](https://github.com/trailofbits/fickling)

[利用 Python pickle](https://davidhamann.de/2020/04/05/exploiting-python-pickle/)

[cpython/pickletools.py 位于 3.10 · python/cpython](https://github.com/python/cpython/blob/3.10/Lib/pickletools.py)

[cpython/pickle.py 位于 3.10 · python/cpython](https://github.com/python/cpython/blob/3.10/Lib/pickle.py)

[CrypTen/serial.py 位于 main · facebookresearch/CrypTen](https://github.com/facebookresearch/CrypTen/blob/main/crypten/common/serial.py)

[CTFtime.org / Balsn CTF 2019 / pyshv1 / 解题报告](https://ctftime.org/writeup/16723)

[修复 Python 的 pickle 模块](https://github.com/moreati/pickle-fuzz)