永不安全的时刻：利用机器学习pickle文件

许多机器学习（ML）模型本质上都是Python pickle文件，这种做法有其合理性。使用pickle可以节省内存、支持启停式模型训练，并使训练好的模型具有可移植性（从而可共享）。Pickle易于实现，内置于Python无需额外依赖，并支持自定义对象的序列化。毫无疑问，选择pickle进行持久化是Python程序员和ML从业者的流行做法。

预训练模型通常被视为ML的“免费”副产品，因为它们允许产生模型的有价值知识产权（如算法和语料库）保持私有。这使许多人放心地在互联网上共享模型，特别是可复用的计算机视觉和自然语言处理分类器。像PyTorch Hub这样的网站促进了模型共享，有些库甚至提供API来自动从GitHub仓库下载模型。

在此，我们讨论仅通过加载不受信任的pickle文件或ML模型可能发生的隐秘诡计。在此过程中，我们介绍一个新工具Fickling，它可以帮助您逆向工程、测试甚至创建恶意pickle文件。如果您是ML从业者，您将了解标准ML实践中固有的安全风险。如果您是安全工程师，您将了解一个可以帮助您构建和取证检查pickle文件的新工具。无论哪种方式，到本文结束时，pickling hopefully会在您口中留下酸涩的味道。

您知道pickle是如何存储的吗？令人震惊！

Python pickle是在称为Pickle Machine（PM）的独特虚拟机中运行的编译程序。PM解释pickle文件的操作码序列以构建任意复杂的Python对象。Python pickle也是一种流格式，允许PM在通过网络下载或从文件读取pickle部分时逐步构建结果对象。

PM使用哈佛架构，将程序操作码与可写数据内存分离，从而防止自修改代码和内存破坏攻击。它也不支持条件、循环甚至算术。在反序列化期间，PM读入pickle程序并执行一系列指令。一旦到达STOP操作码就停止，此时栈顶的任何对象就是反序列化的最终结果。

从这个描述中，人们可能合理得出结论：PM不是图灵完备的。这种格式怎么可能不安全？为了腐蚀三岛由纪夫著名格言的词语：

计算机程序是一种将现实简化为抽象以传输给我们理性的媒介，而在其腐蚀现实的力量中不可避免地潜伏着怪异机器的危险。

PM包含两个可以在PM之外执行任意Python代码的操作码，将结果推入PM的栈：GLOBAL和REDUCE。GLOBAL用于导入Python模块或类，REDUCE用于将一组参数应用于可调用对象，通常先前通过GLOBAL导入。即使pickle文件不使用REDUCE操作码，仅导入模块的行为就可以且将会执行该模块中的任意代码，因此仅GLOBAL就很危险。

例如，可以使用GLOBAL从__builtins__导入exec函数，然后使用REDUCE调用exec，并附带包含要运行的Python代码的任意字符串。对于其他敏感函数如os.system和subprocess.call也是如此。Python程序可以选择通过定义自定义反序列化器来限制此行为；然而，我们检查的ML库都没有这样做。即使它们这样做了，这些保护措施几乎总是可以被绕过；没有保证安全加载不受信任pickle文件的方法，正如官方Python 3.9 Pickle文档中的警告所强调的：

警告 pickle模块不安全。仅反序列化您信任的数据。

可以构造恶意pickle数据，在反序列化期间执行任意代码。切勿反序列化可能来自不受信任来源或可能被篡改的数据。

如果需要确保数据未被篡改，请考虑使用hmac对数据签名。

如果您正在处理不受信任的数据，更安全的序列化格式（如JSON）可能更合适。

我们不知道有任何ML文件格式包含模型的校验和†。

Python pickling的危险已被计算机安全社区所知相当长的时间。

介绍Fickling：pickle文件的反编译器、静态分析器和字节码重写器

Fickling有自己的Pickle虚拟机（PM）实现，并且可以安全地在潜在恶意文件上运行，因为它符号执行代码而不是明显执行它。

让我们看看如何使用Fickling逆向工程一个pickle文件，通过创建一个包含基本Python类型序列化列表的无害pickle：

1
2
3
4
5

$ python3 -c "import sys, pickle; \
  sys.stdout.buffer.write(pickle.dumps([1, ‘2’, {3: 4}]))" \
  > simple_list.pickle
$ python3 -m pickle simple_list.pickle
[1, ‘2’, {3: 4}]

在pickle文件上运行fickling将反编译它并产生人类可读的Python程序，等效于真实PM在反序列化期间运行的代码：

1
2

$ fickling simple_list.pickle
result = [1, ‘2’, {3: 4}]

在这种情况下，由于它是一个简单的序列化列表，代码既不令人惊讶也不非常有趣。通过将--trace选项传递给Fickling，我们可以跟踪PM的执行：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36

$ fickling --trace simple_list.pickle
PROTO
FRAME
EMPTY_LIST
    Pushed []
MEMOIZE
    Memoized 0 -> []
MARK
    Pushed MARK
BININT1
    Pushed 1
SHORT_BINUNICODE
    Pushed '2'
MEMOIZE
    Memoized 1 -> '2'
EMPTY_DICT
    Pushed {}
MEMOIZE
    Memoized 2 -> {}
