使用Semgrep保护机器学习安全

tl;dr: 我们公开的Semgrep规则集现已包含11条专门检测机器学习库误用的规则。立即尝试！

想象一下：您花了数月时间整理图像、尝试不同架构、下载预训练模型、折腾Kubernetes，终于准备好交付闪亮的机器学习（ML）产品。然后您被问到（希望不是令人恐惧的问题）：您采取了哪些安全措施？

或许您已使用Counterfit和PrivacyRaven等工具测试模型以防止模型提取和模型反转，但这不应是终点。您不仅在构建模型，还在构建流水线。而流水线的核心是源代码。ML模型不能被视为独立对象，其创建者必须考虑流水线的每个元素，从预处理程序到底层硬件。

Semgrep是一款强大的静态分析工具，使用户能够大规模搜索潜在有问题的代码模式。此前，我们发布了多条检测goroutine泄漏的Semgrep规则，这些规则已包含在我们的公开规则集中。为加强ML生态系统，我们分析了多个ML库的源代码，识别出一些常见的问题模式。我们决定将这些模式转化为Semgrep规则，以便轻松查找和修复潜在漏洞，而无需依赖开发人员彻底检查每个所用库的源代码或文档。

千个pickle之地中的野生分布式pickle

不幸的是，ML生态系统是“千个pickle之地”。Pickling（一种特别流行的保存ML模型的做法）出现在大量工具和库中。然而，pickling极其不安全，极易导致任意代码执行。为解决此问题，我们去年发布了Fickling，一款用于逆向工程和创建恶意pickle文件的工具。但许多库仍在底层使用pickling，即使安全意识强的开发人员也可能最终使用它们。让我们看一个著名ML库中pickling的微妙实例：PyTorch Distributed。

深度学习的高计算需求推动开发人员使用分布式系统，其中计算可在多台机器和设备上并行化。在这些系统中，多个进程形成一个组，通信在该组内进行。有些组要求开发人员利用点对点通信在进程间传输数据，而其他组则需要集体通信。许多库专注于分布式ML，每个库都支持对象从一個进程广播到组内其他进程。

PyTorch Distributed向用户提供broadcast_object_list函数，该函数将对象列表共享到组中的其他节点。其用法如下列从PyTorch Distributed文档获取的代码片段所示：

1
2
3

# 示例代码：广播对象列表
objects = [torch.tensor([1, 2, 3]), "hello"]
torch.distributed.broadcast_object_list(objects, src=0)

为共享上述对象列表，列表首先被pickle。这使得攻击者可以制作恶意pickle文件，从而在整个系统中执行任意代码。

文档还显示了关于此函数中隐式使用pickle的警告，如下所示：

Warning: This function uses pickle module for serialization, which is insecure. Only use this function if you trust the participants.

尽管有警告，开发人员可能不会立即意识到这是pickling的不安全应用。我们在此包中发现多个具有相同问题的函数，因此创建了以下Semgrep规则帮助用户避免野生pickle的调用。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

rules:
- id: pickles-in-torch-distributed
  patterns:
    - pattern-either:
      - pattern: torch.distributed.broadcast_object_list(...)
      - pattern: torch.distributed.all_gather_object(...)
      - pattern: torch.distributed.gather_object(...)
      - pattern: torch.distributed.scatter_object_list(...)
  message: |
    Functions reliant on pickle can result in arbitrary code execution. 
    For more information, see https://blog.trailofbits.com/2021/03/15/never-a-dill-moment-exploiting-machine-learning-pickle-files/ 
  languages: [python]
  severity: WARNING

大pickle再次出击

NumPy是ML开发人员使用的另一个重要库，它将外部数据转换为NumPy数组并执行不同操作。Pickle也出现在NumPy中，即在numpy.load()中，如下所示：

1
2

import numpy as np
data = np.load('file.npy', allow_pickle=True)

在NumPy 1.16.0及更早版本中，np.load()默认允许pickle，从而导致任意代码执行。在CVE-2019-6446之后，新版本NumPy默认将allow_pickle设置为False。您可以在GitHub上查看此漏洞的报告和NumPy维护者的响应（见上图）。

如下所示，代码将allow_pickle标志设置为True，并更新了针对新版本NumPy的早期PoC：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12

class Malicious:
    def __reduce__(self):
        import os
        return (os.system, ('ls',))

malicious_obj = Malicious()
import pickle
with open('malicious.pkl', 'wb') as f:
    pickle.dump(malicious_obj, f)

# 触发任意代码执行
np.load('malicious.pkl', allow_pickle=True)

