Skip to content

PyTorch weights_only 加载链的内存破坏风险 ​

摘要 ​

One Chain to Own Them All 分析了 PyTorch torch.load(..., weights_only=True) 的两代安全问题。CVE-2025-32434 证明旧实现仍可经受限反序列化路径执行代码;CVE-2026-24747(GHSA-63cw-57p8-fm3p)进一步表明,即使对象构造白名单生效,恶意 checkpoint 仍可借张量存储大小与元数据不一致破坏内存,并在特定条件下发展为代码执行。

公开公告标注 CVE-2026-24747 影响 PyTorch 2.9.1 及更早版本,2.10.0 修复,CVSS 为 8.8。演讲还追踪 vLLM、OpenLLM、NVIDIA Dynamo、SGLang 和 ComfyUI 等下游入口。本文只把已展示的数据流作为受影响路径,不把“能到达 torch.load”等同于已经完成远程代码执行。

两代 weights_only 问题不能混为一谈 ​

CVE-2025-32434(GHSA-53q9-r3pm-6pq6)发生在 PyTorch 2.5.1 及更早版本:即使调用方指定 weights_only=True,旧式归档兼容路径仍可能进入不安全反序列化。PyTorch 2.6.0 修复了这条对象构造绕过。

CVE-2026-24747 则是另一类问题。受限 Unpickler 确实只允许预期的 storage、tensor 和少量容器对象,但这些“允许对象”的底层长度关系没有被完整校验。攻击面从 Python 的任意对象构造下沉到了 C++ 存储对象、归档记录长度、张量偏移与 Pickle 栈操作。官方公告标注影响 < 2.10.0,2.10.0 修复。

问题首先失效的控制主要修复方向
CVE-2025-32434weights_only 没有覆盖旧格式加载路径所有格式统一进入安全加载路径
CVE-2026-24747允许对象内部的长度不变量未校验归档实际字节数必须与声明 storage 大小一致

因此,升级到 2.6.0 只能修复前一问题,不能自动消除后一问题。

正常 checkpoint 是怎样重建张量的 ​

PyTorch ZIP checkpoint 把结构和数据分开保存。data.pkl 描述 storage 类型、键、设备、元素个数以及如何调用 _rebuild_tensor_v2;实际字节位于 data/0 等归档记录中。一个概念化的正常流程是:

text
persistent_load(("storage", LongStorage, "0", "cpu", 10))
  -> 为 10 个 int64 元素取得 80 字节 storage

_rebuild_tensor_v2(storage, offset=0, size=(10,), stride=(1,), ...)
  -> 视图覆盖 [0, 80) 字节,满足不变量

Python 层 persistent_load 会按 numel × element_size(dtype) 计算期望字节数,然后调用 load_tensor。问题在于旧实现继续进入 get_storage_from_record(name, numel, ...),底层 getRecord 返回实际归档数据,却没有在创建 storage 前确认“实际记录长度等于声明长度”。

根因:声明容量与真实分配脱节 ​

演讲者的第一轮尝试只是让张量视图越过 storage 末端。例如,在 80 字节 storage 上使用偏移 1、长度 10 的 int64 张量,需要 88 字节。PyTorch 发现空间不足并尝试调整 storage 大小,但该 storage 标记为不可调整,于是抛出 RuntimeError。这一步没有形成漏洞,反而帮助研究者定位安全检查所在的位置。

随后作者让 Pickle 元数据声明 storage 含有更多元素,而归档中的 data/0 仍只有较少字节。旧 C++ 路径使用实际记录创建底层内存,却在上层保留较大的逻辑容量。后续张量重建和受限 Pickle 指令便可能按照声明容量执行索引或赋值,最终访问真实分配之外的内存。

可利用条件由四个值共同决定:

text
归档记录实际字节数
storage 声明的元素个数 × 元素大小
tensor 的 storage_offset
tensor 的 size 与 stride 所覆盖的最大索引

只检查 shape 上限或只限制 Pickle 全局对象,都无法发现这种不一致。

从越界访问到内存写入 ​

受限 Unpickler 仍需支持权重文件中的字典和列表,因此保留了 SETITEM、SETITEMS 等指令。研究利用链把畸形 tensor 放到这些合法容器操作可触达的位置,使受限解释器执行本应正常的索引赋值,但目标对象的底层边界已被伪造。演讲进一步展示了信息泄露、地址定位和控制流劫持。

稳定利用依赖运行环境。作者指出,某些 Debian Python 构建为非 PIE,初始演示可以较容易定位 system;对启用地址随机化的服务,又需要先从错误响应或内存布局中取得地址,再完成后续阶段。本文保留这一研究结论,但不复刻 Pickle opcode、堆布局、任意地址读写和控制流劫持步骤。

修复机制:检查实际记录长度 ​

演讲展示的 PyTorch 修复(PR 170085)在取出归档记录后获得实际 size,并用 TORCH_CHECK 验证:

text
actual_record_size == numel * element_size(scalar_type)

这项检查位于创建 storage 之前,因此畸形制品在任何张量视图或容器赋值发生前被拒绝。它比在 _rebuild_tensor_v2 处补一个 shape 检查更完整,因为根因是物理记录与逻辑 storage 的身份不一致。

PyTorch 安全说明也删除了“weights_only=True 即安全”的绝对表述,改为强调所有模型都包含攻击面,应对不可信模型做来源校验并在隔离环境中加载。

