Introduction

Let's talk about an often overlooked attack surface in AI systems: model file formats. Sure, everyone focuses on API security and web vulnerabilities, but there's a whole world of potential bugs hiding in how these systems load and process model files. Think about it - after training a model like ChatGPT, you need to save those weights somewhere. These files aren't just simple data blobs - they're complex formats like Pickle, ONNX, Safetensors, and GGUF, each with their own parsing quirks and potential vulnerabilities.

The critical moment? When a user loads that model file into memory. That's where the fun starts - buffer overflows, code execution, all the classics we know and love from traditional binary exploitation, but in a fresh context that hasn't been picked clean by decades of security research.

This guide will walk you through hunting for 0-days in these model formats. We're talking about bugs that could let you hide malicious payloads, manipulate memory, or plant backdoors that only trigger under specific conditions. For bug hunters, this is largely uncharted territory - an expanding attack surface that most people aren't even looking at yet. And here's the best part: you don't need to be a machine learning expert or master reverse engineer to find these bugs. Just bring your curiosity and methodical approach to breaking things.

GGUF Heap Overflow: A Vulnerability in Model Parsing

Understanding the Vulnerability

In frameworks using the new GGUF format (used by libraries like ggml), parsing logic that reads the model's key-value pairs can be misused. If the code doesn't properly validate header fields (like n_kv, the number of key-value entries), it can lead to out-of-bounds writes on the heap.

How the Vulnerability Works 

• The file header includes a n_kv field that tells the loader how many key-value pairs to expect
• The loader allocates an array based on this number
• If n_kv is huge or manipulated, the code writes past the allocated memory, potentially corrupting the heap

Example Code Snippet 

Key Takeaways from the Exploit 

Memory corruption bugs in ML model loaders often occur where file parsing meets memory allocation. Even seemingly simple header parsing can lead to exploitable issues if values aren't properly validated.

Getting Started with Vulnerability Analysis 

• Search for file parsing functions that read headers and model metadata
• Look for memory allocations based on values read from files
• Check for loops that rely on untrusted counters from the file

Finding Similar Vulnerabilities:

Tools You’ll Need for Hunting Vulnerabilities 

• Hex editor (like 010 Editor)
• AddressSanitizer (ASAN)
• GDB/LLDB debugger
• American Fuzzy Lop (AFL++)

Where to Search for Vulnerabilities 

Testing Strategies for Model File Formats 

• Build with sanitizers enabled
• Modify legitimate model files
• Focus on size fields in headers
• Test extreme values (0, MAX_UINT64)
• Use AFL++ with structure-aware fuzzing

Common Security Pitfalls in AI Model Loaders 

• Integer overflows in allocation
• Unchecked array access
• Blind trust in header values
• Missing size validation
• Buffer writes without bounds checking

Keras Lambda Layers: How Code Execution Risks Lurk in AI Models

Understanding the Vulnerability

Keras models can contain custom code in Lambda layers. Loading these models runs the code in these layers. While not always malicious, a Lambda layer can execute arbitrary Python code, making them powerful attack vectors if a model file is tampered with.

How the Vulnerability Works

• Keras Lambda layers let you define custom operations
• When the model is loaded (using model = tf.keras.models.load_model('model.h5'), for example), the Lambda layer code is run
• An attacker can bake malicious code into a Lambda layer, achieving remote code execution upon loading the model

Example Setup:

Key Takeaway

If a ML tool loads Keras models with Lambda layers from untrusted sources, treat it like loading a script file. Look for patterns in code or docs: "We support loading arbitrary Keras models." That can be a red flag if no sandboxing is mentioned.

Getting Started with Detection

• Check for where models are loaded in a project. Are Lambda layers allowed?
• If yes, try replacing a legitimate model file with a malicious one that uses a Lambda to run system commands

Finding Similar Vulnerabilities

Understand Serialization

• How does the format store complex objects?
• Does it support storing functions or custom operations?
• Look for keywords like "custom ops", "extensions", "callbacks"

Check Documentation Flags

• Watch for terms like "arbitrary", "custom", "dynamic"
• Look for security warnings or sandboxing recommendations
• Note any "unsafe" load options in APIs

Code Analysis Strategy

Test Cases

ONNX Custom Operators: Hidden Control Flow Vulnerabilities

Understanding the Vulnerability

ONNX models contain a graph of operators that define their computation. The format supports custom operators and complex control flow that can be manipulated. When a model includes custom ops or intricate branching logic, loading it could lead to unexpected computation paths or even code execution.

How the Vulnerability Works

• ONNX supports registering custom operator implementations
• Models can include complex control flow (If/Loop nodes) with conditional execution
• Custom ops can be implemented in native code (C++/CUDA) for "performance"
• When loading models with custom ops, the code for these ops gets executed

Example Setup

Key Takeaway

When ML tools support loading ONNX models with custom operators or complex control flow, treat it like loading native code. Look for phrases like "full ONNX compatibility" or "custom operator support" in documentation.

Getting Started with Detection

• Search codebases for ONNX model loading and custom op registration
• Check if control flow operators (If/Loop) are fully supported
• Try adding custom ops or complex branching to legitimate models
• Look for ways to trigger different execution paths

Finding Similar Vulnerabilities

Check Format Capabilities

• Does it support custom implementations?
• Are there native code interfaces?
• How is control flow handled?

Review Documentation For

• Custom operator APIs
• Native code integration
• Runtime extensibility
• Performance optimization hooks

Code Analysis

Test Cases

• Add custom ops to legitimate models
• Create complex control flow graphs
• Mix native and interpreted code
• Probe extension mechanisms

Starting Your Hunt

The ML ecosystem is ripe for bug hunters who can combine basic reverse-engineering, code reading, and a pinch of creativity. Start small: pick an ML library on huntr that you're familiar with, review how it loads models, and ask yourself: "What happens if the file is malformed or malicious?"

By applying these techniques and continuously learning, you can carve out a niche in the ML vulnerability research space—an area that's fast-growing and in critical need of skilled security researchers. Happy hunting!