29.3.3 ONNX Format Vulnerabilities

After dissecting the code execution risks in framework-specific formats like PyTorch’s Pickle and TensorFlow’s SavedModel, you might view the Open Neural Network Exchange (ONNX) format as a safer harbor. ONNX is designed for interoperability, defining a model as a static computation graph based on Protocol Buffers (Protobuf). This structure-first approach eliminates the immediate arbitrary code execution threat seen during deserialization. However, this safety is conditional and creates a different, more subtle attack surface.

The danger in ONNX shifts from the loading phase to the interpretation and execution phase. An attacker doesn’t need to hide a Python script in the file; instead, they manipulate the very structure of the model graph to exploit the runtime that executes it.

The ONNX Attack Surface: From Deserialization to Interpretation

Unlike Pickle, an ONNX file itself doesn’t contain executable code. It’s a declarative blueprint of operations (ops) and their connections. The security risk depends entirely on how an ONNX-compliant runtime (like ONNX Runtime, TensorRT, or OpenVINO) interprets and implements this blueprint. This creates a trust boundary problem: the model file is untrusted data that dictates the behavior of a trusted runtime process.

Primary Attack Vectors in ONNX Models

An attacker’s goal is to craft an .onnx file that, while structurally valid, causes malicious side effects when processed by a target runtime. Here are the primary methods.

Malicious Custom Operators: The Trojan Horse

The ONNX standard is extensible, allowing for “custom operators” not defined in the official specification. This is a powerful feature for researchers and developers, but it’s also a gaping security hole. An attacker can create a model that references a custom operator, for example, com.malicious.ExecuteCommand.

The .onnx file itself only contains the reference. The real payload lies in a separate, malicious shared library (.dll, .so) that implements this operator. If a victim’s environment is tricked into loading this library to satisfy the model’s dependency, the attacker gains code execution within the runtime’s process every time the model is run.

Graph-Based Denial of Service (DoS)

The model graph itself can be weaponized. An attacker can construct a model that is computationally explosive, designed to consume disproportionate resources during graph parsing or execution. This can include:

• Excessive Memory Allocation: Defining tensors with absurdly large dimensions (e.g., [2147483647, 2147483647, 3]) that the runtime might try to allocate, leading to an out-of-memory crash.
• Computational Bombs: Creating graphs with exponential complexity, such as deeply nested loops or recursive-like structures that, while small in file size, unroll into an unmanageable number of operations at runtime.
• Infinite Loops: Certain combinations of operators like Loop and control flow primitives can be crafted to create non-terminating execution paths.

Exploiting Runtime-Specific Operator Bugs

The security of an ONNX model is only as strong as the runtime executing it. A C++ implementation of a standard operator like Conv or Gather in a specific version of ONNX Runtime could have a classic memory safety vulnerability (e.g., a buffer overflow or integer overflow). An attacker can craft a model with specific tensor shapes, strides, or attributes that trigger this latent bug in the runtime, leading to a crash or potentially code execution.

This turns model files into exploit delivery vehicles for vulnerabilities in the underlying ML infrastructure, a critical vector for red teams to investigate.

Inspecting a Model for Suspicious Operators

A basic first step in defense is programmatic inspection. Before loading a model into a runtime for execution, you can parse it with a library like onnx to audit its contents, particularly the operator domains.

import onnx

# Load the untrusted ONNX model without executing it
model_path = "untrusted_model.onnx"
model = onnx.load(model_path)

# Define the standard, expected ONNX domain
# An empty string "" denotes the standard ONNX opset
standard_domain = ""
allowed_domains = {standard_domain, "ai.onnx.ml"}

print(f"Analyzing operators in {model_path}...")

# Iterate through all nodes in the model graph
for node in model.graph.node:
    if node.domain not in allowed_domains:
        print(f"[!] WARNING: Non-standard operator domain found!")
        print(f"    Node Name: {node.name}")
        print(f"    Op Type:   {node.op_type}")
        print(f"    Domain:    '{node.domain}' <-- Suspicious")
    else:
        # You could add further checks here for allowed op types
        pass

print("nAnalysis complete.")

This script won’t find every threat, but it immediately flags the most obvious sign of a custom operator attack: a non-standard operator domain.

Defensive Postures and Mitigation Strategies

Securing your pipeline against malicious ONNX models requires a multi-layered approach focused on validation and containment.

• Model Provenance is Paramount: The most effective defense is to only use models from cryptographically verified sources. Use digital signatures to ensure the model you’re loading is the one the publisher intended.
• Strict Operator Allow-listing: Configure your ONNX runtime to execute only a pre-approved list of standard operators. Explicitly disable support for custom operators unless they are from a vetted internal source.
• Resource-Constrained Sandboxing: Execute all untrusted models in a heavily sandboxed environment (e.g., a container with strict cgroup limits, seccomp filters, and no network access). This contains the blast radius of a DoS attack or a successful runtime exploit.
• Static Graph Analysis: Before execution, perform deep static analysis on the model graph. Check for insane tensor dimensions, excessive node counts, unusual graph depth, and other structural anomalies that suggest a DoS attempt.
• Keep Runtimes Patched: Treat your ONNX runtime as a critical piece of infrastructure. Monitor it for CVEs and apply security patches diligently to protect against known implementation bugs.

Ultimately, an ONNX file should never be treated as inert data. It is a set of instructions for your high-performance computing hardware. Granting it trust by default is an invitation for attack.