使用Semgrep保护机器学习安全

tl;dr: 我们公开可用的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
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.

尽管有警告，开发人员并不立即清楚这是pickle的不安全应用。我们在此包中发现了多个具有相同问题的函数，因此我们创建了以下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
13
14

import numpy as np
import pickle

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

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

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

更具体地说，我们创建了一个类，将其pickle，然后在NumPy中加载它。因为类的__reduce__方法包含pwd命令，NumPy在尝试加载文件时会运行ls，从而演示任意代码执行。可以在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):
        # 有问题的代码
        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库规则：