Hub 文档
Pickle 扫描
并获得增强的文档体验
开始使用
Pickle 扫描
Pickle 是机器学习中广泛使用的序列化格式。最值得注意的是，它是 PyTorch 模型权重的默认格式。
当你加载一个 pickle 文件时，可能会遭受危险的任意代码执行攻击。我们建议从你信任的用户和组织加载模型，依赖经过签名的提交，和/或使用 from_tf=True 自动转换机制从 TF 或 Jax 格式加载模型。我们还通过直接在 Hub 上显示/“审查”任何 pickle 文件中的导入列表来缓解此问题。最后，我们正在试验一种名为 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 文件并在你的终端打印以下指令：
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×. │
└────────┴─────────────────────────┴─────────────────────────┴────────┴────────┘基本上，这就是 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 告诉 unpickler 用提供的参数执行函数，而 *GLOBAL 指令告诉 unpickler 去 `import` 东西。
总而言之，pickle 是危险的，因为：
- 当导入一个 python 模块时，可以执行任意代码
- 你可以导入像 `eval` 或 `exec` 这样的内置函数，它们可以用来执行任意代码
- 当实例化一个对象时，构造函数可能会被调用
这就是为什么大多数使用 pickle 的文档中都说明，不要 unpickle 来自不受信任来源的数据。
缓解策略
不要使用 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` 或任何其他 pickle 文件时，都会运行此扫描。
在 Hub 上，导入列表将显示在每个包含导入的文件旁边。如果任何导入看起来可疑，它将被高亮显示。
我们通过 pickletools.genops 获得这些数据，它允许我们读取文件内容而无需执行潜在的危险代码。
请注意，这使得我们能够知道，当 unpickling 文件时，它是否会对一个通过 `*GLOBAL` 导入的潜在危险函数执行 `REDUCE` 操作。
免责声明：这并非 100% 万无一失。作为用户，你有责任检查某些内容是否安全。我们不会主动审计 python 包的安全性，我们的安全/不安全导入列表是以尽力而为的方式维护的。如果你认为某些东西不安全，而我们将其标记为安全，请通过发送电子邮件至 website at huggingface.co 与我们联系。
潜在解决方案
有人可能会想到创建一个自定义的 Unpickler，类似于这个。但正如我们在这个复杂的漏洞利用中看到的，这行不通。
幸运的是，`eval` 导入总会留下痕迹，所以直接读取操作码应该能捕捉到恶意使用。
我目前提出的解决方案是创建一个类似于 `.gitignore` 的文件，但用于导入。
这个文件将是一个导入的白名单，如果 `pytorch_model.bin` 文件中有未包含在白名单中的导入，该文件将被标记为危险。
可以想象使用一种类似正则表达式的格式，例如通过一行简单的 `numpy.*` 来允许所有的 numpy 子模块。
进一步阅读
pickle - Python object serialization - Python 3.10.6 documentation
Dangerous Pickles - Malicious Python Serialization
GitHub - trailofbits/fickling: A Python pickling decompiler and static analyzer
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
Rehabilitating Python’s pickle module
< > 在 GitHub 上更新