FrankLin00/keras-mfv-poc-006-pickle-unsafe-load

[internal-id: KERAS-MFV-006] KerasSaveable.__reduce__ pickle path forces safe_mode=False, defeating Keras's documented safe-loading guarantees

Severity: CVSS 3.1 8.8 High baseline (CVSS:3.1/AV:N/AC:L/PR:N/UI:R/S:U/C:H/I:H/A:H). The UI:R reflects the typical post-load forward pass that triggers embedded Lambda invocation. In automated pipelines that immediately predict() after load (model registries, Dask/Ray workers, batch inference services), UI:N applies — yielding CVSS 9.8 Critical (CVSS:3.1/AV:N/AC:L/PR:N/UI:N/S:U/C:H/I:H/A:H).

CWE: CWE-502 (Deserialization of Untrusted Data) plus CWE-693 (Protection Mechanism Failure). CWE-502 covers the pickle entry point; CWE-693 covers the bypass of safe_mode — the documented defensive control that this code path explicitly disables.

Affected versions: Keras >= 3.0.0, <= 3.15.0 (HEAD 42b66280e, 2026-05-01). KerasSaveable.__reduce__ and _unpickle_model were introduced together with the unified saving stack and have shipped with safe_mode=False hardcoded since the first 3.x release. git show v3.0.0:keras/saving/keras_saveable.py and HEAD's keras/src/saving/keras_saveable.py both contain the identical sink.

Component:

• keras/src/saving/keras_saveable.py:21-23 — the sink: _load_model_from_fileobj(..., safe_mode=False).
• keras/src/saving/keras_saveable.py:25-38 — the __reduce__ hook that emits this sink as the pickle reconstructor.
• keras/src/saving/saving_lib.py:434 — _load_model_from_fileobj (the loader being invoked unsafely).
• keras/src/layers/core/lambda_layer.py:170-180 — one of the safe_mode-gated defenses bypassed (Lambda bytecode RCE). Other gated defenses also bypassed: keras/src/layers/preprocessing/index_lookup.py:443-454 (vocabulary file read), keras/src/utils/torch_utils.py:170-180, keras/src/export/tfsm_layer.py:158-180.

Reporter: Independent security research (NTU academic vulnerability research project).

Summary

Keras advertises safe_mode=True (the default of keras.saving.load_model, and the value KerasSaveable.from_config / Lambda.from_config / IndexLookup.set_vocabulary all read via serialization_lib.in_safe_mode()) as the gate that prevents an attacker-controlled artifact from running arbitrary code. Loading a .keras file from an untrusted source is treated as a soft-warned but bounded operation: Lambda bytecode is rejected, TextVectorization vocabulary file paths are rejected, custom torch modules are rejected, etc. The user opts into the dangerous behaviour by an explicit enable_unsafe_deserialization() call or safe_mode=False parameter.

That guarantee does not survive the pickle path. KerasSaveable.__reduce__ (keras/src/saving/keras_saveable.py:25) emits (_unpickle_model, (BytesIO(model_bytes),)) as the pickle reconstructor. _unpickle_model (line 17) calls saving_lib._load_model_from_fileobj(bytesio, custom_objects=None, compile=True, safe_mode=False) — safe_mode=False is hardcoded, with no override path for the caller. Whenever a Keras model is reconstructed through Python's pickle protocol (or any wrapper such as joblib.load, cloudpickle.loads, or stdlib multiprocessing queue traffic), every safe_mode-gated defense in the entire deserialization pipeline is silently disabled.

The comment at keras_saveable.py:20 ("pickle is not safe regardless of what you do") acknowledges pickle's general untrustworthiness but mischaracterizes the bug. The harm is not "pickle can run arbitrary code if you unpickle bytes you don't trust" (a well-known property). The harm is that the documented safe_mode invariant is broken: even pickle allowlists / pickle scanners that explicitly approve Keras-internal classes (reasonable: "loading a Keras model is bounded by safe_mode") are bypassed, because the dangerous payload is buried inside an opaque BytesIO blob whose contents only get classifier-deserialized after the unpickler hands control back to Keras.

Root cause

keras/src/saving/keras_saveable.py:17-23:

@classmethod
def _unpickle_model(cls, bytesio):
    import keras.src.saving.saving_lib as saving_lib

    # pickle is not safe regardless of what you do.
    return saving_lib._load_model_from_fileobj(
        bytesio, custom_objects=None, compile=True, safe_mode=False
    )

The literal safe_mode=False overrides the package-wide default safe_mode=True that all other public load entries (saving_api.load_model, saving_lib.load_model, legacy_h5_format.load_model_from_hdf5) inherit. There is no way for a caller of pickle.loads to inject safe_mode=True; the reducer's tuple form (_unpickle_model, (bytesio,)) hardwires the call's positional/keyword arguments.

Once safe_mode=False propagates into _load_model_from_fileobj and from there into deserialize_keras_object (saving_lib.py:439), every nested deserialization runs with in_safe_mode() == False:

• Lambda.from_config (lambda_layer.py:184) — __lambda__ bytecode marshalled-load is permitted.
• IndexLookup.set_vocabulary (index_lookup.py:444) — vocabulary path string is permitted to read arbitrary host files.
• TFSMLayer (tfsm_layer.py:168), TorchModuleWrapper (torch_utils.py:170) — both gates bypassed identically.

Exploitation

