jxn17/vulnerability_benchmark_ops

# 高性能机器学习系统工程与漏洞基准测试 ## 目录 1. [概述](#overview) 2. [项目架构](#project-architecture) 3. [项目结构](#project-structure) 4. [Google Colab 故障排除与修复](#google-colab-troubleshooting--fixes) 5. [环境设置](#environment-setup) 6. [构建与运行](#build--run) 7. [预期终端输出](#expected-terminal-output) 8. [AI 协作日志](#ai-collaboration-log) 9. [执行输出（PDF）](#execution-output-pdf) ## 概述 本项目作为机器学习系统的安全测试环境。当工程师编写自定义 C++ 操作以使 ML 模型在 GPU 上运行更快时，有时会引入内存管理缺陷。 为了展示这些缺陷如何发生以及如何检测它们，我们故意在自定义 PyTorch C++ 扩展中注入了一个经典漏洞——堆栈缓冲区溢出（CWE-120）。然后，我们使用 AddressSanitizer（ASAN）构建了一个自动化管道。这个管道证明我们的安全工具可以在将过大的“坏数据”输入系统时成功捕获该错误，同时验证在正常、安全数据的情况下，底层数学运算仍然可以完美执行。 ## 项目架构 本项目分为四个独立的层： ### 1. 核心操作 * `modules/custom_operator.cpp`：包含将我们的自定义矩阵乘法代码连接到 PyTorch 的核心逻辑。关键的是，此文件还包含故意注入的 CWE-120 堆栈缓冲区溢出漏洞。 * `modules/kernel.cu`：一个高度优化的 CUDA 脚本，允许数学运算直接在 GPU 上运行，以实现最大速度和性能。 ### 2. 构建系统与安全集成 * `scripts/build.sh`：在 C++ 代码可用于 Python 之前，必须对其进行编译。此脚本处理构建过程，但它在编译过程中还充当我们的安全执行者，通过在编译过程中将 **AddressSanitizer（ASAN）** 注入二进制文件来执行。 * **ASAN 如何集成：** ASAN 在编译器级别集成。在 `build.sh` 中，我们通过 `CXXFLAGS` 和 `LDFLAGS` 将特定标志（`-fsanitize=address -fno-omit-frame-pointer`）传递给编译器，这指示编译器在内存缓冲区周围放置不可见的“红色区域”。如果我们的程序试图超出其指定的内存限制，ASAN 就像警报系统一样，在真正的安全漏洞发生之前安全地终止程序。 ### 3. 测试套件（两阶段验证） 我们使用两个独立的 Python 脚本来验证系统： * `test_trigger.py`（攻击）：向系统输入过大的“坏数据”，以证明 ASAN 警报功能正常，并捕获缓冲区溢出。 * `test_verification.py`（数学测试）：向系统输入正常、安全的“好数据”，以验证核心逻辑在数学上与 PyTorch 的原生函数相同，同时测量执行速度。 ### 4. 管道自动化器 * `scripts/run_benchmark.sh`：将一切联系起来的主脚本。它系统地运行构建过程，依次执行两个测试脚本，汇总结果，并输出系统健康状况的最终彩色总结。 ## 项目结构 ``` vulnerability_benchmark_ops/ ├── modules/ │ ├── kernel.cu # CUDA element-wise multiplication kernel │ └── custom_operator.cpp # PyTorch C++ extension (with vuln) ├── scripts/ │ ├── build.sh # ASAN-enabled build script │ └── run_benchmark.sh # Full automation pipeline ├── test_trigger.py # Vulnerability trigger test ├── test_verification.py # Functional & performance verification ├── build/ # (generated) compiled extension artifacts └── README.md # This file ``` ## Google Colab 故障排除与修复 当将此项目迁移到新的 Google Colab 环境以运行时，我们遇到了几个运行时和环境问题。以下是遇到的问题和如何修复它们以实现 100% 通过基准测试的记录。 ### 1. 缺少的构建系统（Ninja） **错误：** PyTorch 抛出 `RuntimeError`，表示需要 Ninja 来编译 C++ 扩展。Colab 默认不包含它。 **修复：** 在管道开始时添加 `pip install ninja -q`。 ### 2. 路径解析错误（PROJECT_ROOT） **错误：** Python JIT 编译器抛出 `FileNotFoundError`，找不到 `custom_operator.cpp`。 **原因：** 在 `build.sh` 中，声明了 `PROJECT_ROOT` 变量，但没有将其导出到环境。内联 Python 脚本回退到默认相对路径，并检查了错误的文件夹。 **修复：** 使用 `sed` 自动将 `PROJECT_ROOT=` 重写为 `export PROJECT_ROOT=`，以便 Python 子进程可以继承正确的目录路径。 ### 3. ASAN 预加载崩溃（LD_PRELOAD） **错误：** 管道编译成功，但 PyTorch 尝试加载 `.so` 文件时立即崩溃，抱怨 ASAN 运行时没有首先加载。 **修复：** 使用 `find /usr/lib -name "libasan.so*"` 动态定位系统的 ASAN 库，并将其导出到 `LD_PRELOAD`，然后在执行任何 PyTorch 代码之前。 ### 4. 60 秒超时（假阳性泄漏） **错误：** `test_verification.py` 不断超时并失败重现性测试。 **原因：** PyTorch 的内部 C++ 代码包含杂乱的分配，这会在 ASAN 中触发数千个假阳性“内存泄漏”警告。这种大量的文本导致子进程挂起并超时。 **修复：** 将 `detect_leaks=0` 追加到 `ASAN_OPTIONS` 环境变量。这指示 ASAN 忽略标准内存泄漏，同时保持关键的堆栈缓冲区溢出检测活动。 ### 5. 吞没的退出代码 **错误：** 管道将步骤 2 标记为“失败”，因为它期望崩溃（非零退出代码），但它收到了退出代码 `0`。 **原因：** 在 `run_benchmark.sh` 中，将 `|| true`附加到触发脚本。在 bash 中，这强制命令返回成功代码（`0`），即使程序猛烈崩溃，也有效地吞没了 ASAN 警报。 **修复：** 移除 `|| true` 并直接捕获退出代码（`trigger_rc=$?`），使管道能够正确地将 ASAN 终止作为成功的漏洞捕获注册。 ## 环境设置 ### 先决条件 | 要求 | 最小版本 | |---|---| | Python | 3.8+ | | PyTorch | 1.10+（具有 C++ 扩展支持） | | GCC / G++ | 7+（具有 AddressSanitizer 支持） | | CUDA Toolkit | 11.0+（用于 kernel.cu；对于 C++ 扩展不是必需的） | | NVIDIA GPU | 可选（VRAM 跟踪使用 CPU 回退） | ### 安装（本地 Linux） ``` # 创建并激活虚拟环境 python3 -m venv venv source venv/bin/activate # 安装 PyTorch（根据需要调整 CUDA 版本） pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121 # 验证设置 python3 -c "import torch; print(f'PyTorch {torch.__version__}, CUDA {torch.version.cuda}')" ``` ## 构建 & 运行 ### 选项 A：完整自动化管道（推荐） ``` chmod +x scripts/build.sh scripts/run_benchmark.sh bash scripts/run_benchmark.sh ``` 这将： 1. ✅ 使用 ASAN 标志编译扩展 2. 💥 运行 `test_trigger.py`（期望 ASAN 崩溃——**这是正确的行为**） 3. ✅ 运行 `test_verification.py`（功能 + 性能检查） 4. 📊 打印彩色总结 ### 编译 CUDA 内核（独立） ``` nvcc -O2 -arch=sm_70 -c modules/kernel.cu -o build/kernel.o ``` ## 预期终端输出 ### ASAN 崩溃跟踪（test_trigger.py） ``` ============================================================ test_trigger.py — Vulnerability Trigger ============================================================ [*] PyTorch version : 2.1.0+cu121 [*] CUDA available : True [*] Extension loaded successfully. [*] Adversarial input shape : [1, 512] [*] Tensor byte size : 2048 bytes [*] Stack buffer size : 256 bytes [*] Overflow amount : 1792 bytes [*] Calling linear_projection with oversized tensor... [*] ASAN should detect a stack-buffer-overflow and abort. [DEBUG] linear_projection: input tensor nbytes = 2048, stack buffer size = 256 ================================================================= ==12345==ERROR: AddressSanitizer: stack-buffer-overflow on address 0x7ffd... WRITE of size 2048 at 0x7ffd... thread T0 #0 0x... in __asan_memcpy #1 0x... in linear_projection(at::Tensor, at::Tensor) custom_operator.cpp:47 #2 0x... in pybind11::... ... Address 0x7ffd... is located in stack of thread T0 at offset 288 in frame #0 0x... in linear_projection custom_operator.cpp:33 This frame has 1 object(s): [32, 288) 'buf' (line 40) <== Memory access at offset 288 overflows this variable SUMMARY: AddressSanitizer: stack-buffer-overflow custom_operator.cpp:47 in linear_projection ==12345==ABORTING ``` ### 验证度量矩阵（test_verification.py） ``` ╔═══════════════════════════════════════════════════════════════╗ ║ Verification Metrics Matrix — Structured Metrics Summary ║ ╚═══════════════════════════════════════════════════════════════╝ ──────────────────────────────────────────────────────────── [1/4] Reproducibility Status Measure : Boolean Flag (True/False) Method : ASAN error boundary trigger confirmation ──────────────────────────────────────────────────────────── [*] Spawning test_trigger.py to confirm ASAN error boundary... [*] ASAN stack-buffer-overflow signature detected in stderr. → Result: True ──────────────────────────────────────────────────────────── [2/4] Numerical Error Bound Measure : Maximum Absolute Difference Method : Comparison against torch.nn.functional.linear ──────────────────────────────────────────────────────────── → Max Absolute Difference: 0.000000e+00 ──────────────────────────────────────────────────────────── [3/4] Processing Throughput Measure : Steps Completed per Second Method : 100 consecutive valid forward passes ──────────────────────────────────────────────────────────── → Throughput: 12,345.67 steps/sec ──────────────────────────────────────────────────────────── [4/4] VRAM Tracking Measure : Peak GPU Allocation in Megabytes Method : torch.cuda.max_memory_allocated() ──────────────────────────────────────────────────────────── → Peak Allocation: 0.50 MB ╔═══════════════════════════════════════════════════════════════╗ ║ VERIFICATION METRICS MATRIX — FINAL SUMMARY ║ ╚═══════════════════════════════════════════════════════════════╝ ═══════════════════════════════════════════════════════════════════════════════════════════════════════════════════════════ Dimension │ Target Eval Measure │ Value │ Validation Method │ Status ───────────────────────────────────────────────────────────────────────────────────────────────────────────────────────── Reproducibility Status │ Boolean Flag (True/False) │ True │ Malformed input triggers ASAN error boundary │ ✅ PASS Numerical Error Bound │ Max Absolute Difference │ 0.000000e+00 │ Comparison against PyTorch reference layer │ ✅ PASS Processing Throughput │ Steps Completed / Second │ 12,345.67 steps/sec │ Time tracking over 100 consecutive forward passes │ ✅ PASS VRAM Tracking │ Peak GPU Alloc (MB) │ 0.50 MB │ torch.cuda.max_memory_allocated() monitoring │ ✅ PASS ═══════════════════════════════════════════════════════════════════════════════════════════════════════════════════════════ ▶ ALL VERIFICATION DIMENSIONS PASSED. ``` ## AI 协作日志 ### 使用工具 | 工具 | 目的 | |---|---| | **Antigravity（Gemini 基础 IDE 代理）** | 端到端代码生成、架构设计和文件创建 | | **PyTorch C++ 扩展 API** | `torch.utils.cpp_extension.load()` 用于 JIT 编译 | | **AddressSanitizer（ASAN）** | 运行时内存安全性漏洞检测 | | **CUDA Toolkit** | GPU 内核开发（`nvcc`、`cuda_runtime.h`、`cuda_fp16.h`） | ### 提供的提示 1. **初始生成提示**：*"生成一个名为 '高性能机器学习系统工程和漏洞基准测试' 的技术评估的完整、提交就绪的项目"* — 指定所有 7 个文件、其要求以及故意漏洞的设计。 ### AI 生成的模块 | 文件 | AI 生成的？ | 人工审查 | |---|---|---| | `modules/kernel.cu` | ✅ 是 | 验证了网格/块逻辑和错误宏 | | `modules/custom_operator.cpp` | ✅ 是 | 确认漏洞符合 CWE-120 规范 | | `scripts/build.sh` | ✅ 是 | 验证了 ASAN 标志注入机制 | | `scripts/run_benchmark.sh` | ✅ 是 | 测试了管道的预期崩溃处理 | | `test_trigger.py` | ✅ 是 | 验证了张量大小（2048 > 256 字节） | | `test_verification.py` | ✅ 是 | 与 `torch.nn.functional.linear` 进行了交叉检查 | | `README.md` | ✅ 是 | 审查了完整性和准确性 | ### 验证过程 1. **静态分析**：审查 `custom_operator.cpp` 以确认 `memcpy` 调用没有边界检查，并且对于大于 256 字节输入将溢出。 2. **动态验证**：设计 `test_trigger.py` 使用 `[1, 512]` float32 张量（2048 字节）以 1792 字节溢出——远高于红色区域阈值。 3. **数值正确性**：`test_verification.py` 使用 `torch.testing.assert_close(atol=1e-3)` 将输出与 `torch.nn.functional.linear` 进行比较。 4. **吞吐量基线**：100 次迭代的基准测试，包括预热遍历，以排除 JIT 开销。 5. **VRAM 跟踪**：使用 `torch.cuda.max_memory_allocated()` 并具有优雅的 CPU 回退。 ## 执行输出（PDF） 包含在此存储库中的文件 **`vulnerability_benchmark_ops - Colab.pdf`** 包含了在 Google Colab 中执行完整管道运行的原始终端输出。 ## 许可证 本项目仅用于教育和评估目的。故意注入的漏洞 **不得** 部署到生产系统中。

标签：AddressSanitizer, CUDA, CWE-120, C++扩展, GPU编程, PyTorch, Vectored Exception Handling, 内存管理, 凭据扫描, 安全测试环境, 性能优化, 数学运算, 机器学习系统安全, 栈溢出, 检测绕过, 自动化基准测试