使用eBPF为PyTorch模型构建安全沙箱

本文介绍了如何利用eBPF技术，通过secimport工具包为PyTorch等Python运行时创建安全沙箱环境，有效防范供应链攻击和恶意代码执行，提供了从策略创建到行为监控的完整安全解决方案。

使用eBPF保护PyTorch模型

本文并非由GPT生成

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

我将从阐述其必要性开始（如果您已了解，可以直接跳过这部分），然后演示如何安全地运行PyTorch模型。

（图片说明：由Hitesh Choudhary在Unsplash上发布的照片）

在本系列的第一部分，我介绍了针对Python的操作系统和应用追踪与沙箱化。我写了一个基于dtrace的最小可行解决方案（MVP），它能将Python运行时的安全防护一直延伸到系统调用级别。

若想深入了解现有的沙箱解决方案——请查看该文！

目录

  • 评估不安全的代码
  • 关于Pickle协议和供应链攻击的说明
  • PyTorch沙箱示例
  • 使用secimport阻止PyTorch代码执行
  • 结论

评估不安全的代码

在当今的软件开发环境中，向我们的代码库添加新的库可能充满挑战。我们缺乏对软件包预期行为的清晰认知（它需要做什么才能正常工作），并且导入的包可能在不知不觉中操纵我们的环境。

如果你查看HuggingFace，会发现仓库通常存储着模型的PyTorch定义，也就是Python代码。你可能觉得……已经有人审查过它了，难道这还不够安全吗？它有足够多的星标……

我们依赖“星标”作为可信度指标，这一点需要改变。

我们给仓库加星标是为了收藏它们——很少有人真的打算贡献代码并深入研究代码。许多人会不加思索地使用它。我们在没有审查的情况下使用别人的代码。

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

对于自LLM兴起以来，短短几天内获得2K星标并不罕见，而跟踪或审查其代码几乎是不可能的。

安全措施应该在Python的运行时中实施。

一个例子就是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
5

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进程。

（图片说明：由Jase Bloor在Unsplash上发布的照片）

让我们以Pytorch为例。这是PyTorch官方包文档中的一条信息，我相信如果你使用过PyTorch，很多人可能都错过了：

（图片说明：PyTorch文档警告截图）

PyTorch模型非常容易被利用。例如，这篇博客展示了如何用基于pickle的漏洞利用程序来修补任何torch模型，这与上面的例子非常相似，但更复杂，并且能通过所有代码静态安全扫描。

我们应该努力避免执行非预期的代码。

你可以使用CLI来追踪你允许的pickle文件的逻辑， 或者使用secimport的Python API来确保在加载特定文件时，“pickle”不会运行奇怪的系统调用：

1
2
3
4
5

import secimport

pickle = secimport.secure_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.")
[1]    28027 killed     ipython

运行此代码后，会自动创建一个日志文件，其中包含有关该进程以及它为何被终止的所有信息：

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

$ less /tmp/.secimport/sandbox_pickle.log
  @posix_spawn from /Users/avilumelsky/Downloads/Python-3.10.0/Lib/threading.py
    DETECTED SHELL:
        depth=8
        sandboxed_depth=0
        sandboxed_module=/Users/avilumelsky/Downloads/Python-3.10.0/Lib/pickle.py
    TERMINATING SHELL:
        libsystem_kernel.dylib`__posix_spawn+0xa
        ...
                libsystem_kernel.dylib`__posix_spawn+0xa
                libsystem_c.dylib`system+0x18b
                python.exe`os_system+0xb3
    KILLED:

我希望你理解了这个问题，许多工业项目都依赖这种不安全的格式。无论如何——程序的行为，即使是AI模型的行为，也应该总是预先可知的。

AI中的供应链攻击

有很多方法可以用最小的努力来利用Python用户。 PyTorch存在许多设计问题。 图像文件 拼写错误 HuggingFace、AutoGPT和类似的SaaS AI公司都以此方式工作（依赖于设计上就不安全的pickle和框架）。 模型通常是开源的，我们需要它们是可移植的。 PyTorch nn.Module 实例（模型）通过代码进行管理。 使用别人的模型会导致在你的私有环境中加载不安全的代码。