The shipped PoC produces a pickle file (/tmp/keras_evil.pkl) containing a benign-looking graph header (builtins.getattr → keras.src.models.functional.Functional → ...) plus a Lambda layer whose function is anonymous (so it serializes via the bytecode path at lambda_layer.py:152 rather than by-name) and whose body invokes os.system. Three victim scenarios are demonstrated:

• Scenario A — bare pickle.loads (run_a_pickle.py): the simplest reproduction; any process that unpickles the file gets RCE on the next forward pass. Models are routinely shipped between processes via multiprocessing.Queue, concurrent.futures worker handoff, Dask/Ray cross-node transport, Spark broadcast variables, or downloaded from a model registry as *.pkl.
• Scenario B — joblib.load (run_b_joblib.py): joblib is the de-facto scikit-learn / tabular-ML model persistence helper. It calls pickle.load internally, hits the same __reduce__ path. Demonstrates the bug is reachable through the dominant non-Keras loader in the Python ML stack.
• Scenario C — allowlist-restricted unpickler (run_c_restricted.py): the strongest scenario for impact. A defender writes the textbook pickle defense — pickle.Unpickler subclass with find_class restricted to a module allowlist — and reasonably allows keras.src.saving.keras_saveable/keras.src.models.functional/etc. because their threat model says "we trust Keras to load Keras models safely (it has safe_mode)". The allowlist does not save them: the only top-level reference is KerasSaveable._unpickle_model, the malicious payload lives inside a BytesIO blob the unpickler treats as opaque, and the embedded _load_model_from_fileobj(safe_mode=False) call inside _unpickle_model deserializes the dangerous Lambda layer freely.

The marker file /tmp/PWNED_DETAILS written by the Lambda body contains pwned_pid_<victim-pid>_at_<unix-ts> — confirming code ran inside the victim process, not the attacker process.

Impact

• Documented defense bypass. Keras's headline security property — that loading a model with safe_mode=True (the default) blocks code-execution primitives such as Lambda bytecode — is silently undone the moment the same bytes flow through pickle.loads instead of keras.saving.load_model. No opt-in, no flag, no warning.
• Dominates ProtectAI's own product surface. Pickle scanners that inspect opcode streams (picklescan, fickling, etc.) cannot see the marshalled Lambda bytecode that lives inside a BytesIO blob — to the scanner the pickle is just getattr + a Keras classmethod + a binary buffer. The attack defeats the same allowlist defense ProtectAI's customers deploy.
• Realistic delivery channels are common, not pathological. Distributed ML frameworks (Dask, Ray, Spark, joblib, multiprocessing, concurrent.futures, cloudpickle) all serialize models via pickle. Any of these consumed from an untrusted upstream — colleague's model.pkl, model-registry artifact, message queue from another tenant, broadcast variable in a multi-tenant Spark cluster — yields RCE.
• Chains with all other safe_mode-gated bugs. Since this bypass disables every safe_mode check globally inside the load pipeline, every other safe_mode-protected primitive — IndexLookup arbitrary file read, TFSMLayer load, TorchModuleWrapper load — also fires, on the same pickle.loads. A single artifact can deliver RCE + arbitrary file read in a single shot.

Affected Versions

• First vulnerable release: Keras 3.0.0 (2023-11-27). The __reduce__ reducer and _unpickle_model classmethod were introduced as part of the initial 3.x saving overhaul and have always carried safe_mode=False.
• Last vulnerable release: HEAD 42b66280e (__version__ = "3.15.0", keras/src/version.py:4).
• The text of keras_saveable.py is unchanged from the file's first commit through HEAD aside from import-line whitespace.

Suggested fix

There is no scenario in which _unpickle_model legitimately needs to pass safe_mode=False. The reducer should preserve the same default-safe contract that keras.saving.load_model advertises:

@classmethod
def _unpickle_model(cls, bytesio):
    import keras.src.saving.saving_lib as saving_lib
    # Default-safe: respect the same invariant load_model() does.
    return saving_lib._load_model_from_fileobj(
        bytesio, custom_objects=None, compile=True, safe_mode=True
    )

Users who genuinely need to round-trip Lambda-bytecode models through pickle (rare, and already requires enable_unsafe_deserialization() to create the bytecode form in the first place) can perform the unsafe load explicitly through keras.saving.load_model(..., safe_mode=False) after manually unpacking the BytesIO. They lose __reduce__ ergonomics but gain the documented security guarantee.

A complementary mitigation, cheap and orthogonal: emit a warnings.warn on every _unpickle_model invocation pointing at this advisory and the recommendation to migrate to keras.saving.save_model / keras.saving.load_model for cross-process model transport.

Reproduction

The PoC repository contains a self-contained build + three victim scenarios:

poc/
├── build.py             # Phase 1 — build malicious /tmp/keras_evil.pkl
├── run_a_pickle.py      # Scenario A — bare pickle.loads
├── run_b_joblib.py      # Scenario B — joblib.load
├── run_c_restricted.py  # Scenario C — allowlist-restricted Unpickler
└── result.md            # captured outputs of all three scenarios

Run order:

python build.py
python run_a_pickle.py        # → /tmp/PWNED_A
python run_b_joblib.py        # → /tmp/PWNED_B
python run_c_restricted.py    # → /tmp/PWNED_C

Marker file existence after each run is the evidence of RCE inside the victim process. Inspect /tmp/PWNED_DETAILS for the victim PID and Unix timestamp written from inside the Lambda body.