CVE-2021-29514: TensorFlow GPU - Heap Buffer Overflow in RaggedBincount

Background

TensorFlow is an end-to-end open source platform for machine learning. It provides a comprehensive ecosystem of tools, libraries, and resources for developing and deploying machine learning models. TensorFlow GPU is a variant of TensorFlow that leverages GPU acceleration for faster computation. It is widely used in various domains for training and inference of deep learning models.

Vulnerability Detail

The vulnerability (CVE-2021-29514) exists in the RaggedBincount operation of TensorFlow GPU. When the splits argument of RaggedBincount does not specify a valid SparseTensor, it can trigger a heap buffer overflow. This leads to a read operation outside the bounds of the splits tensor buffer. The issue arises due to incorrect handling of the batch_idx variable. An attacker can exploit this vulnerability by setting splits(0) to an arbitrary value, causing the while loop to be skipped and resulting in a write operation to out(-1, bin), which is before the allocated buffer for the output tensor. This can lead to memory corruption and potential code execution.

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

Workaround

As a workaround, users can avoid using the RaggedBincount operation or ensure that the splits argument specifies a valid SparseTensor to prevent triggering the vulnerability. Here is an example of how to use a valid SparseTensor with RaggedBincount:

import tensorflow as tf

splits = tf.sparse.SparseTensor(indices=[[0, 0], [1, 0], [2, 0]], values=[1, 1, 1], dense_shape=[3, 1])
result = tf.RaggedTensor.from_sparse(splits)
bincount = tf.RaggedTensor.from_value_rowids(result, [0, 0, 0])
output = tf.RaggedTensor.from_row_lengths(bincount, [3])

Conclusion

The vulnerability (CVE-2021-29514) in TensorFlow GPU's RaggedBincount operation can be exploited to trigger a heap buffer overflow, leading to memory corruption and potential code execution. Users should update their TensorFlow GPU installation to one of the fixed versions (2.5.0, 2.4.2, or 2.3.3) to mitigate the risk. It is recommended to follow the provided workaround or apply the fix as soon as possible. For further assistance or support in vulnerability remediation, please contact DevNack.com.

See every vulnerability recorded for tensorflow-gpu →

FAQ