PyTorch沙箱示例

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

  1. 我们将运行一个来自PyTorch文档的示例，包含沙箱和不包含沙箱两种情况。
  2. 我们将在代码中添加一行恶意代码。
  3. 沙箱将记录违规行为，然后在活动发生之前将其阻止（IDS模式与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
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62

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

class Polynomial3(torch.nn.Module):
    def __init__(self):
        """
        In the constructor we instantiate four parameters and assign them as
        member parameters.
        """
        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):
        """
        In the forward function we accept a Tensor of input data and we must return
        a Tensor of output data. We can use Modules defined in the constructor as
        well as arbitrary operators on Tensors.
        """
#        import os; os.system('ps')
        return self.a + self.b * x + self.c * x ** 2 + self.d * x ** 3

    def string(self):
        """
        Just like any class in Python, you can also define custom method on PyTorch modules
        """
        return f'y = {self.a.item()} + {self.b.item()} x + {self.c.item()} x^2 + {self.d.item()} x^3'

start_time = time.time()
# Create Tensors to hold input and outputs.
x = torch.linspace(-math.pi, math.pi, 2000)
y = torch.sin(x)

# Construct our model by instantiating the class defined above
model = Polynomial3()

# Construct our loss function and an Optimizer. The call to model.parameters()
# in the SGD constructor will contain the learnable parameters (defined
# with torch.nn.Parameter) which are members of the model.
criterion = torch.nn.MSELoss(reduction='sum')
optimizer = torch.optim.SGD(model.parameters(), lr=1e-6)

for t in range(2000):
    # Forward pass: Compute predicted y by passing x to the model
    y_pred = model(x)

    # Compute and print loss
    loss = criterion(y_pred, y)
    if t % 100 == 99:
        print(t, loss.item())

    # Zero gradients, perform a backward pass, and update the weights.
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

print(f'Result: {model.string()}')
print("--- %s seconds ---" % (time.time() - start_time))

运行代码：

1
2
3
4
5
6
7

root@3ecd9c9b5613:/workspace# Python-3.10.0/python -m pip install torch
root@3ecd9c9b5613:/workspace# Python-3.10.0/python pytorch_example.py
99 674.6323852539062
...
1999 9.19102668762207
Result: y = -0.013432367704808712 + 0.8425596952438354 x + 0.0023173068184405565 x^2 + -0.09131323546171188 x^3
--- 0.6940326690673828 seconds ---

代码运行得很好。

为该代码创建量身定制的沙箱

现在，我想为那段代码创建一个安全策略，使得除了该代码本身之外，任何其他代码都不能运行。这是通过 secimport trace 命令来追踪代码，以及 secimport build 命令来从追踪记录构建沙箱完成的。

您可以通过此处或此处阅读它的工作原理，但请继续往下看。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

root@ec15bafca930:/workspace/examples/cli/ebpf/torch_demo# secimport trace --entrypoint pytorch_example.py
  >>> secimport trace
  TRACING: ['/root/.local/lib/python3.10/site-packages/secimport/profiles/trace.bt', '-c', 'bash -c "/workspace/Python-3.10.0/python pytorch_example.py"', '-o', 'trace.log']
                          Press CTRL+D/CTRL+C to stop the trace;
/workspace/examples/cli/ebpf/torch_demo/pytorch_example.py:36: UserWarning: Failed to initialize NumPy: No module named 'numpy' (Triggered internally at ../torch/csrc/utils/tensor_numpy.cpp:84.)
  x = torch.linspace(-math.pi, math.pi, 2000)
......
Result: y = -0.04786265641450882 + 0.8422093987464905 x + 0.008257105946540833 x^2 + -0.09126341342926025 x^3
--- 1.5200915336608887 seconds ---
 TRACING DONE;

让我们从追踪的代码构建一个沙箱。

它将根据追踪期间代码中运行的模块创建系统调用映射。在代码新位置出现的任何新系统调用，或添加逻辑的代码更改，都将导致secimport视为“违规”行为。

在以下示例中，secimport通过分析追踪记录，为你的代码构建YAML/JSON策略。

 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

root@ec15bafca930:/workspace/examples/cli/ebpf/torch_demo# secimport build
  >>> secimport build
  SECIMPORT COMPILING...
 CREATED JSON TEMPLATE:  sandbox.json
