Featured image of post 使用eBPF保护PyTorch模型安全：实战指南与漏洞防御

使用eBPF保护PyTorch模型安全：实战指南与漏洞防御

本文深入探讨如何利用eBPF技术构建Python沙箱环境secimport，实现对PyTorch模型的运行时安全防护。通过分析pickle协议漏洞、供应链攻击风险，并演示如何监控系统调用级行为，有效阻断恶意代码执行。

使用eBPF保护PyTorch模型安全

由Hitesh Choudhary拍摄的Unsplash图片

本文并非由GPT生成

在本篇技术博客中，我将介绍secimport——一个基于eBPF（bpftrace）的Python沙箱工具包，用于创建和运行受保护的Python应用程序。我将从需求背景开始（可选择跳过），然后演示如何安全运行PyTorch模型。

在本系列的第一部分中，我介绍了Python的操作系统和应用级追踪与沙箱技术，并讨论了使用dtrace实现的最小可行方案（MVP），该方案可在系统调用级别全面保护Python运行时环境。关于现有沙箱解决方案的深入解析，请查阅该文章！

目录

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

评估不安全代码

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

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

我们依赖星标作为可信度指标，这种观念需要改变。我们给存储库加星标是为了收藏——很少有人真正打算贡献代码或深入审查。许多人会不加思索地直接使用这些代码，这意味着我们在未审查的情况下运行了他人编写的代码。

我认为重大安全事件的发生只是时间问题。对于LLM（大语言模型）而言，几天内获得2K星标并不罕见，但跟踪和审查这些模型的代码几乎不可能。此外，如今伪造星标也很容易——只需为项目准备一份优秀的README（借助GPT），就会有人尝试运行你的代码。

安全措施应在Python运行时层面实施。一个典型的例子是pickle协议。许多主流框架（以及Python的多进程模块）都依赖pickle作为基础组件。

为什么pickle存在安全问题？
其设计本身存在漏洞：

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

import pickle

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

# 序列化恶意对象
pickle.dumps(Demo())
# 输出: 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数据
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进程。

由Jase Bloor拍摄的Unsplash图片

以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

# -*- 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
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命令实际运行之前记录了两个违规：

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用户能够以不同的权限和规则限制代码

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