CVE-2021-37679 (published 2021-08-12T23:15Z)

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 v2 base score: 4.6 (AV:L/AC:L/Au:N/C:P/I:P/A:P)
CVSS v3 base score: 7.8, HIGH (CVSS:3.1/AV:L/AC:L/PR:L/UI:N/S:U/C:H/I:H/A:H)

Weakness information:
cwe_id: CWE-681
cwe_name: Incorrect Conversion between Numeric Types
description: 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.
extended_description: 
            
            
            
         
         

Fix information:
repo_url: https://github.com/tensorflow/tensorflow
commit_details: repo_url: https://github.com/tensorflow/tensorflow
author: Edward Loper
author_date: 2021-07-29 09:50:01-07:00
msg: 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
repository: repo_url: https://github.com/tensorflow/tensorflow
repo_name: tensorflow/tensorflow
description: An Open Source Machine Learning Framework for Everyone
repo_language: C++
owner: tensorflow
file_changes: filename: ragged_tensor_from_variant_op.cc
change_type: ModificationType.MODIFY
diff: @@ -174,7 +174,23 @@ Status NestedStackRaggedTensors(
   auto output_values_flat =
       output_ragged->mutable_values()->flat_outer_dims<VALUE_TYPE, 2>();
   int values_index = 0;
+
+  TensorShape expected_value_shape = component_values_shape;
+  expected_value_shape.RemoveDim(0);
+
   for (int i = 0; i < ragged_components.size(); i++) {
+    // 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.");
+    }
+
     auto component_values_flat =
         ragged_components[i].values().flat_outer_dims<VALUE_TYPE, 2>();
     int num_inner_elements = ragged_components[i].values().NumElements();

programming_language: C++
filename: ragged_map_fn_op_test.py
change_type: ModificationType.MODIFY
diff: @@ -21,9 +21,11 @@ from absl.testing import parameterized
 import numpy as np
 
 from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
 from tensorflow.python.framework import sparse_tensor
 from tensorflow.python.framework import test_util
 from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import map_fn as map_fn_lib
 from tensorflow.python.ops import math_ops as mo
 from tensorflow.python.ops import string_ops
 from tensorflow.python.ops.ragged import ragged_factory_ops
@@ -309,6 +311,27 @@ class RaggedMapOpTest(test_util.TensorFlowTestCase,
     )
     self.assertAllEqual(id_t2, [[0, 5], [0, 4]])
 
+  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)
+
 
 if __name__ == '__main__':
   googletest.main()

programming_language: Python