CVE-2021-37679: Heap OOB in nested `tf.map_fn` with `RaggedTensor`s in TensorFlow

Background

TensorFlow is an open-source machine learning framework developed by Google. It provides a comprehensive ecosystem of tools, libraries, and resources for building and deploying machine learning models. TensorFlow supports various data structures, including RaggedTensor, which allows for handling irregular and nested data. However, a vulnerability has been discovered in the nested usage of tf.map_fn with RaggedTensors, leading to a heap out-of-bounds (OOB) issue.

Vulnerability Detail

The vulnerability arises when nesting a tf.map_fn within another tf.map_fn call with a RaggedTensor as input. If no function signature is provided, the code assumes the output is a fully specified tensor and fills the output buffer with uninitialized contents from the heap. This can result in data leakage and potential memory information disclosure. The issue originates from the conversion from a Variant tensor to a RaggedTensor, where the implementation fails to check that all inner shapes match, leading to additional dimensions. The same vulnerability can also cause data loss if the input tensor is modified.

pip install --upgrade "tensorflow==2.21.0"

Workaround

As a workaround, you can provide a function signature explicitly when using nested tf.map_fn with RaggedTensors. This ensures that the output tensor is correctly specified and avoids the uninitialized buffer issue. Here's an example:

import tensorflow as tf

x = tf.ragged.constant([[1,2,3], [4,5], [6]])

# Provide function signature
@tf.function
def map_fn_inner(y):
    return y

@tf.function
def map_fn_outer(r):
    return tf.map_fn(map_fn_inner, r)

t = tf.map_fn(map_fn_outer, x)

# t and z outputs should be identical
z = tf.ragged.constant([[[1,2,3],[1,2,3],[1,2,3]],[[4,5],[4,5]],[[6]]])

By explicitly defining the function signature, the issue of uninitialized buffer contents can be avoided.

Conclusion

The vulnerability in nested tf.map_fn with RaggedTensors in TensorFlow can lead to data leakage and potential memory information disclosure. It is crucial to update to the patched versions of TensorFlow (2.6.0, 2.5.1, 2.4.3, or 2.3.4) to mitigate this issue. By following the recommended steps for updating your TensorFlow installation, you can ensure the security and stability of your machine learning applications. For further support on vulnerability remediation, please contact DevNack.com.

See every vulnerability recorded for tensorflow →

FAQ