Repository: tensorflow
Issue: Potential Remote Code Execution (RCE) Vulnerability in Custom Layers Handling
Labels: type:others,TF 2.18

### Issue type

Bug

### Have you reproduced the bug with TensorFlow Nightly?

Yes

### Source

source

### TensorFlow version

 2.18.0

### Custom code

Yes

### OS platform and distribution

_No response_

### Mobile device

Kali Linux 2024.1

### Python version

Python 3.11.9

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

CUDA: 11.2 

### GPU model and memory

NVIDIA GTX 1080

### Current behavior?

Currently, TensorFlow (primarily through Keras) allows users to create custom layers that can contain malicious code. When a model containing these custom layers is loaded or run, the malicious code in the layer can be executed without any restrictions or filtering. This opens up potential exploits because the executed code could potentially damage the system or steal sensitive data. For example, in this report, the saved model contains a MaliciousLayer layer that calls the os.system() system command to execute malicious code when the model is loaded, which could result in corruption or unauthorized data acquisition.



I hope TensorFlow can provide protection that limits or filters the execution of malicious code in registered custom layers. For example, by blocking malicious functions such as os.system() or prohibiting the use of code that can execute external commands when the model is loaded. In addition, the use of custom layers that contain unwanted external code should be protected by restrictions or sandboxing. This will ensure that the loaded model cannot inject malicious code into the user's system, thereby increasing the security of using TensorFlow models.

The main issue here is that TensorFlow (in this case Keras) allows users to define custom layers that can execute malicious code when the model is loaded and used. This is a potential security hole that can be exploited if malicious code is injected into the model.

The challenge is the ability to include custom layers that execute malicious code. TensorFlow/Keras does not filter or restrict the use of malicious code in registered custom layers, which allows exploits like this.

So, the vulnerability is in the implementation of custom layers in TensorFlow/Keras

Who can exploit the vulnerability?

An attacker or malicious individual who gains access to a TensorFlow model saved in the .keras format and has knowledge of how to create malicious custom layers can exploit this vulnerability. Typically, this would involve actors with an understanding of TensorFlow programming or machine learning who can inject harmful code into the model to be executed when the model is loaded.

What do they gain from it?

An attacker who successfully exploits this vulnerability can execute malicious code on the victim's system when the model is loaded. This allows them to steal sensitive personal data, such as login credentials, financial data, or other private information. Additionally, they can leverage remote code execution to gain control of the system or cause further dama

### Standalone code to reproduce the issue

```shell
1.Create a Python file, for example kontol13.py, containing the following code to create a model with a malicious layer (MaliciousLayer):


import tensorflow as tf
from tensorflow.keras.layers import Layer
import os

# Defines a custom layer containing malicious code
@tf.keras.utils.register_keras_serializable()
class MaliciousLayer(Layer):
    def __init__(self, **kwargs):
        super(MaliciousLayer, self).__init__(**kwargs)

    def call(self, inputs):
        # Code that is executed when the model runs
        os.system('echo "Malicious Code Executed!"')
        return inputs

# Creating a model with dangerous layers
model = tf.keras.Sequential([
    tf.keras.layers.InputLayer(input_shape=(10,)),
    tf.keras.layers.Dense(10, activation='relu'),
    MaliciousLayer(),  # Layer berbahaya
    tf.keras.layers.Dense(1)
])

# Save the model in Keras format (.keras)
model.save('kontol13.keras')

print("Malicious model created and saved as 'kontol13.keras'.")
```
2.Run the above script to create and save the model containing the malicious layers

I also saved the result in google colab:

https://colab.research.google.com/drive/1IwpwNGOeTPUgYu4Y4bAl7WufOXBvOQy9?authuser=1#scrollTo=EYbfvtj6O6wU&line=1&uniqifier=1
```


### Relevant log output

```shell
I run this code on kali linux this is the output

2024-12-04 17:45:53.327622: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:477] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
WARNING: All log messages before absl::InitializeLog() are called are written to STDERR
E0000 00:00:1733309153.640400 154 cuda_dnn.cc:8310] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
E0000 00:00:1733309153.730398 154 cuda_blas.cc:1418] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
2024-12-04 17:45:54.592724: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: AVX2 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
/home/funsociety/.local/lib/python3.11/site-packages/keras/src/layers/core/input_layer.py:27: UserWarning: Argument `input_shape` is deprecated. Use `shape` instead.
 warnings.warn(
2024-12-04 17:46:03.552814: E external/local_xla/xla/stream_executor/cuda/cuda_driver.cc:152] failed call to cuInit: INTERNAL: CUDA error: Failed call to cuInit: UNKNOWN ERROR (303)
Malicious Code Executed!
```