BININT1
    Pushed 3
BININT1
    Pushed 4
SETITEM
    Popped 4
    Popped 3
    Popped {}
    Pushed {3: 4}
APPENDS
    Popped {3: 4}
    Popped '2'
    Popped 1
    Popped MARK
STOP
    result = [1, '2', {3: 4}]
    Popped [1, '2', {3: 4}]

您可以运行Fickling的静态分析来检测某些类别的恶意pickle，通过传递--check-safety选项：

1
2
3
4

$ fickling --check-safety simple_list.pickle
Warning: Fickling failed to detect any overtly unsafe code,
but the pickle file may still be unsafe.
Do not unpickle this file if it is from an untrusted source!

如果pickle文件是恶意的，它会是什么样子？好吧，为什么不制作一个！我们可以通过将任意Python代码注入pickle文件来实现：

1
2
3
4

$ fickling --inject 'print("Hello World!")' testpickle > testpickle.pwn3d
$ python3 -m pickle testpickle.pwn3d
Hello World!
[1, '2', {3: 4}]

它工作了！所以让我们看看Fickling的反编译：

1
2
3

$ fickling testpickle.pwn3d
_var0 = eval('print("Hello World!")')
result = [1, '2', {3: 4}]

及其分析：

1
2
3

$ fickling --check-safety testpickle.pwn3d
Call to `eval('print("Hello World!")')` is almost certainly
evidence of a malicious pickle file

Fickling也可以用作Python库，并具有编程接口来反编译、分析、修改和合成Pickle文件。它是开源的，您可以通过运行以下命令安装它：

1

pip3 install fickling

制作恶意ML模型

由于大多数ML模型广泛使用pickling，因此存在对模型进行权重/神经元扰动的潜在攻击面，包括故障注入、实时特洛伊木马和权重毒化攻击等。例如，在反序列化期间，注入pickle的代码可以根据本地环境（如时间、时区、主机名、系统区域设置/语言或IP地址）以编程方式对模型进行更改。这些更改可能是细微的，如位翻转攻击，或更明显的，如在反序列化中注入任意延迟以拒绝服务。

Fickling有一个基于官方PyTorch教程的概念验证，将任意代码注入现有的PyTorch模型。此示例显示如何将生成的模型加载到PyTorch中会自动列出当前目录中的所有文件（可能包含专有模型和代码）并将其外泄到远程服务器。

这对于像Microsoft的Azure ML这样的服务来说令人担忧，该服务支持在其云实例中运行用户提供的模型。恶意的“Fickled”模型可能导致拒绝服务，和/或在Microsoft可能假设为专有的环境中实现远程代码执行。如果多个用户的作业没有充分隔离，还存在外泄其他用户专有模型的潜在风险。

我们如何应对？

理想的解决方案是完全避免pickling。有几种不同的编码——JSON、CBOR、ProtoBuf——比pickling安全得多，并且足以编码这些模型。事实上，PyTorch已经包括state_dict和load_state_dict函数，这些函数将模型权重保存和加载到字典中，可以轻松序列化为JSON格式。为了完全加载模型，还需要模型结构（多少层、层类型等）。如果PyTorch实现模型结构的序列化/反序列化方法，整个模型可以更安全地编码到JSON文件中。

在PyTorch之外，还有其他框架避免使用pickle进行序列化。例如，开放神经网络交换（ONNX）旨在为编码AI模型提供通用标准以提高互操作性。ONNX规范使用ProtoBuf编码其模型表示。

负责任披露

我们在1月25日向PyTorch和PyTorch Hub维护者报告了关于共享ML模型的担忧，并在两天后收到了回复。维护者表示，他们将考虑向PyTorch和PyTorch Hub添加额外警告。他们还解释说，提交给PyTorch Hub的模型经过质量和效用审查，但维护者在将GitHub仓库链接添加到PyTorch Hub索引页面之前，不会对发布模型的人进行任何背景调查，也不会仔细审核代码的安全性。维护者似乎没有遵循我们转向更安全序列化形式的建议；他们说，责任在于用户确保第三方模型的来源和可信度。

我们认为这不够，特别是在面对日益普遍的typosquatting攻击时（参见pip和npm的攻击）。此外，供应链攻击可以很容易地将恶意代码注入合法模型，即使相关源代码看起来良性。检测此类攻击的唯一方法是使用像Fickling这样的工具手动检查模型。

结论

随着ML继续增长 popularity 和大多数从业者依赖通用框架，我们必须确保框架安全。许多用户没有计算机科学背景，更不用说计算机安全，可能不了解信任未知来源模型文件的危险。对于大多数框架来说，放弃pickling作为数据序列化形式相对简单，并且是安全的轻松胜利。我们期待有一天pickling不再用于反序列化不受信任的文件。同时，尝试Fickling并告诉我们您如何使用它！

致谢

非常感谢我们团队对此工作的辛勤工作：Sonya Schriner, Sina Pilehchiha, Jim Miller, Suha S. Hussain, Carson Harmon, Josselin Feist, and Trent Brunson

† 一些库如Tensorflow确实具有验证下载校验和的能力，但是，验证默认禁用，并且基于嵌入文件名中的校验和，这可以轻松伪造。