使用eBPF保护PyTorch模型安全的技术实践

本文详细介绍如何使用eBPF技术通过secimport工具包为PyTorch模型创建安全沙箱环境，防止恶意代码执行，涵盖从代码追踪到策略构建的完整实现流程。

使用eBPF保护PyTorch模型

评估不安全代码

在现代软件开发中，向代码库添加新库可能具有挑战性。我们缺乏对包期望的清晰认识，导入的包可以在我们不知情的情况下操纵环境。

以HuggingFace为例，存储库通常存储模型的PyTorch定义，即Python代码。你可能会认为有人已经审查过它…这难道不安全吗？它有足够的星星…

我们依赖星星作为可信度指标，这应该改变。我们给存储库加星是为了收藏它们——很少有人真正打算贡献并深入研究代码。许多人会不假思索地使用它。我们在没有审查的情况下使用别人的代码。

我认为重大安全事件的发生只是时间问题。对于LLMs来说，几天内获得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

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（左侧）与Secimport（右侧）对比。

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

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

PyTorch示例

让我们以Pytorch为例。这是PyTorch包文档中的官方消息，我相信如果你曾经使用过Pytorch，很多人都会错过：

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

我们应该努力避免执行意外的代码。

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

1
2
3
4

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中的供应链攻击

有许多方法可以以最小的努力利用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

# -*- 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'

start_time = time.time()
x = torch.linspace(-math.pi, math.pi, 2000)
y = torch.sin(x)
model = Polynomial3()
criterion = torch.nn.MSELoss(reduction='sum')
optimizer = torch.optim.SGD(model.parameters(), lr=1e-6)

for t in range(2000):
    y_pred = model(x)
    loss = criterion(y_pred, y)
    if t % 100 == 99:
        print(t, loss.item())
    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

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']
...
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）——但在相同的解释器和相同的操作系统上，它总是相同的行为。

最终，对于内核来说，一切都是系统调用。

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

 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只记录了策略违规。

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

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.

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

你不会在每个项目中使用这些标志，你可能默认想要记录，但这种Kill和Stop行为可以解决许多问题，或者给安全团队在Python运行时中限制第三方代码的能力。

结论

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

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

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

感谢你阅读到这里。

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

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

参考文献

第1部分：在你的代码中沙箱化Python模块 第3部分：使用eBPF保护FastAPI 源代码和示例：https://github.com/avilum/secimport

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