下游远程入口:并非只有模型上传 ​

研究者没有停留在本地 checkpoint,而是搜索哪些服务把网络输入变成 BytesIO 后交给 torch.load。

vLLM:聊天请求中的 image_embeds ​

演讲追踪的调用链从 /v1/chat/completions 开始。多模态消息允许 type: image_embeds,其字符串经聊天解析、媒体连接器和 base64 解码后进入图像加载函数,最终执行:

text
API 请求
  -> serving_chat / serving_engine
  -> chat_utils 的多模态内容解析
  -> _parse_image_embeds
  -> media connector
  -> image.py: load_bytes(base64.b64decode(data))
  -> torch.load(BytesIO(data), weights_only=True)

这使一个看似“图片嵌入”的字段具备 checkpoint 解析语义。是否能远程利用仍取决于 vLLM 版本、PyTorch 版本、接口是否暴露和多模态配置。

OpenLLM:权重更新与兼容格式 ​

材料检查了从磁盘更新权重的管理入口以及 load_pt_file。正常路径使用 weights_only=True,但下载来源、管理端点暴露范围和旧格式兼容路径共同决定攻击面。它更接近“被授权的权重更新导致供应链执行”,不能简单归类为未认证远程代码执行。

SGLang:错误传播变成进程级拒绝服务 ​

同一畸形文件在 SGLang 中的直接后果与 vLLM 不同。解析异常发生在 worker,进程退出后调度器向进程树发送终止信号,导致服务级中断,而不是只返回一次 HTTP 500。该案例说明加载器漏洞的影响受进程监督模型显著影响。

ComfyUI:本地工作流中的多阶段利用 ​

CheckpointLoaderSimple 经 load_checkpoint_guess_config 和 comfy.utils.load_torch_file 到达 torch.load(..., weights_only=True)。演讲展示先泄露 libtorch_python.so 地址,再定位 libtorch_cpu.so 和可用调用点,最终完成代码执行。ComfyUI 自动下载的最新 PyTorch 在作者验证时已经包含修复,但旧环境或固定依赖仍需单独核对。

NVIDIA Dynamo:提示词嵌入跨进程进入 Python worker ​

Rust 前端接收编码的 prompt_embeds,转交 Python worker 解码并加载。错误响应还可能成为地址信息的返回通道。材料显示后续修复通过配置项 enable_prompt_embeds 把该能力改为显式启用,默认拒绝外部嵌入;这属于减少入口暴露,仍应同时升级底层 PyTorch。

去武器化 PoC:验证拒绝路径而不制造内存破坏 ​

以下测试格式不是 PyTorch checkpoint,只模拟同类长度不变量。它适合验证网关或预解析器在把字节交给运行时前能否拒绝矛盾元数据:

python
from dataclasses import dataclass

@dataclass(frozen=True)
class StorageMeta:
    elements: int
    element_size: int
    record_size: int

def validate_storage(meta: StorageMeta) -> None:
    expected = meta.elements * meta.element_size
    if expected != meta.record_size:
        raise ValueError("DRY_RUN_STORAGE_LENGTH_MISMATCH")

validate_storage(StorageMeta(elements=10, element_size=8, record_size=80))

负向夹具将 record_size 改为 72,断言解析器返回固定错误、主服务不崩溃、错误响应不包含指针或模块地址,并且没有进入 torch.load。本站不提供可被 PyTorch 接受的畸形 data.pkl、归档补丁脚本或利用链。

资产盘点与验证方法 ​

仅在代码中搜索“上传模型”会漏掉大量入口。建议建立来源到加载点清单:

输入类型常见伪装应检查的转换
模型或 checkpoint.pt、.pth、工作流节点ZIP/Pickle 直接加载
多模态嵌入image_embeds、prompt_embedsbase64 → bytes → torch.load
适配器或热更新LoRA、update weights下载或共享磁盘 → 加载器
缓存与进程间消息预计算 tensorIPC 字节流 → Python worker
  • 升级到 PyTorch 2.10.0 或厂商确认包含同等补丁的版本,并在实际镜像内读取运行时版本。
  • 仅加载来源可信且摘要固定的制品;weights_only 不提供发布者认证或完整性保证。
  • 对归档实际长度、storage 元素数、dtype、偏移、shape 与 stride 做联合一致性检查,并设置资源上限。
  • 在无网络、无长期凭据、低权限、短生命周期进程中解析外部模型;解析器崩溃不得带走调度器或主 API。
  • 错误响应只返回稳定错误码,不回显指针、模块基址、内部异常栈或原始 Pickle 内容。
  • 默认关闭网络可达的二进制嵌入和权重更新功能,需要时为入口单独认证、限流并记录制品摘要。

关键结果与实际影响 ​

核心结果是:weights_only=True 只缩小了对象级攻击面,并没有把 Python/C++ 混合加载器变成内存安全解析器。受影响的下游项目共享同一个加载点,但前置权限、默认配置和进程模型不同,影响可以表现为单请求错误、服务级拒绝服务、信息泄露或代码执行。

局限与待验证问题 ​

从越界访问到稳定代码执行依赖分配器、操作系统、PIE/ASLR、编译选项和错误回显。演讲提供了多环境利用过程,但没有形成可外推的成功率。下游项目变化快,本文记录的是材料与 2026-08-13 公告可核对的状态;部署方仍需根据自身镜像、功能开关和进程监督方式复测。

参考链接 ​