标签 torch.load 下的文章

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 安全
赏金报告


从 VCTF2025 ez_train学习torch.load反序列化绕过

前言

感觉挺有意思的一道题,赛后看了一下,这个题目主要考察了代码审计和 torch.load 在 weights_only=True 条件下的利用。

题目环境搭建

下载题目附件并解压,

file-20260113204022005.png



题目直接给了 docker 环境,但是是 linux 环境的,为了方便待会调试这里按照题目版本去官网上下载一个 windows 版本的,

file-20260113203804115.png



然后按照题目给的 requirements.txt 中的依赖进行下载,下载好后运行 server.py 启动,启动后如下,

file-20260113204526100.png



这里还注意到题目在 text-generation-webui-3.13\user_data\training\datasets 目录中放了一些文件,一起复制过来。

题目分析

查看题目描述:

file-20260113204951881.png



结合题目名称我们应该主要看 train 模块了,全局搜索发现就只有个 trianing.py 文件

file-20260113205156648.png



然后像这类微调训练大模型的本地 web 应用比较耳熟能详的就是 torch.load 引起的反序列化漏洞了,而这个 training 中就恰好存在这个方法的调用,不难看出这里应该就是出题人的考点了,

file-20260113205625132.png



这段代码大概意思是从 adapter_model.bin 加载 LoRA 的权重参数(只包含 LoRA,不包含基础模型),接着将加载的 LoRA 权重注入到当前的 lora_model 中,如果这个文件内容我们可以控制这里是不是就能进行反序列化了呢?

我们先看这个 {lora_file_path} 能不能控制,简单溯源一下,

file-20260113210229786.png



发现 {lora_file_path} 是由 {Path(shared.args.lora_dir)}{lora_file_path} 拼接而成的,{lora_file_path}lora_name 的值,而这个 lora_name 的值其实就是我们训练时传入的参数,

file-20260113211019396.png


file-20260113211055508.png



至于 {Path(shared.args.lora_dir)} 的值我们调试发现为 user_data/loras 目录,是个固定值,

file-20260113210937792.png



那么现在知道 {lora_file_path} 就一小部分能控制,感觉没什么用啊,而且还需要我们能把文件 adapter_model.bin 文件上传到这个目录下。

后面找了一圈没找到哪里可以上传,但是发现这个应用可以从 huggingface 远程下载 model,而且下载目录就在我们的 user_data/loras 下面吗。简单测试一下,这里填入“paulinsider/llamafactory-hack”,然后点击下载,

file-20260113212126010.png



确实成功把文件下载到了我们要的目录。

file-20260113212353132.png



那么这样思路就清晰了,我们就可以在远程 model 中放入恶意的 adapter_model.bin 文件,然后下载到本地,通过控制训练时的 lora_name 参数使得最后 torch.load 加载的文件是我们的恶意文件,

这里简单尝试一下,假设刚刚下载的 paulinsider_llamafactory-hack 目录下存在恶意 adapter_model.bin 文件,我们填写训练 lora_name 为 paulinsider_llamafactory-hack

file-20260113213321532.png



点击训练后,看到成功使得 torch.load 的文件是我们能控制的文件了,

file-20260113213310006.png



但是这里还是没法进行反序列化,因为这里得 lora_name 设置了weights_only 参数为 True。

torch.load 反序列化绕过

查看题目的requirements.txt 文件发现给的 torch 版本为 2.5.1,网上简单搜了下不难发现其实还是存在绕过方法的,

file-20260113214328275.png



参考 https://i.blackhat.com/BH-USA-25/Presentations/US-25-Jian-Lishuo-Safe-Harbor-or-Hostile-Waters.pdf ,我们先看看这个 weights_only 的工作原理

生成个恶意的 pickle 反序列化文件

torch.load demo

进入 torch.load 中看到当 weights_only 为 true 时会调用到 _legacy_load 方法,

file-20260113215502623.png



然后一直到 pickle_module.load 方法,

file-20260113215923584.png



和正常的 pickle.load() 不一样,这里会调用到 Unpickler.load() 方法,这里面会依次读取字节,然后对每个字节进行判断,对全局变量和函数都有白名单进行限制

file-20260113220408194.png



只能用白名单中提供的,

file-20260113220439125.png



这个白名单是绕不过了,其中也没什么有用的,根据参考文章知道在 torch.load 中其实还调用了 torch.jit.load 方法,而且不受weights_only 参数影响,这个方法可以加载一个已经用 TorchScript 保存好的模型,用于直接推理(不需要原始 Python 代码)

file-20260113220646171.png



所以现在关键是看怎么生成个恶意的 TorchScript 模型,这里面涉及到了 TorchScript 运算符的概念,根据参考文章知道作者是发现了 torch.save 最后可以调用到 aten::save 运算符实现写文件操作,所以最后生成恶意的 TorchScript 模型代码如下

这里注意到在 torch.load 代码下面还有个 newModule.items() 方法调用 ,这里是漏洞触发点,如果只返回了 TorchScript 模型对象是没办法触发到里面的 torch.save 方法,

回到题目,我们可以让adapter_model.bin 为恶意的 TorchScript 模型,然后我们还需要找下触发点,

file-20260113221624043.png



看到返回 state_dict_peft 后调用了 set_peft_model_state_dict 方法进行处理,跟进这个 set_peft_model_state_dict ,把 state_dict_peft 赋值给了state_dict

file-20260113221730914.png



后面又调用到了 _insert_adapter_name_into_state_dict

file-20260113221818019.png



最后发现调用的还是 items 方法

file-20260113221958131.png



所以直接用上面的 poc 就可以了,这里就直接把生成的恶意文件复制到 paulinsider_llamafactory-hack 目录下,然后和上面一下训练

file-20260113223421839.png



成功实现写文件操作

file-20260113223444024.png



查看题目给的 docker 文件不难发现最后应该是通过写定时任务实现 rce。

参考

https://i.blackhat.com/BH-USA-25/Presentations/US-25-Jian-Lishuo-Safe-Harbor-or-Hostile-Waters.pdf

https://github.com/ChaMd5Team/Venom-WP/blob/main/2025VenomCTF/2025_venomctf_web_ez_train.pdf

https://mp.weixin.qq.com/s/2VsxBTIiX5P8b16Rl1ylcA