Lambda Layers Code Execution Flaw Leads To Supply Chain On AI/ML Applications

A new supply-chain vulnerability has been identified in the Lambda Layers of third-party TensorFlow-based Keras models. This vulnerability could allow threat actors to inject arbitrary code into any AI/ML application.

Any Lambda Layers that were built before version Keras 2.13 are susceptible to a supply chain attack.

A threat actor can create and distribute a trojanized popular model among AI/ML developers.

If the attack succeeds, the threat actor can execute untrusted arbitrary code on the vulnerable environments with the same privileges as the running application. 

Free Webinar | Mastering WAAP/WAF ROI Analysis | Book Your Spot

Lambda Layers Code Execution Flaw

The Keras framework provides a high-level interface for TensorFlow and offers several features for designing, training, validating, and packaging ML models.

The building blocks used for building neural networks are called Layers. Keras provides an API for these layers.

There are many layer types available in Keras, one of which is the Lambda Layer type. This type allows a developer to add arbitrary code to a model as a lambda function.

This can be done using the model.save() or save_model () method as described in the Keras Documentation.

Additionally, the Keras 2 documentation describes an additional mechanism for disallowing the loading of a native version 3 Keras model, which has the option to add a lambda layer during safe_mode.

This safe_mode is enabled by default which is responsible for allowing/disallowing unsafe lambda deserialization and has the potential to trigger arbitrary code execution.

However, in Keras versions 2.13 and later, there is an exception that is raised in a program when there is an attempt to load a model with Lambda Layers stored in version 3 of the format, there will be an exception raised.

This particular mechanism was absent in versions prior to 2.13, making the earlier versions deserialize untrusted code.

According to the TensorFlow documentation, a statement is provided as a warning to developers, which is probably not fully understood by new AI/ML community members.

The statement says, “Since models are practically programs that TensorFlow executes, using untrusted models or graphs is equivalent to running untrusted code”.

However, the Kensar Framework documentation for the load_model function states under the “Arguments” section about an option called safe_mode, which is “Boolean, whether to disallow unsafe lambda deserialization.

When safe_mode=False, loading an object has the potential to trigger arbitrary code execution. This argument is only applicable to the Keras v3 model format. Defaults to True.”

This code injection vulnerability when packaging data together with code is not new; there have been several instances in the past, including the Pickle mechanism in the standard Python library, allowing the serialization of code in line with this data.

To prevent these kinds of supply chain attacks, it is recommended that developers upgrade to the latest Keras versions, 2.13 or later, and ensure no valuable assets are in the scope of the running application.

This will reduce the potential data exfiltration in case of pre-2.13 applications in a sandbox.

Looking to Safeguard Your Company from Advanced Cyber Threats? Deploy TrustNet to Your Radar ASAP.