Hub documentation
Pickle 扫描
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 文档
GitHub - trailofbits/fickling:一个 Python 反序列化反编译器和静态分析工具
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 / 解题报告
Update on GitHub