29.3.3. Az ONNX formátum sebezhetőségei

Miért más az ONNX, mint a pickle vagy a SavedModel?

Gondolj az ONNX-ra úgy, mint a mesterséges intelligencia modellek „PDF-jére”. Míg a PyTorch .pt vagy a TensorFlow SavedModel formátumai egy adott ökoszisztémához kötődnek (mint a .docx a Wordhöz), az ONNX egy univerzális, keretrendszer-független leíró nyelv. Nem Python-objektumokat szerializál, mint a pickle, és nem is egy komplex könyvtárstruktúra, mint a SavedModel.

Az ONNX lényegében egy számítási gráfot definiál, amely operátorokból (pl. Conv, Relu, MatMul) és a köztük lévő adatfolyamot leíró tenzorokból áll. Ez a deklaratív jelleg elméletben biztonságosabbnak tűnik, hiszen „csak” matematikai műveleteket ír le. A probléma a részletekben, pontosabban az operátorok implementációjában rejlik.

Hol a sebezhetőség? Az egyéni operátorok (Custom Ops)

Az ONNX specifikáció lehetővé teszi az úgynevezett „custom operator”-ok használatát. Ezek olyan, a standard készletben nem szereplő műveletek, amelyeket a fejlesztők saját maguk implementálhatnak a specifikus igényeikhez. A támadási vektor itt nyílik meg: mi történik, ha egy támadó egy olyan egyéni operátort hoz létre, amely nem matematikai műveletet, hanem rendszerszintű parancsot hajt végre?

Az ONNX futtatókörnyezetek (mint az onnxruntime) feladata, hogy a gráfban definiált operátorokat leképezzék a tényleges, végrehajtható kódra. Ha a futtatókörnyezet nincs megfelelően „homokozóba” (sandbox) zárva, vagy ha naivan megbízik a modellben definiált operátorokban, egy rosszindulatúan preparált .onnx fájl tetszőleges kódfuttatást (Arbitrary Code Execution – ACE) érhet el a gyanútlan áldozat rendszerén.

Példa: A parancsfuttató operátor

Tételezzük fel, hogy egy támadó létrehoz egy látszólag ártalmatlan modellt. A modell belsejébe azonban elrejt egy egyéni operátort, mondjuk ShellExecute néven. A támadás két lépésből áll:

1. A rosszindulatú ONNX modell létrehozása

A támadó egy Python szkript segítségével legenerálja a mérgezett modellt. Nem kell hozzá komplex neurális háló, elég egyetlen csomópont.

# támadó oldali kód
import onnx
from onnx import helper, TensorProto

# Létrehozunk egy csomópontot (node), amely egy nem létező,
# de beszédes nevű egyéni operátort használ.
# A 'command' attribútumban adjuk meg a futtatandó parancsot.
node = helper.make_node(
 'ShellExecute', # Operátor neve, amit a runtime majd keres
 inputs=['input'], # Bemeneti tenzor neve
 outputs=['output'], # Kimeneti tenzor neve
 domain='com.redteam.ops', # Egyéni domain a névütközések elkerülésére
 command='touch /tmp/pwned_by_onnx' # A VÉGREHAJTANDÓ PAYLOAD
)

# Definiáljuk a gráf be- és kimeneteit
graph = helper.make_graph(
 [node],
 'malicious-graph',
 [helper.make_tensor_value_info('input', TensorProto.FLOAT, [1])],
 [helper.make_tensor_value_info('output', TensorProto.FLOAT, [1])]
)

# Létrehozzuk és elmentjük a modellt
model = helper.make_model(graph)
onnx.save(model, 'malicious_model.onnx')

2. Az áldozat végrehajtja a modellt

Az áldozat letölti a malicious_model.onnx fájlt egy megbízhatatlan forrásból (pl. egy kevésbé ismert modell-zoológiai kertből). A kódja, amivel futtatni próbálja, teljesen standard és ártalmatlannak tűnik. A sebezhetőség kihasználásához a támadónak rá kell vennie az áldozatot, hogy egy olyan módosított vagy sebezhető onnxruntime-ot használjon, amely implementálja a ShellExecute operátort. Ez társadalmi mérnökösködéssel vagy egy kompromittált futtatókörnyezet terjesztésével érhető el.

# áldozat oldali kód (feltételezve egy kompromittált runtime-ot)
import onnxruntime as ort
import numpy as np

try:
 # A gyanútlan felhasználó betölti a modellt.
 # A háttérben a runtime felismeri a 'ShellExecute' op-t és előkészíti.
 session = ort.InferenceSession('malicious_model.onnx')

 # Bemeneti adat létrehozása
 input_data = np.array([1.0], dtype=np.float32)

 # A modell futtatása. EZ A PONT, AHOL A PAYLOAD VÉGREHAJTÓDIK!
 # A 'session.run' hívás aktiválja a ShellExecute operátort.
 session.run(None, {'input': input_data})

 print("A modell lefutott... vagy mégsem?")
 # Ekkor a /tmp/pwned_by_onnx fájl már létrejött a rendszeren.
except Exception as e:
 print(f"Hiba történt: {e}")
 # Valószínűleg a runtime panaszkodni fog, ha nem ismeri az operátort.
 # Egy kifinomultabb támadás ezt is elrejtheti.

Védekezési stratégiák

Az ONNX-alapú támadások elleni védekezés a bizalom és az ellenőrzés elvén alapul. Mivel a formátum maga csak egy leírás, a felelősség a futtatókörnyezetre és a felhasználóra hárul.

Összefoglalva, az ONNX interoperabilitása kétélű fegyver. Red Teamerként ezt a felületet kell vizsgálnunk: nem a modellt magát, hanem azt, ahogyan a célrendszer értelmezi és életre kelti azt. A védekező oldalon pedig a „soha ne bízz a bemenetben” elvét kell kiterjeszteni a gépi tanulási modellekre is.