CVE-2023-25659: Out-of-bounds Read (CWE-125)
value: TensorFlow is an open source platform for machine learning. Prior to versions 2.12.0 and 2.11.1, if the parameter `indices` for `DynamicStitch` does not match the shape of the parameter `data`, it can trigger an stack OOB read. A fix is included in TensorFlow version 2.12.0 and version 2.11.1.

CWE description: The product reads data past the end, or before the beginning, of the intended buffer.
CVSS v3 base score: 7.5
Published: 2023-03-25T00:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: C++
Files changed: tensorflow/compiler/tf2xla/kernels/dynamic_stitch_op.cc, tensorflow/core/kernels/dynamic_stitch_op.cc, tensorflow/python/kernel_tests/data_structures/dynamic_stitch_op_test.py
Commit date: 2023-02-01 14:13:23 -0800
Commit message: Add out of bounds array check to dynamic_stitch_op.

PiperOrigin-RevId: 506418249

Diff: diff --git a/tensorflow/compiler/tf2xla/kernels/dynamic_stitch_op.cc b/tensorflow/compiler/tf2xla/kernels/dynamic_stitch_op.cc
index 7566cc79742d0..cd03b617158e7 100644
--- a/tensorflow/compiler/tf2xla/kernels/dynamic_stitch_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/dynamic_stitch_op.cc
@@ -144,10 +144,14 @@ class DynamicStitchOp : public XlaOpKernel {
     std::vector<bool> src_index_used(number_of_indices);
     int index_used_count = 0;
     for (int input_num = 0; input_num < indices.size(); input_num++) {
       for (int i = 0; i < indices[input_num].shape().dimensions(0); ++i) {
         int index = indices[input_num].Get<int>({i});
+        OP_REQUIRES(
+            ctx, index >= 0,
+            errors::InvalidArgument("indices[", index, "] is out of range"));
+
         src_input_vector[index] = input_num;
         src_slice_vector[index] = i;
         if (!src_index_used[index]) {
           src_index_used[index] = true;
           ++index_used_count;
diff --git a/tensorflow/core/kernels/dynamic_stitch_op.cc b/tensorflow/core/kernels/dynamic_stitch_op.cc
index b47180ef831a9..23fc151b4ebd4 100644
--- a/tensorflow/core/kernels/dynamic_stitch_op.cc
+++ b/tensorflow/core/kernels/dynamic_stitch_op.cc
@@ -95,10 +95,21 @@ class DynamicStitchOpImplBase : public OpKernel {
       }
     }
 
     *first_dim_size = max_index + 1;
 
+    for (const Tensor& indices : *indices_inputs) {
+      auto indices_vec = indices.flat<int32>();
+
+      for (int i = 0; i < indices_vec.size(); i++) {
+        int32_t index = internal::SubtleMustCopy(indices_vec(i));
+        OP_REQUIRES(
+            c, FastBoundsCheck(index, *first_dim_size),
+            errors::InvalidArgument("indices[", i, "] is out of range"));
+      }
+    }
+
     // Validate that data[i].shape = indices[i].shape + constant
     OP_REQUIRES_OK(c, c->input_list("data", data_inputs));
     const Tensor& data0 = (*data_inputs)[0];
     const Tensor& indices0 = (*indices_inputs)[0];
     for (int input_num = 0; input_num < indices_inputs->size(); input_num++) {
@@ -263,25 +274,19 @@ class DynamicStitchOpImplCPU : public DynamicStitchOpImplBase<T> {
         if (DataTypeCanUseMemcpy(DataTypeToEnum<T>::v())) {
           T* merged_base = merged_flat.data();
           const T* data_base = data_flat.data();
           for (int i = 0; i < indices_vec.size(); i++) {
             int32_t index = internal::SubtleMustCopy(indices_vec(i));
-            OP_REQUIRES(
-                c, FastBoundsCheck(index, first_dim_size),
-                errors::InvalidArgument("indices[", i, "] is out of range"));
             memcpy(merged_base + index * slice_size, data_base + i * slice_size,
                    slice_bytes);
           }
         } else {
           Eigen::DSizes<Eigen::DenseIndex, 2> sizes(1, slice_size);
           for (int i = 0; i < indices_vec.size(); i++) {
             // Copy slice data[i] to merged[indices[i]]
             Eigen::DSizes<Eigen::DenseIndex, 2> data_indices(i, 0);
             int32_t index = internal::SubtleMustCopy(indices_vec(i));
-            OP_REQUIRES(
-                c, FastBoundsCheck(index, first_dim_size),
-                errors::InvalidArgument("indices[", i, "] is out of range"));
             Eigen::DSizes<Eigen::DenseIndex, 2> merged_indices(index, 0);
             merged_flat.slice(merged_indices, sizes) =
                 data_flat.slice(data_indices, sizes);
           }
         }
diff --git a/tensorflow/python/kernel_tests/data_structures/dynamic_stitch_op_test.py b/tensorflow/python/kernel_tests/data_structures/dynamic_stitch_op_test.py
index a573cd8bf65b4..6e2ffb0fced9a 100644
--- a/tensorflow/python/kernel_tests/data_structures/dynamic_stitch_op_test.py
+++ b/tensorflow/python/kernel_tests/data_structures/dynamic_stitch_op_test.py
@@ -224,10 +224,23 @@ class DynamicStitchTestBase(object):
         constant_op.constant([10, 60, 20, 30])
     ]
     with self.assertRaises(ValueError):
       self.stitch_op(indices, data)
 
+  def testOutOfBoundsIndexRaisesInvalidArgument(self):
+    with self.assertRaisesRegex(errors.InvalidArgumentError, "out of range"):
+      indices = [[-1000], [405], [519], [758], [1015]]
+      data = [
+          [110.27793884277344],
+          [120.29475402832031],
+          [157.2418212890625],
+          [157.2626953125],
+          [188.45382690429688],
+      ]
+
+      self.evaluate(self.stitch_op(indices, data))
+
 
 class DynamicStitchTest(DynamicStitchTestBase, test.TestCase):
 
   def __init__(self, *test_case_args):
     test.TestCase.__init__(self, *test_case_args)

Vulnerable code:             OP_REQUIRES(
                c, FastBoundsCheck(index, first_dim_size),
                errors::InvalidArgument("indices[", i, "] is out of range"));
            OP_REQUIRES(
                c, FastBoundsCheck(index, first_dim_size),
                errors::InvalidArgument("indices[", i, "] is out of range"));
Fixed code:         OP_REQUIRES(
            ctx, index >= 0,
            errors::InvalidArgument("indices[", index, "] is out of range"));

    for (const Tensor& indices : *indices_inputs) {
      auto indices_vec = indices.flat<int32>();

      for (int i = 0; i < indices_vec.size(); i++) {
        int32_t index = internal::SubtleMustCopy(indices_vec(i));
        OP_REQUIRES(
            c, FastBoundsCheck(index, *first_dim_size),
            errors::InvalidArgument("indices[", i, "] is out of range"));
      }
    }

  def testOutOfBoundsIndexRaisesInvalidArgument(self):
    with self.assertRaisesRegex(errors.InvalidArgumentError, "out of range"):
      indices = [[-1000], [405], [519], [758], [1015]]
      data = [
          [110.27793884277344],
          [120.29475402832031],
          [157.2418212890625],
          [157.2626953125],
          [188.45382690429688],
      ]

      self.evaluate(self.stitch_op(indices, data))