CREATED YAML TEMPLATE:  sandbox.yaml
compiling template sandbox.yaml
[debug] adding syscall write to blocklist for module general_requirements
[debug] adding syscall writev to blocklist for module general_requirements
...
[debug] adding syscall stat to allowlist for module /workspace/Python-3.10.0/lib/python3.10/site-packages/py.py
[debug] adding syscall clock_gettime to allowlist for module /workspace/Python-3.10.0/lib/python3.10/site-packages/py.py
[debug] adding syscall exit_group to allowlist for module /workspace/Python-3.10.0/lib/python3.10/site-packages/py.py
...
[debug] adding syscall mmap to allowlist for module /workspace/Python-3.10.0/lib/python3.10/site-packages/torch/ao/
[debug] adding syscall brk to allowlist for module /workspace/Python-3.10.0/lib/python3.10/site-packages/torch/aut
[debug] adding syscall close to allowlist for module /workspace/Python-3.10.0/lib/python3.10/site-packages/torch/aut
...
[debug] adding syscall mmap to allowlist for module /workspace/Python-3.10.0/lib/python3.10/site-packages/torch/jit
[debug] adding syscall brk to allowlist for module /workspace/Python-3.10.0/lib/python3.10/site-packages/torch/lib
[debug] adding syscall clock_gettime to allowlist for module /workspace/Python-3.10.0/lib/python3.10/site-packages/torch/nn/
[debug] adding syscall write to allowlist for module /workspace/Python-3.10.0/lib/python3.10/site-packages/torch/nn/
...
DTRACE SANDBOX:  sandbox.d
BPFTRCE SANDBOX:  sandbox.bt
 SANDBOX READY: sandbox.bt

现在，让我们在沙箱中运行原始代码：

1
2
3
4
5
6
7

root@3ecd9c9b5613:/workspace# secimport run --entrypoint pytorch_example.py
99 3723.3251953125
...
1999 11.318828582763672
Result: y = -0.04061822220683098 + 0.8255564570426941 x + 0.007007318548858166 x^2 + -0.08889468014240265 x^3
--- 0.8806719779968262 seconds ---
SANDBOX EXITED;

很好！正如预期的那样，它在沙箱内运行，没有任何错误。

现在，我们都想看看沙箱的实际运作——让我们修改代码，让它做一些新的事情。

使用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）。 其他Python多进程库将调用其他系统调用（Fork/Spawn）——但在相同操作系统和相同解释器上，行为将始终一致。 归根结底，对于内核来说，一切都是系统调用。

（图片说明：由Thomas Park在Unsplash上发布的照片）

在我们注入了任意命令后，让我们再次运行代码。 我们应该期望沙箱记录违规行为。

 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)

在其他机器上，Python可能单独使用系统调用编号 59 (EXECVE) 而不是 56 然后 61。被调用的系统调用取决于操作系统，可能会有所不同，但在相同的操作系统和解释器上，secimport始终保持一致。

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

在上述案例中，secimport只记录了策略违规。 使用这两个标志可以轻松实现代码执行预防：

secimport run … — stop_on_violation

1
2

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

secimport run … — kill_on_violation

1
2
3

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

只需添加其中一个标志，你就可以在生产运行时真正阻止代码执行，防患于未然。这是人们可以期望的最强保护类型。

你未必会在每个项目中使用这些标志，默认情况下你可能只想记录日志，但这种“终止”和“停止”行为可以解决很多问题，或者让安全团队有能力限制Python运行时中的第三方代码。

结论

在这篇博客文章中，我们了解了secimport的用法，以及如何通过secimport CLI，将torch模型运行时的安全防护一直延伸到内核的系统调用级别，且针对每个模块进行。

Secimport使Python用户能够在其代码中，为不同的模块设置不同的权限和规则。

我鼓励您为自己的用例尝试secimport。

感谢您阅读到这里。

如果您喜欢我的工作，请告诉我。

如果您有任何问题或想法，欢迎与我联系或评论，当然，也欢迎您在GitHub上给它加星标并贡献代码！

参考资料：

您可以在此处阅读更多关于eBPF和bpftrace的信息。

顺便说一下，我是在业余时间做这些的。我确实也非常喜欢咖啡！

（作者签名及推广链接、其他文章链接等内容，根据格式要求省略不翻译）

comments powered by Disqus
© 2020 - 2026 qife
使用 Hugo 构建
主题 Stack 由 Jimmy 设计
本站总访问量次 本站访客数人次