使用Semgrep保护机器学习代码安全

概述

我们的公开Semgrep规则集现在包含11条专门针对机器学习库误用的规则。立即尝试吧！

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

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

Semgrep是一个强大的静态分析工具，使用户能够大规模搜索潜在有问题的代码模式。之前，我们发布了几条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

# 示例代码

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

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

1

# 警告示例

尽管有警告，开发人员并不立即清楚这是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: |
   依赖pickle的函数可能导致任意代码执行。
   更多信息，请参见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

# 示例代码

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

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

1

# 示例代码

更具体地说，我们创建了一个类，将其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: | 
      依赖pickle的函数可能导致任意代码执行。
      考虑使用fickling或切换到更安全的序列化方法。
      更多信息，请参见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

# 示例代码

由于错误的随机数生成，此代码将导致相同的增强！去年，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: |
    在Torch数据集中使用NumPy RNG可能导致许多数据加载问题，包括相同的增强。
    相反，使用Python和PyTorch内置的随机数生成器。
    更多细节，请参见https://tanelp.github.io/posts/a-bug-that-plagues-thousands-of-open-source-ml-projects/
  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库规则：