使用eBPF保护PyTorch模型安全

本文介绍如何使用secimport工具包结合eBPF技术为PyTorch模型创建安全沙箱环境，防止恶意代码执行，详细分析pickle协议的安全风险并演示实时系统调用监控和防护方案。

使用eBPF保护PyTorch模型安全

引言

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

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

Photo by Hitesh Choudhary on Unsplash

在本系列的第1部分中，我介绍了Python的操作系统和应用程序追踪与沙箱化。我写了一个使用dtrace的最小可行解决方案（MVP），该方案一直保护到系统调用级别。

关于现有沙箱解决方案的深入解释——请查看！

目录

  • 评估不安全代码
  • 关于pickle协议和供应链攻击的说明
  • PyTorch沙箱示例
  • 使用secimport阻止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进程。

Photo by Jase Bloor on Unsplash

让我们以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模型，也应该始终预先知道。

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
63

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

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

归根结底，对于内核来说，一切都是系统调用。

Photo by Thomas Park on 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）

在其他

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