Overview

PAIT-TCHST-300

TorchScript Model Arbitrary Code Execution Detected at Model Load Time

Overview

TorchScript is PyTorch's intermediate representation format that allows models to be serialized and run independently of Python. The format saves models as directories containing both Python source files and Pickle files, creating multiple attack vectors for malicious code injection.

Models flagged for this threat meet the following criteria:

1. The model format is detected as TorchScript.
2. The model contains potentially malicious code in internal Python files or Pickle files which will run when the model is executed.

Pickle is the original serialization Python module used for serializing and deserializing Python objects to share between processes or other computers. While convenient, Pickle poses significant security risks when used with untrusted data, as it can execute arbitrary code during deserialization. This makes it vulnerable to remote code execution attacks if an attacker can control the serialized data.

In this case, running the model will execute the code, and whatever malicious instructions have been inserted into it.

Key Points:

1. TorchScript models should be treated as executable code packages
2. The TorchScript format saves model code, data, and instructions for inference in a directory, containing both Python files and Pickle files.
3. Both Python and Pickle files can contain malicious code that will run automatically when the model is executed.
4. Only use/load models from trusted sources.

Background Information

TorchScript Format

TorchScript uses a directory-based serialization format with the following base structure:

1model_directory/
2├── code/
3│   ├── __torch__.py
4│   └── __torch__/
5│       └── module_implementations.py
6├── data.pkl
7├── constants.pkl
8└── version
9

The format includes:

• Python source files (.py) containing the model's computational graph and custom operations
• Pickle files (.pkl) storing serialized tensors, parameters, and metadata
• Version files for compatibility checking

This code-inclusive serialization extends portability but creates security vulnerabilities that attackers can exploit to distribute malicious code.

Further reading:

1. PyTorch TorchScript Documentation
2. TorchScript Serialization Format
3. Never a dill moment: Exploiting machine learning pickle files.

Impact

This attack can harm your organization in any of the following ways:

1. Sharing data scientist / company github credentials with the attacker.
2. Sharing data scientist / company cloud credentials, resulting in IP theft of data and other models, all the way to cloud take over attacks.
3. Altering the behavior of expected models; allowing fraud, model performance impacts, and adversarial by-passes.
4. Corruption of other models, by leveraging credentials to rewrite known model assets within your environments.

How The Attack Works

Best Practices

You should:

1. Only load and execute models from trusted sources.
2. Implement a vetting process for third-party models before use, including scanning.
3. Use sandboxing techniques when loading untrusted models.
4. Use model formats that don't allow arbitrary code execution, such as SafeTensors, which provides a safe alternative to traditional serialization formats

Remediation

If possible, use a different model format like SafeTensors in order to remove this type of code injection attack from impacting your work entirely.

If not possible, reach out to the model creator and alert them that the model has failed our scan. You can even link to the specific page on our Insights Database to provide our most up to date findings.

The model provider should also report what they did to correct this issue as part of their release notes.