Hub 文档

Pickle 扫描

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

Pickle 扫描

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

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

什么是 Pickle?

来自 官方文档

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

这意味着 pickle 是一种序列化协议,您可以使用它在各方之间高效地共享数据。

我们称 pickle 为在 pickle 过程中生成的二进制文件。

pickle 的核心基本上是一堆指令或操作码。您可能已经猜到,它不是人类可读的。操作码在 pickle 时生成,并在反 pickle 时按顺序读取。根据操作码,执行给定的操作。

这是一个小例子

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

    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 模块对于分析 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×.  │
└────────┴─────────────────────────┴─────────────────────────┴────────┴────────┘

基本上,这是您在反 pickle 时发生的事情

# ...
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_GLOBALGLOBALREDUCE

REDUCE 告诉反 pickle 器使用提供的参数执行函数,而 *GLOBAL 指令告诉反 pickle 器 import 内容。

总而言之,pickle 是危险的,因为

  • 当导入 python 模块时,可以执行任意代码
  • 您可以导入内置函数,如 evalexec,它们可用于执行任意代码
  • 当实例化对象时,可能会调用构造函数

这就是为什么在使用 pickle 的大多数文档中都声明,不要反 pickle 来自不受信任来源的数据。

缓解策略

不要使用 pickle

Luc 的建议听起来不错,但 pickle 被广泛使用,并且不会很快消失:找到每个人都满意的新格式并启动更改将需要一些时间。

那么我们现在可以做什么呢?

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

在 Hub 上,您可以使用 GPG 密钥签署您的提交。这并不能保证您的文件是安全的,但它可以保证文件的来源。

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

从 TF 或 Flax 加载模型权重

TensorFlow 和 Flax 检查点不受影响,可以使用 from_tffrom_flax kwargs 为 from_pretrained 方法加载它们到 PyTorch 架构中,以规避此问题。

例如:

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 或任何其他 pickle 文件时,都会运行此扫描。

在 hub 上,导入列表将显示在每个包含导入的文件旁边。如果任何导入看起来可疑,它将被突出显示。

我们通过 pickletools.genops 获取此数据,它允许我们在不执行潜在危险代码的情况下读取文件。

请注意,这使我们能够知道,当反 pickle 文件时,它是否会在由 *GLOBAL 导入的潜在危险函数上 REDUCE

免责声明:这并非 100% 可靠。作为用户,您有责任检查某事物是否安全。我们不会主动审核 python 包的安全性,我们维护的安全/不安全导入列表是以尽力而为的方式进行的。如果您认为某些内容不安全,并且我们将其标记为不安全,请通过发送电子邮件至 website at huggingface.co 联系我们

潜在的解决方案

有人可能会想到创建一个自定义的 Unpickler,类似于这个。但正如我们在这个复杂的漏洞利用中看到的那样,这行不通。

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

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

这个文件将是一个导入白名单,如果 pytorch_model.bin 文件中存在未包含在白名单中的导入,则该文件将被标记为危险。

可以想象使用类似正则表达式的格式,例如,您可以通过像 numpy.* 这样的简单行来允许所有 numpy 子模块。

延伸阅读

pickle - Python 对象序列化 - Python 3.10.6 文档

危险的 Pickles - 恶意的 Python 序列化

GitHub - trailofbits/fickling: 一个 Python pickle 反编译器和静态分析器

利用 Python pickles

cpython/pickletools.py at 3.10 · python/cpython

cpython/pickle.py at 3.10 · python/cpython

CrypTen/serial.py at main · facebookresearch/CrypTen

CTFtime.org / Balsn CTF 2019 / pyshv1 / Writeup

修复 Python 的 pickle 模块

< > 在 GitHub 上更新