TensorFlow AI models may be at risk of supply chain attacks due to a flaw in the Keras API that enables execution of potentially unsafe code.
Keras is an API for neural networks, which is written in Python and provides a high-level interface for deep learning software libraries like TensorFlow and Theano.
A vulnerability tracked as CVE-2024-3660 affects Keras versions prior to 2.13 and was disclosed by the CERT Coordination Center last Tuesday. The flaw lies in the handling of Lambda Layers, a type of AI “building block” that enables developers to add arbitrary Python code to a model as an anonymous lambda function.
In earlier Keras versions, code included in Lambda Layers could be deserialized and executed without any checks, meaning an attacker could potentially distribute a trojanized version of a popular model that includes malicious Lambda Layers and execute code on the system of anyone who loads the model.
“This is just another in a long line of model injection vulnerabilities dating back more than a decade, including previous command injections in Keras models,” Dan McInerney, lead AI threat researcher at Protect AI, told SC Media in an email.
Keras 2.13 and later versions include a “safe_mode” parameter that is set to “True” by default and prevents the deserialization of unsafe Lambda Layers that may trigger arbitrary code execution. However, this check is only performed for models serialized in the Keras version 3 format (file extension .keras), meaning Keras models in older formats may still pose a risk.
The vulnerability poses a potential supply chain risk for developers working with TensorFlow models in Keras. A developer could unknowingly incorporate a third-party model with a malicious Lambda Layer in their own application or build their own model on a base model that includes the malicious code.
Model users and creators are urged to upgrade Keras to at least version 2.13 and ensure the safe_mode parameter is set to “True” to avoid arbitrary code execution from Lambda Layers. Models should also be saved and loaded in the Keras version 3 serialization format.
“Model users should only use models developed and distributed by trusted sources, and should always verify the behavior of models before deployment. They should follow the same development and deployment best practices to applications that integrate ML models as they would to any application incorporating any third party component,” the CERT researchers wrote.
Open-source software hosting platforms like Hugging Face, GitHub, npm and PyPI are popular targets for supply chain attacks due to the extent to which modern software depends on open-source third-party code. With the boom in AI development over the last couple years, supply chain threats focused on compromising AI models are likely to increase.
“The risks are compounded by the fact that unsafe model formats such as pickle were simply accepted by the machine learning community as the default model format for many years and the massive rise in the usage of third party models downloaded from online repositories such as Hugging Face,” McInerney said.
Indeed, earlier this month, malicious models in the insecure pickle format were found to be circulating on the Hugging Face platform.
“There are useful open source tools such as ModelScan that can detect malicious models, but this is unlikely the end of novel ways to force models to execute malicious code without the end user even being aware,” McInerney concluded.