Pickle 扫描

Pickle 是 ML 中广泛使用的序列化格式。最值得注意的是，它是 PyTorch 模型权重的默认格式。

当你加载 pickle 文件时，可能会发生危险的任意代码执行攻击。我们建议从你信任的用户和组织加载模型，依赖签名提交，和/或使用 from_tf=True 自动转换机制从 TF 或 Jax 格式加载模型。我们还通过在 Hub 上直接显示/"审查"任何 pickled 文件中的导入列表来缓解此问题。最后，我们正在试验一种新的简单权重序列化格式，称为 safetensors。

什么是 pickle？

来自官方文档：

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

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

我们将 pickling 时生成的二进制文件称为 pickle。

就其核心而言，pickle 基本上是一堆指令或操作码。正如你可能已经猜到的那样，它不是人类可读的。操作码在 pickling 时生成，在 unpickling 时按顺序读取。根据操作码，执行给定的操作。

以下是一个小示例：

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 文件并在终端中打印以下指令：

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

现在不要太担心指令，只需知道 pickletools 模块对于分析 pickle 非常有用。它允许你读取文件中的指令，而无需执行任何代码。

Pickle 不仅仅是序列化协议，它通过在反序列化时允许用户运行 python 代码来提供更大的灵活性。听起来不太好，对吧？

为什么它很危险？

正如我们上面所说，反序列化 pickle 意味着可以执行代码。但这有一定的限制：你只能引用顶级模块中的函数和类；你不能将它们嵌入到 pickle 文件本身中。

回到绘图板：

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。当我们检查文件内容时：

# 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.  │
└────────┴─────────────────────────┴─────────────────────────┴────────┴────────┘

我们可以看到里面没有太多内容，只有一些操作码和关联数据。你可能在想，那么 pickle 有什么问题呢？

让我们尝试其他方法：

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 库来简化操作。它允许我们添加 pickle 指令，通过 exec 函数执行包含在字符串中的代码。这就是你如何规避无法在 pickle 中定义函数或类的事实：你在保存为字符串的 python 代码上运行 exec。

运行此代码时，它会创建一个 payload.pkl 并打印以下内容：

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

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

# 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×.  │
└────────┴─────────────────────────┴─────────────────────────┴────────┴────────┘

基本上，这就是当你反序列化时发生的情况：

# ...
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 告诉 unpickler 使用提供的参数执行函数，而 *GLOBAL 指令告诉 unpickler import 内容。

总而言之，pickle 很危险，因为：

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

这就是为什么在使用 pickle 的大多数文档中都说明，不要从不受信任的源反序列化数据。

缓解策略

不要使用 pickle

好建议 Luc，但 pickle 被广泛使用，短期内不会消失：找到一个每个人都满意的新格式并启动更改需要一些时间。

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

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

在 Hub 上，你可以使用 GPG 密钥签名你的提交。这不能保证你的文件是安全的，但它确实保证了文件的来源。

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

从 TF 或 Flax 加载模型权重

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

例如：

from transformers import AutoModel

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

使用你自己的序列化格式

最后一种格式 safetensors 是我们目前正在开发和试验的简单序列化格式！如果你能提供帮助或贡献，请参与 🔥。

改进 `torch.load/save`

PyTorch 正在进行关于默认安全地从 *.pt 文件加载仅权重的安全方法的公开讨论 - 请在那里发表意见！

Hub 的安全扫描器

我们现在拥有的

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

ClamAV 扫描
Pickle 导入扫描

对于 ClamAV 扫描，文件通过开源防病毒软件 ClamAV 运行。虽然这涵盖了相当多的危险文件，但它不涵盖 pickle 漏洞利用。

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

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

我们通过 pickletools.genops 获得这些数据，它允许我们读取文件而无需执行潜在危险的代码。

请注意，这允许我们知道，在反序列化文件时，它是否会在由 *GLOBAL 导入的潜在危险函数上执行 REDUCE。

免责声明：这不是 100% 万无一失的。作为用户，你有责任检查某些内容是否安全。我们没有主动审计 python 包的安全性，我们拥有的安全/不安全导入列表是以最佳努力方式维护的。如果你认为某些内容不安全，而我们将其标记为不安全，请通过发送电子邮件至 [email protected] 联系我们