使用eBPF保护PyTorch模型安全

引言

在这篇博客中，我将介绍secimport——一个利用eBPF（bpftrace）保护Python运行时的工具包，用于创建和运行沙箱化应用程序。

我将从为什么需要它开始（可以跳过这部分），然后演示如何安全地运行PyTorch模型。

评估不安全代码

在当今的软件开发环境中，向代码库添加新库可能具有挑战性。我们缺乏对包期望的清晰认识（它应该做什么才能正常运行），导入的包可以在我们不知情的情况下操纵我们的环境。

如果你查看HuggingFace，存储库通常存储模型的PyTorch定义，这是Python代码。你可能认为有人已经审查过它…这不算安全吗？它有足够的星星…

我们依赖星星作为可信度指标，这应该改变。

我们给存储库加星是为了收藏它们——我们中很少有人真正打算贡献代码并深入研究它。我们中的许多人毫不犹豫地使用它。我们在没有审查的情况下使用别人的代码。

我认为重大安全事件的发生只是时间问题。

关于pickle协议和供应链攻击

一个例子是pickle协议。许多主要框架（和Python的多进程处理）都依赖pickle作为构建块。

为什么pickle是个问题？

它的设计就是脆弱的，这就是原因：

1
2
3
4
5
6
7
8


import pickle

class Demo:
    def __reduce__(self):
        return (eval, ("__import__('os').system('echo Exploited!')",))

In: pickle.dumps(Demo())
Out: b"\x80\x04\x95F\x00\x00\x00\x00\x00\x00\x00\x8c\x08builtins\x94\x8c\x04eval\x94\x93\x94\x8c*__import__('os').system('echo Exploited!')\x94\x85\x94R\x94."

在另一个终端或环境中，加载别人的pickle代码将导致——你猜对了——一个漏洞：

1
2
3
4


import pickle
pickle.loads(b"\x80\x04\x95F\x00\x00\x00\x00\x00\x00\x00\x8c\x08builtins\x94\x8c\x04eval\x94\x93\x94\x8c*__import__('os').system('echo Exploited!')\x94\x85\x94R\x94.")
Exploited!
0

让我们看看secimport如何使用eBPF阻止这个pickle漏洞：

常规Python（左侧）vs Secimport（右侧）。

在上图中，secimport能够阻止pickle漏洞，因为我们预先定义了策略。

我们使用"secimport run"运行Python进程——它在eBPF监督下实时运行Python进程。

PyTorch沙箱示例

假设你现在已经理解在Python中我们可以相当容易地运行任意代码，让我们尝试保护一个给定的PyTorch模型。

我们将运行PyTorch文档中的一个示例，有沙箱和没有沙箱的情况
我们将向代码添加恶意行
沙箱将记录违规行为，然后在活动发生之前阻止它（IDS模式vs IPS模式）

我们将使用torch（==2.0.1）的随机示例：

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


# -*- coding: utf-8 -*-
import time
import torch
import math

class Polynomial3(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.a = torch.nn.Parameter(torch.randn(()))
        self.b = torch.nn.Parameter(torch.randn(()))
        self.c = torch.nn.Parameter(torch.randn(()))
        self.d = torch.nn.Parameter(torch.randn(()))

    def forward(self, x):
        # import os; os.system('ps')
        return self.a + self.b * x + self.c * x ** 2 + self.d * x ** 3

    def string(self):
        return f'y = {self.a.item()} + {self.b.item()} x + {self.c.item()} x^2 + {self.d.item()} x^3'

# ... 其余代码

运行代码：

1
2


root@3ecd9c9b5613:/workspace# Python-3.10.0/python -m pip install torch
root@3ecd9c9b5613:/workspace# Python-3.10.0/python pytorch_example.py

代码运行良好。

使用secimport阻止代码执行

现在，让我们取消注释"os.system"命令，看看secimport是否能识别该变化。“os.system"也可以被混淆或使用’subprocess’模块来运行命令——但由于我们在系统调用级别监控，我们不在乎！我们的eBPF沙箱应该能看到一切。

我们将使用与上一步相同的沙箱。这次，程序将在模型的forward()中执行"ps"命令。

让我们改变这个：

1
2


def forward(self, x):
    return self.a + self.b * x + self.c * x ** 2 + self.d * x ** 3

变成：

1
2
3


def forward(self, x):
    import os; os.system('ps')
    return self.a + self.b * x + self.c * x ** 2 + self.d * x ** 3

Secimport能够检测到包被导入。但它不会在机器上做任何事情。之后，secimport将捕获"os.system"调用，导致系统调用。

在Mac上，将使用系统调用56和61（CLONE和WAIT4）。在Linux上，将单独使用系统调用59（EXECVE）。

让我们在注入任意命令后再次运行代码。

我们应该期望沙箱记录违规行为。

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


root@3ecd9c9b5613:/workspace# secimport run --entrypoint pytorch_example.py
>>> secimport run
...
[SECURITY PROFILE VIOLATED]: /workspace/examples/cli/ebpf/torch_demo/pytorch_example.py called syscall 56 at depth 561032
[SECURITY PROFILE VIOLATED]: /workspace/examples/cli/ebpf/torch_demo/pytorch_example.py called syscall 61 at depth 561032
  PID TTY          TIME CMD
    1 pts/0    00:00:00 sh
   11 pts/0    00:00:00 bash
 4279 pts/0    00:00:00 python
 4280 pts/0    00:00:00 sh
 4281 pts/0    00:00:06 bpftrace
 4285 pts/0    00:00:06 python
 8289 pts/0    00:00:00 sh
 8290 pts/0    00:00:00 ps
1999 9.100260734558105
Result: y = 0.017583630979061127 + 0.8593376278877258 x + -0.003033468732610345 x^2 + -0.09369975328445435 x^3

太棒了！secimport在ps命令实际运行之前记录了2个违规：

1
2


[SECURITY PROFILE VIOLATED]: /workspace/examples/cli/ebpf/torch_demo/pytorch_example.py called syscall 56 at depth 561032
[SECURITY PROFILE VIOLATED]: /workspace/examples/cli/ebpf/torch_demo/pytorch_example.py called syscall 61 at depth 561032

系统调用编号56（clone）系统调用编号61（sys_wait4）

不仅仅是检测——如何防止代码执行？

在上面的例子中，secimport只记录了策略违规。

使用这两个标志可以轻松防止代码执行：

1
2
3
4
5
6
7
8


secimport run … --stop_on_violation
[SECURITY PROFILE VIOLATED]: <stdin> called syscall 56 at depth 8022
^^^ STOPPING PROCESS 85918 DUE TO SYSCALL VIOLATION ^^^

secimport run … --kill_on_violation
[SECURITY PROFILE VIOLATED]: <stdin> called syscall 56 at depth 8022
^^^ KILLING PROCESS 86466 DUE TO SYSCALL VIOLATION ^^^
KILLED.

通过简单地添加其中一个标志，你可以在生产运行时真正阻止代码，在它发生之前。这是人们可以期望的最强保护类型。

结论

在这篇博客文章中，我们介绍了secimport的用法，以及如何通过secimport的CLI一路保护torch模型运行时到内核的系统调用，按模块进行。

Secimport使Python用户能够在其代码中限制不同模块具有不同的特权和规则。

我鼓励你为你的用例尝试secimport。