Hub documentation

Pickle 扫描

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

Pickle 扫描

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

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

什么是 pickle?

来自官方文档:

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

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

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

从本质上讲,pickle 基本上是一系列指令或操作码(opcodes)的堆栈。正如你可能已经猜到的那样,它并不是人类可读的。操作码在序列化(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 文件,并在终端中打印以下说明:

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

我们可以看到,其中并没有太多内容,只有几个操作码(opcodes)以及与之关联的数据。你可能会想,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 库以简化操作。它允许我们通过 exec 函数添加 pickle 指令,以执行包含在字符串中的代码。这就是你绕过 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×.  │
└────────┴─────────────────────────┴─────────────────────────┴────────┴────────┘

基本上,当你执行反序列化(unpickle)时,发生的就是以下过程:

# ...
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 密钥对提交进行签名。这并不能保证你的文件是安全的,但它可以保证文件的来源。

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

从 TF 或 Flax 加载模型权重

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

E.g.:

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,它允许我们在不执行潜在危险代码的情况下读取文件。

请注意,正是这一点使得我们能够判断,在反序列化文件时,是否会针对由 *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 反序列化反编译器和静态分析工具

利用 Python pickle

cpython/pickletools.py 位于 3.10 · python/cpython

cpython/pickle.py 位于 3.10 · python/cpython

CrypTen/serial.py 位于 main · facebookresearch/CrypTen

CTFtime.org / Balsn CTF 2019 / pyshv1 / 解题报告

修复 Python 的 pickle 模块

Update on GitHub