CVE-2021-37679 (Python, published 2021-08-12T23:15Z)
Description: value: TensorFlow is an end-to-end open source platform for machine learning. In affected versions it is possible to nest a `tf.map_fn` within another `tf.map_fn` call. However, if the input tensor is a `RaggedTensor` and there is no function signature provided, code assumes the output is a fully specified tensor and fills output buffer with uninitialized contents from the heap. The `t` and `z` outputs should be identical, however this is not the case. The last row of `t` contains data from the heap which can be used to leak other memory information. The bug lies in the conversion from a `Variant` tensor to a `RaggedTensor`. The [implementation](https://github.com/tensorflow/tensorflow/blob/460e000de3a83278fb00b61a16d161b1964f15f4/tensorflow/core/kernels/ragged_tensor_from_variant_op.cc#L177-L190) does not check that all inner shapes match and this results in the additional dimensions. The same implementation can result in data loss, if input tensor is tweaked. We have patched the issue in GitHub commit 4e2565483d0ffcadc719bd44893fb7f609bb5f12. The fix will be included in TensorFlow 2.6.0. We will also cherrypick this commit on TensorFlow 2.5.1, TensorFlow 2.4.3, and TensorFlow 2.3.4, as these are also affected and still in supported range.
Severity: MEDIUM (CVSS v3 base score: 7.8)
Weakness: CWE-681 \u2014 Incorrect Conversion between Numeric Types
When converting from one data type to another, such as long to integer, data can be omitted or translated in a way that produces unexpected values. If the resulting values are used in a sensitive context, then dangerous behaviors may occur.
Repository: https://github.com/tensorflow/tensorflow
Files changed:
tensorflow/core/kernels/ragged_tensor_from_variant_op.cc
tensorflow/python/ops/ragged/ragged_map_fn_op_test.py
Commit message: Fix bug that could cause map_fn to produce incorrect results (rather than an error)
when mapping over a ragged tensor with an inappropriate fn_output_signature.  (Note: there are cases where the default value for fn_output_signature is not appropriate, so the user needs to explicitly specify the correct output signature.)

PiperOrigin-RevId: 387606546
Change-Id: Ib4ea27b9634e6ab413f211cfe809a69a90f0e2cd


Vulnerable code:

Fixed code:

  TensorShape expected_value_shape = component_values_shape;
  expected_value_shape.RemoveDim(0);

    // Check that the flat_values tensor shape is compatible.
    TensorShape value_shape = ragged_components[i].values().shape();
    value_shape.RemoveDim(0);
    if (value_shape != expected_value_shape) {
      return errors::InvalidArgument(
          "All flat_values must have compatible shapes.  Shape at index 0: ",
          expected_value_shape, ".  Shape at index ", i, ": ", value_shape,
          ".  If you are using tf.map_fn, then you may need to specify an "
          "explicit fn_output_signature with appropriate ragged_rank, and/or "
          "convert output tensors to RaggedTensors.");
    }

from tensorflow.python.framework import errors
from tensorflow.python.ops import map_fn as map_fn_lib
  def testRaggedMapWithIncorrectFnOutputSignature(self):
    x = ragged_factory_ops.constant([[1, 2, 3, 4], [1]])
    with self.assertRaisesRegex(errors.InvalidArgumentError,
                                'All flat_values must have compatible shapes'):
      y = map_fn_lib.map_fn(lambda r: map_fn_lib.map_fn(lambda y: r, r), x)
      self.evaluate(y)

  def testNestedRaggedMapWithFnOutputSignature(self):
    ragged1d = ragged_tensor.RaggedTensorSpec([None], dtypes.int32)
    ragged2d = ragged_tensor.RaggedTensorSpec([None, None], dtypes.int32)

    x = ragged_factory_ops.constant([[1, 2, 3, 4], [1]])
    # pylint: disable=g-long-lambda
    y = map_fn_lib.map_fn(
        lambda r: map_fn_lib.map_fn(
            lambda y: r, r, fn_output_signature=ragged1d),
        x,
        fn_output_signature=ragged2d)
    expected = [[[1, 2, 3, 4], [1, 2, 3, 4], [1, 2, 3, 4], [1, 2, 3, 4]], [[1]]]
    self.assertAllEqual(y, expected)