Hub 文档
Pickle 扫描
并获得增强的文档体验
开始使用
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_GLOBAL
、GLOBAL
和 REDUCE
。
REDUCE
告诉反 pickle 器使用提供的参数执行函数,而 *GLOBAL
指令告诉反 pickle 器 import
内容。
总而言之,pickle 是危险的,因为
- 当导入 python 模块时,可以执行任意代码
- 您可以导入内置函数,如
eval
或exec
,它们可用于执行任意代码 - 当实例化对象时,可能会调用构造函数
这就是为什么在使用 pickle 的大多数文档中都声明,不要反 pickle 来自不受信任来源的数据。
缓解策略
不要使用 pickle
Luc 的建议听起来不错,但 pickle 被广泛使用,并且不会很快消失:找到每个人都满意的新格式并启动更改将需要一些时间。
那么我们现在可以做什么呢?
从您信任的用户和组织加载文件
在 Hub 上,您可以使用 GPG 密钥签署您的提交。这并不能保证您的文件是安全的,但它可以保证文件的来源。
如果您认识并信任用户 A,并且 Hub 上包含该文件的提交由用户 A 的 GPG 密钥签名,那么可以相当安全地假设您可以信任该文件。
从 TF 或 Flax 加载模型权重
TensorFlow 和 Flax 检查点不受影响,可以使用 from_tf
和 from_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 文档
GitHub - trailofbits/fickling: 一个 Python pickle 反编译器和静态分析器
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
< > 在 GitHub 上更新