Honor safe_mode on the legacy/.h5 path and restrict Lambda function-gadget module resolution

[adilburaksen]

Summary

load_model(..., safe_mode=True) (the default) is documented to "disallow unsafe lambda deserialization", but two paths do not honor it. This brings tf-keras in line with the safe_mode hardening that landed in Keras 3 (which had the same gaps).

1. Arbitrary module/gadget resolution under safe_mode

For a Lambda layer whose function config is function_type=="function", _retrieve_class_or_fn (serialization_lib.py) imports the module named in the config and fetches a symbol from it with no restriction — so a config naming os / subprocess / builtins resolves an arbitrary callable (e.g. os.system) and installs it as the layer's function, even under safe_mode=True. Only the function_type=="lambda" path was guarded.

Fix: under safe mode, restrict the module-import fall-through to TF-Keras / TensorFlow modules (keras, tf_keras, tensorflow, numpy); legitimate built-in symbols are still resolved by the existing fast paths above. Anything else raises with the standard safe_mode message and the safe_mode=False escape hatch.

2. Legacy (.h5 / SavedModel) load silently ignores safe_mode

load_model(path, safe_mode=True) dispatches .h5/.hdf5 to the legacy loader without passing or scoping safe_mode (saving_api.py). The legacy path therefore never enters a SafeModeScope, in_safe_mode() defaults to falsy, and unsafe Lambda bytecode loads even when safe_mode=True.

Fix: wrap the legacy load in SafeModeScope(safe_mode) so the existing unsafe-lambda guard fires on the .h5 path too (matching the .keras path's behavior).

Tests

Added test_function_gadget_blocked_in_safe_mode (gadget from a non-Keras module is rejected under safe mode, allowed under safe_mode=False). Verified locally on tf-keras==2.21.0: with these changes, loading a crafted .h5 Lambda model and resolving an os.system gadget both raise under the default safe_mode=True (previously both succeeded).

Compatibility

safe_mode=False continues to allow all of the above (the documented opt-out). Legitimate Keras/TensorFlow Lambda functions resolve unchanged.

[gemini-code-assist]

Code Review

This pull request ensures that safe_mode is honored in legacy model loading paths by wrapping legacy load_model calls in a SafeModeScope. Additionally, it restricts module resolution under safe mode to allowed TF-Keras, TensorFlow, and NumPy modules to prevent arbitrary code execution vectors. A corresponding unit test has been added to verify this behavior. There are no review comments, and I have no feedback to provide.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[adilburaksen]

@hertschuh could you take a look when you get a chance? On the legacy .h5 path, safe_mode isn't honored and Lambda layers can resolve arbitrary modules as function gadgets, which reopens the deserialization concern safe_mode is meant to close. The change honors safe_mode there and restricts the module resolution. Glad to adjust the allowlist if there are legit modules I've missed.