永不安全的时刻：利用机器学习Pickle文件 - Trail of Bits博客

机器学习模型中的Pickle安全隐患

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

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

本文将探讨加载不受信任的pickle文件或ML模型时可能发生的隐秘恶意行为。在此过程中，我们介绍一个新工具Fickling，它可以帮助您逆向工程、测试甚至创建恶意pickle文件。如果您是ML从业者，您将了解标准ML实践中固有的安全风险。如果您是安全工程师，您将了解一个可以帮助构建和取证检查pickle文件的新工具。无论哪种情况，希望本文能让您对pickle产生警惕。

Pickle的存储机制：令人震惊！

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

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

从这个描述来看，人们可能合理推断PM不是图灵完备的。这种格式怎么可能不安全？用三岛由纪夫著名格言的变体来说：

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

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文件格式包含模型的校验和†。

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

Python pickle的危险性早已为计算机安全社区所知。

介绍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}]

在这种情况下，由于是简单的序列化列表，代码既不令人惊讶也不很有趣。通过向Fickling传递–trace选项，我们可以跟踪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}]

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

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

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

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

如何应对？

理想的解决方案是完全避免使用pickle。有几种不同的编码方式——JSON、CBOR、ProtoBuf——比pickle安全得多，并且足以编码这些模型。事实上，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索引页面之前，不会对发布模型的人进行任何背景调查，也不会仔细审核代码安全性。维护者似乎没有遵循我们转向更安全序列化形式的建议；他们说责任在于用户确保第三方模型的来源和可信度。

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

结论

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

致谢

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

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