CVE-2024-5480 Vulnerability in PyTorch's torch

 CVE-2024-5480: Remote Code Execution Vulnerability in PyTorch's torch.distributed.rpc Framework

June 7, 2024

    A critical remote code execution (RCE) vulnerability has been identified in the PyTorch's torch.distributed.rpc framework, affecting versions prior to 2.2.2. This vulnerability allows attackers to execute arbitrary commands by leveraging built-in Python functions such as eval during multi-cpu RPC communication. The issue arises from the lack of restriction on function calls when a worker node serializes and sends a Python User Defined Function (PythonUDF) to the master node, which then deserializes and executes the function without validation. This flaw can be exploited to compromise master nodes initiating distributed training, potentially leading to the theft of sensitive AI-related data.

Background

    The torch.distributed.rpc framework is used in distributed training scenarios to facilitate communication between worker nodes and the master node. The framework is designed to handle RPC operations efficiently, but it lacks proper verification of the functions being called during these operations. This oversight allows attackers to execute arbitrary commands by leveraging built-in Python functions such as eval.

Vulnerability Details

    The vulnerability arises from the following sequence of events:

    Serialization and Sending: A worker node serializes and sends a PythonUDF to the master node.

    Deserialization and Execution: The master node deserializes and executes the PythonUDF without validation.

    This sequence of events allows attackers to execute arbitrary commands by crafting a malicious PythonUDF that includes built-in Python functions such as eval. The eval function can be used to execute arbitrary code, which can lead to remote code execution.

Impact

    The impact of this vulnerability is significant, as it can be exploited to compromise master nodes initiating distributed training. This could potentially lead to the theft of sensitive AI-related data, including model weights, training data, and other confidential information.

Mitigation

    To mitigate this vulnerability, it is recommended that users upgrade to the latest version of PyTorch (2.2.2 or later) that includes the necessary fixes. Additionally, users should ensure that their distributed training environments are properly configured to prevent unauthorized access to the master node.

Conclusion

    The CVE-2024-5480 vulnerability highlights the importance of proper verification and validation of functions during RPC operations in distributed training scenarios. It is crucial for developers and users to stay up-to-date with the latest security patches and best practices to prevent such vulnerabilities from being exploited.

How can I patch my PyTorch installation to fix CVE-2024-5480

    To patch your PyTorch installation and fix CVE-2024-5480, you need to upgrade to PyTorch version 2.2.2 or later. The vulnerability affects versions prior to 2.2.2. Here are the steps to upgrade PyTorch:

1. Open your terminal or command prompt. 2. Run the following command to install the latest version of PyTorch:

pip install --upgrade torch

3. Verify the installed version by running:

python -c 

import torch; 

print (torch.__version__)

    The output should display version 2.2.2 or higher, indicating that the vulnerability has been patched.     It is crucial to upgrade PyTorch as soon as possible to protect your distributed training environments from potential exploitation of CVE-2024-5480. The vulnerability allows attackers to execute arbitrary commands by leveraging built-in Python functions like eval during multi-cpu RPC communication.     By upgrading to the latest version, you will ensure that your PyTorch installation is secure and that the torch.distributed.rpc framework properly verifies the functions being called during RPC operations.