CVE-2021-29536: TensorFlow GPU - Heap Buffer Overflow in QuantizedReshape

Background

TensorFlow is an end-to-end open source platform for machine learning. It provides a comprehensive ecosystem of tools, libraries, and community resources that enable researchers and developers to build and deploy machine learning models. TensorFlow GPU is a variant of TensorFlow that leverages GPU hardware for accelerated computation. It is widely used in various industries for training and inference of machine learning models.

Vulnerability Detail

The vulnerability identified as CVE-2021-29536 affects the QuantizedReshape function in TensorFlow GPU. An attacker can exploit this vulnerability by passing invalid thresholds for the quantization process. The implementation assumes that the provided arguments are valid scalars and directly accesses their numeric values. However, if any of these tensors is empty, accessing the element at position 0 results in a heap buffer overflow. This can lead to arbitrary code execution or a denial-of-service condition. The vulnerability has been assigned a CVSS score of X.X, indicating its severity.

pip install --upgrade "tensorflow-gpu==2.12.0"

Workaround

As a temporary workaround, users can avoid passing empty tensors as arguments to the QuantizedReshape function. It is recommended to validate the input tensors before invoking the function to prevent the heap buffer overflow. Here is an example of how to perform the validation in Python:

import tensorflow as tf

def safe_quantized_reshape(input_tensor, shape_tensor):
    if input_tensor.size == 0 or shape_tensor.size == 0:
        raise ValueError('Input tensors must not be empty.')
    return tf.quantization.quantized_reshape(input_tensor, shape_tensor)

Conclusion

The vulnerability identified as CVE-2021-29536 in TensorFlow GPU's QuantizedReshape function can be exploited by passing invalid thresholds for quantization, leading to a heap buffer overflow. This can result in arbitrary code execution or a denial-of-service condition. Users are strongly advised to update their TensorFlow GPU installations to the fixed versions mentioned above. Additionally, implementing the recommended workaround of validating input tensors can provide temporary protection against this vulnerability. For further assistance or support in vulnerability remediation, please contact DevNack.com.

See every vulnerability recorded for tensorflow-gpu →

FAQ