TensorFlow Keras“安全模式”真的安全吗？绕过 safe_mode 缓解措施，实现任意代码执行

机器学习框架通常依赖序列化和反序列化机制来存储和加载模型，然而模型中不恰当的代码隔离和可执行组件可能会导致严重的安全风险。

TensorFlow 中的 Keras v3 ML 模型结构

对于基于 TensorFlow 的 Keras 模型，存在一个严重的反序列化漏洞，编号为CVE-2024-3660。攻击者可利用此漏洞通过构建恶意 Keras 模型文件注入并执行任意代码。该漏洞利用了 Lambda 层的反序列化，这些层存在于模型中，包含已编排的 Python 代码。因此，攻击者可以将恶意代码嵌入 Keras 模型文件中，从而在毫无戒心的受害者加载该模型时导致任意代码执行 (ACE)。

包含嵌入式 Python 字节码的 Keras Lambda 层示例

为了应对这一风险，Keras在 2.13 版本中引入了**安全模式参数来降低此风险。默认情况下， safe_mode设置为True ，以防止可能执行任意代码的 Lambda 层被反序列化。此修复解决了特定Lambda**层的解组问题，如果用户尝试在安全模式下加载包含已编组 Lambda 层的模型，则会引发异常。

 if (isinstance(fn_config, dict)and "class_name" in fn_configand fn_config["class_name"] == "__lambda__"):cls._raise_for_lambda_deserialization("function", safe_mode)inner_config = fn_config["config"]fn = python_utils.func_load(inner_config["code"],defaults=inner_config["defaults"],closure=inner_config["closure"],)

Lambda 层实现的安全模式

这次修复是一个很好的开端，但问题仍然存在：即使在安全模式下，是否仍然可以创建一个利用反序列化来执行恶意负载的模型？为了解决这个问题，我们应该首先研究 Keras 模型反序列化的实现方式。

Keras反序列化机制

Keras 中的模型加载过程始于函数saving_lib.load_model，该函数识别模型源（文件、目录或 Hugging Face 存储库），然后调用相应的加载器。无论源是什么，每个方法都会从提取自模型的config.json文件中反序列化模型，并调用**_model_from_config**。该函数反过来调用deserialize_keras_object，反序列化逻辑就是在这里执行的。

在此过程中，该方法以递归方式解释字典结构，将其元素转换为相应的 Python 对象。从攻击者的角度来看，有两个关键组件特别有趣。第一个是**_retrieve_class_or_fn**，它是一个基于 JSON 配置实例化任意对象的函数。第二个涉及 Keras Functional API，它为开发人员提供了一种强大而灵活的方式来构建复杂的图形拓扑，同时提供了在加载时执行反序列化对象的原语。**让我们一起构建一个函数模型，演示在加载时执行任意代码。**以下清单包含具有必填字段的函数模型的转储：

 {    "class_name": "Functional","name": "functional","config": {