- 作者: 纯情
- 时间:
- 分类: 资讯
PyTorch 最新版本反序列化漏洞分析 前言 AI入门学习Transformers架构的时候，最忘不掉的无过于主播手把手教你用PyTorch手搓一个Transformers了，哈哈哈哈哈哈 一、漏洞描述 根本原因在于调用 torch.load 时使用了 weights_only=False，这启用了通过 pickle 模块进行的不安全反序列化，从而导致任意代码执行。torch.load 底层使用了 Python 的 pickle。当 weights_only=False 时，它会反序列化并重构文件中存储的任意 Python 对象。对来自不受信任或被篡改的检查点数据进行反序列化，会执行攻击者控制的代码，从而导致远程代码执行 (RCE)。 当加载 .pt2 归档文件时，torch.export.load() 遵循以下路径：它调用 load_pt2() 来枚举归档内容 → 进入 _load_exported_programs() → 然后调用 _load_constants()。如果归档中的任何常量被标记为 use_pickle=True 且其文件名以 tensor 开头，它就会执行 torch.load(io.BytesIO(constant_bytes), weights_only=False)，从而触发 Python Pickle 反序列化。因此，如果 .pt2 文件包含恶意的 Pickle 常量，运行 torch.export.load() 可能会导致任意代码执行。 torch.export.load() 不允许配置 weights_only=True 参数。它总是执行 torch.load(io.BytesIO(constant_bytes), weights_only=False) —— 这个值是硬编码的。因此，加载恶意模型的用户不可避免地会面临被入侵的风险。 二、环境搭建 这个没有什么好搭建的，因为就是一个组件，下载最新版本的包就好了 三、漏洞分析 完整调用链追踪 漏洞代码分析 torch.export.load() 文件: torch/export/__init__.py 分析: 入口函数接收文件路径，调用 load_pt2() 加载 PT2 存档文件。 load_pt2() - 加载 PT2 存档 文件: torch/export/pt2_archive/_package.py 分析: load_pt2() 打开 PT2 存档文件，调用 _load_exported_programs() 处理模型文件。 _load_exported_programs() - 加载导出的程序 文件: torch/export/pt2_archive/_package.py 分析: _load_exported_programs() 遍历模型文件，调用 _load_constants() 加载常量数据。 _load_constants() - 加载常量（漏洞触发点） 文件: torch/export/pt2_archive/_package.py 这一步很关键，因为是有很多限制条件的 总结一下能够触发的条件 1path_name.startswith("tensor") 2payload_meta.use_pickle == True 3 weights_only=False 硬编码: 该值无法通过 torch.export.load() API 覆盖 torch.load() - Pickle 反序列化 文件: torch/serialization.py 当 weights_only=False 时，torch.load() 使用标准 pickle 反序列化，会触发对象的 __reduce__() 方法，导致任意代码执行。 四、漏洞复现 根据漏洞原理，步骤就是生成恶意文件和触发了 恶意模型生成脚本 文件: model.py 注意构造需要满足刚刚上面提到的条件 触发脚本 现在就是我们需要去触发了 文件: poc.py 然后我们查看文件 五、漏洞修复 由于最新版本还是存在这个漏洞，可能是正常功能需要????？ 测试版本: PyTorch 2.11.0.dev20260110 六、参考资料 资源 | 链接 |
PyTorch 官方文档 | |
CWE-502: 不安全的反序列化 | |
Python Pickle 安全 | |
赏金报告 | |
标签: PyTorch, 远程代码执行, 反序列化漏洞, torch.load, Pickle, CWE-502, torch.export.load
