PyTorch 最新版本反序列化漏洞分析 作者: 纯情 时间: 2026-01-19 分类: 资讯 评论 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 官方文档https://pytorch.org/docs/stable/export.htmlCWE-502: 不安全的反序列化https://cwe.mitre.org/data/definitions/502.htmlPython Pickle 安全https://docs.python.org/3/library/pickle.html#restricting-globals赏金报告https://huntr.com/bounties/9eb184db-09c2-4a98-977d-d967342300e8
从 VCTF2025 ez_train学习torch.load反序列化绕过 作者: 纯情 时间: 2026-01-15 分类: 开源 评论 返回文档 从 VCTF2025 ez_train学习torch.load反序列化绕过 前言 感觉挺有意思的一道题,赛后看了一下,这个题目主要考察了代码审计和 torch.load 在 weights_only=True 条件下的利用。 题目环境搭建 下载题目附件并解压, 题目直接给了 docker 环境,但是是 linux 环境的,为了方便待会调试这里按照题目版本去官网上下载一个 windows 版本的, 然后按照题目给的 requirements.txt 中的依赖进行下载,下载好后运行 server.py 启动,启动后如下, 这里还注意到题目在 text-generation-webui-3.13\user_data\training\datasets 目录中放了一些文件,一起复制过来。 题目分析 查看题目描述: 结合题目名称我们应该主要看 train 模块了,全局搜索发现就只有个 trianing.py 文件 然后像这类微调训练大模型的本地 web 应用比较耳熟能详的就是 torch.load 引起的反序列化漏洞了,而这个 training 中就恰好存在这个方法的调用,不难看出这里应该就是出题人的考点了, 这段代码大概意思是从 adapter_model.bin 加载 LoRA 的权重参数(只包含 LoRA,不包含基础模型),接着将加载的 LoRA 权重注入到当前的 lora_model 中,如果这个文件内容我们可以控制这里是不是就能进行反序列化了呢? 我们先看这个 {lora_file_path} 能不能控制,简单溯源一下, 发现 {lora_file_path} 是由 {Path(shared.args.lora_dir)} 和 {lora_file_path} 拼接而成的,{lora_file_path} 是 lora_name 的值,而这个 lora_name 的值其实就是我们训练时传入的参数, 至于 {Path(shared.args.lora_dir)} 的值我们调试发现为 user_data/loras 目录,是个固定值, 那么现在知道 {lora_file_path} 就一小部分能控制,感觉没什么用啊,而且还需要我们能把文件 adapter_model.bin 文件上传到这个目录下。 后面找了一圈没找到哪里可以上传,但是发现这个应用可以从 huggingface 远程下载 model,而且下载目录就在我们的 user_data/loras 下面吗。简单测试一下,这里填入“paulinsider/llamafactory-hack”,然后点击下载, 确实成功把文件下载到了我们要的目录。 那么这样思路就清晰了,我们就可以在远程 model 中放入恶意的 adapter_model.bin 文件,然后下载到本地,通过控制训练时的 lora_name 参数使得最后 torch.load 加载的文件是我们的恶意文件, 这里简单尝试一下,假设刚刚下载的 paulinsider_llamafactory-hack 目录下存在恶意 adapter_model.bin 文件,我们填写训练 lora_name 为 paulinsider_llamafactory-hack 点击训练后,看到成功使得 torch.load 的文件是我们能控制的文件了, 但是这里还是没法进行反序列化,因为这里得 lora_name 设置了weights_only 参数为 True。 torch.load 反序列化绕过 查看题目的requirements.txt 文件发现给的 torch 版本为 2.5.1,网上简单搜了下不难发现其实还是存在绕过方法的, 参考 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 方法, 然后一直到 pickle_module.load 方法, 和正常的 pickle.load() 不一样,这里会调用到 Unpickler.load() 方法,这里面会依次读取字节,然后对每个字节进行判断,对全局变量和函数都有白名单进行限制 只能用白名单中提供的, 这个白名单是绕不过了,其中也没什么有用的,根据参考文章知道在 torch.load 中其实还调用了 torch.jit.load 方法,而且不受weights_only 参数影响,这个方法可以加载一个已经用 TorchScript 保存好的模型,用于直接推理(不需要原始 Python 代码) 所以现在关键是看怎么生成个恶意的 TorchScript 模型,这里面涉及到了 TorchScript 运算符的概念,根据参考文章知道作者是发现了 torch.save 最后可以调用到 aten::save 运算符实现写文件操作,所以最后生成恶意的 TorchScript 模型代码如下 这里注意到在 torch.load 代码下面还有个 newModule.items() 方法调用 ,这里是漏洞触发点,如果只返回了 TorchScript 模型对象是没办法触发到里面的 torch.save 方法, 回到题目,我们可以让adapter_model.bin 为恶意的 TorchScript 模型,然后我们还需要找下触发点, 看到返回 state_dict_peft 后调用了 set_peft_model_state_dict 方法进行处理,跟进这个 set_peft_model_state_dict ,把 state_dict_peft 赋值给了state_dict 后面又调用到了 _insert_adapter_name_into_state_dict 最后发现调用的还是 items 方法 所以直接用上面的 poc 就可以了,这里就直接把生成的恶意文件复制到 paulinsider_llamafactory-hack 目录下,然后和上面一下训练 成功实现写文件操作 查看题目给的 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