torch.load攻击新手法
前言
博客太久没更了,找点东西更更XD。这次看BlackHat ppt传出来了就看到了个有意思的Pytorch漏洞,在四五月份的时候已经有相关cve了(CVE-2025-32434),效果是即使torch.load的参数weights_only=True,也能够成功的RCE,不过能否成功rce也要看情况,细节看下文
背景介绍
weights_only
此参数更多的是限制在torch.load中pickle反序列化的能力,在torch早版本时期,利用pickle常规payload能够直接进行代码注入,命令执行等高危操作。于是官方推出此参数来保护用户的安全。在2.6.0版本中正式引入,并且此后版本的pickle.load此参数的默认值为True,即开启保护。
TorchScript
TorchScript 是 PyTorch 提供的一种中间表示形式,它把原本依赖 Python 的动态图模型转换为可序列化、可优化、可独立运行的静态图。通过 torch.jit.trace 或 torch.jit.script 可以生成 TorchScript 模型,并保存为 .pt 文件,用于跨平台部署(如 C++、移动端)。它既保留了 PyTorch 的灵活性,又解决了性能优化和摆脱 Python 环境依赖的问题。
分析
为了方便理解,先拿出一个简单的demo出来
import torch
class MyModule(torch.nn.Module):
def __init__(self):
super(MyModule, self).__init__()
self.linear = torch.nn.Linear(10, 5)
def forward(self):
print(1)
return torch.zeros(0)
module=MyModule()
sc=torch.jit.script(module)
sc.save("pytorch_model.bin")
newModule=torch.load("pytorch_model.bin",weights_only=True)
modins=newModule()
运行demo可以发现
此时我们可以通过这段demo理解到了一部份的漏洞实质,即这个漏洞并非只需要一个简单的load就能够进行攻击的,而是需要load后创造出来的对象进行一些方法调用才可以进行漏洞攻击,demo中使用的是对该对象进行实例化操作,该对象继承自nn.Module,在此过程中会自动的调用对应的forward方法,看下图代码
为了下一步的了解,我们理解到了forward方法能够被调用,那是否意味着我们能够写入任意代码呢,此时我们将forward方法改写成以下情况
def forward(self):
eval("print(1)")
return torch.zeros(0)
运行后产生下图报错
Python builtin <built-in function eval> is currently not supported in Torchscript
这里就了解到了我们的Torchscript,可以看出我们的eval方法并不被Torchscript所支持,因为eval没有对应的TorchScript操作符。所以现在我们再次理解到了漏洞的又一部份实质,即只有在Torchscript中存在操作符的方法才能够被编译成TorchScript使用