更具体地说，我们创建了一个类，将其pickle，然后在NumPy中加载。由于类的__reduce__方法包含pwd命令，NumPy在尝试加载文件时将运行ls，从而演示任意代码执行。类似的proof of concept可在fickling仓库中找到。为查找此类错误模式，我们开发了以下Semgrep规则：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14

rules:
  - id: pickles-in-numpy
    patterns:
      - pattern: numpy.load(..., allow_pickle=$VALUE)
      - metavariable-regex:
          metavariable: $VALUE
          regex: (True|^\d*[1-9]\d*$)
    message: | 
      Functions reliant on pickle can result in arbitrary code execution. 
      Consider using fickling or switching to a safer serialization method. 
      For more information, see https://blog.trailofbits.com/2021/03/15/never-a-dill-moment-exploiting-machine-learning-pickle-files/
    languages:
      - python
    severity: ERROR

通常，开发人员仅在想要使用参数时将其设置为True。然而，当您将其设置为任何正数时，此参数也将有效。因此，为使我们的规则对变体具有鲁棒性，我们使用Semgrep的metavariable-regex功能编写正则表达式，搜索启用pickling的不同值。否则，此规则可能导致对上述情况的假阴性或对allow_pickle=False的假阳性。

不那么随机的随机性

您能发现以下代码中的错误吗？

1
2
3
4
5
6
7
8

from torch.utils.data import Dataset
import numpy as np

class MyDataset(Dataset):
    def __getitem__(self, index):
        # 错误：使用NumPy RNG
        augmentation = np.random.randint(0, 10)
        return augmentation

由于错误的随机数生成，此代码将导致相同的增强！去年，Tanel Pärnamaa写了一篇博客文章，讨论使用PyTorch和NumPy时出现的微妙错误。更具体地说，使用多个工作进程并行加载数据，每个工作进程继承NumPy随机数生成器的相同初始状态。据Pärnamaa称，此错误已在PyTorch新版本中修复，但在Keras中仍然存在。我鼓励您探索其他框架并尝试复现此问题。

显然，ML中许多问题可能源于并行和分布式计算与随机性的差异。Trail of Bits对随机数生成问题并不陌生。它们众所周知是密码学实现的幽灵。为攻击ML模型，Ilia Shumailov等人开发了数据排序攻击，其中攻击者控制数据提供给模型的顺序，以防止模型学习或注入后门。此外，多项研究检查了ML工具中的非确定性和数值不稳定性，我们必须从这些工具中移除这些问题，并在我们自己的ML代码中避免。

我们编写了以下规则以避免此错误：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24

rules:
- id: numpy-in-torch-datasets
  patterns:
    - pattern-either:
      - pattern: |
          class $X(Dataset):
            ...
            def __getitem__(...):
              ...
              np.random.randint(...)
              ...
      - pattern: |
          class $X(Dataset):
            ...
            def __getitem__(...):
              ...
              $Y = np.random.randint(...)
              ...
  message: |

    Instead, use the random number generators built into Python and PyTorch. 
    See https://tanelp.github.io/posts/a-bug-that-plagues-thousands-of-open-source-ml-projects/ for more details       
  languages: [python]
  severity: WARNING

此错误并非避免在Torch数据集和模型中使用NumPy操作的唯一原因。开放神经网络交换（ONNX）是一种基于Protobufs存储ML模型的文件格式。然而，在模型内部使用NumPy操作可能导致不正确的跟踪和格式错误的ONNX文件。此外，PyTorch提供FX，其中包括PyTorch模型的符号跟踪器；在模型内部使用NumPy和其他非Torch函数可能禁止符号跟踪。这也可能限制ML模型的效率，并阻止您使用TorchScript。因此，我们创建了一条规则来捕获Torch模型中NumPy操作的实例。

发酵安全未来

开发Semgrep规则以查找和修复漏洞的努力正在进行中。随着这些ML库以及整个社区的持续成熟，我们希望像Semgrep这样的工具可以帮助防止误用。

鉴于机器学习是Trail of Bits的重要优先事项，我们完全期望在ML库中发现更多错误，并发布工具使其他人能够发现和纠正它们。但我们需要您的帮助。如果您有任何关于使用Semgrep等工具可以预防的其他潜在问题的建议，请告知我们。请在我们的仓库中提出问题或通过suha.hussain@trailofbits.com与我联系。

附录1：规则

以下是我们迄今为止添加到仓库中的ML库规则：

| waiting-with-torch-distributed | Python | Possible undefined behavior when not waiting for requests |