CVE-2021-29609: Out-of-bounds Write (CWE-787)
value: TensorFlow is an end-to-end open source platform for machine learning. Incomplete validation in `SparseAdd` results in allowing attackers to exploit undefined behavior (dereferencing null pointers) as well as write outside of bounds of heap allocated data. The implementation(https://github.com/tensorflow/tensorflow/blob/656e7673b14acd7835dc778867f84916c6d1cac2/tensorflow/core/kernels/sparse_add_op.cc) has a large set of validation for the two sparse tensor inputs (6 tensors in total), but does not validate that the tensors are not empty or that the second dimension of `*_indices` matches the size of corresponding `*_shape`. This allows attackers to send tensor triples that represent invalid sparse tensors to abuse code assumptions that are not protected by validation. The fix will be included in TensorFlow 2.5.0. We will also cherrypick this commit on TensorFlow 2.4.2, TensorFlow 2.3.3, TensorFlow 2.2.3 and TensorFlow 2.1.4, as these are also affected and still in supported range.
CWE description: The product writes data past the end, or before the beginning, of the intended buffer.
CVSS v3 base score: 7.8
Published: 2021-05-14T20:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: C++
Files changed: tensorflow/core/kernels/sparse_add_op.cc
Commit date: 2021-05-11 15:58:01 -0700
Commit message: Validate that a and b are proper sparse tensors

PiperOrigin-RevId: 373248068
Change-Id: I0a2041a0747901b3f00387a6a3bce9bca6b0b3b1

Diff: diff --git a/tensorflow/core/kernels/sparse_add_op.cc b/tensorflow/core/kernels/sparse_add_op.cc
index 346206365af8d..2bd05fa41adc2 100644
--- a/tensorflow/core/kernels/sparse_add_op.cc
+++ b/tensorflow/core/kernels/sparse_add_op.cc
@@ -42,10 +42,15 @@ class SparseAddOp : public OpKernel {
                     "Input indices should be matrices but received shapes: ",
                     a_indices->shape().DebugString(), " and ",
                     b_indices->shape().DebugString()));
     const int64 a_nnz = a_indices->dim_size(0);
     const int64 b_nnz = b_indices->dim_size(0);
+    const int num_dims = a_indices->dim_size(1);
+    OP_REQUIRES(ctx, b_indices->dim_size(1) == num_dims,
+                errors::InvalidArgument(
+                    "Input indices must have the same dimension, got ",
+                    num_dims, " and ", b_indices->dim_size(1)));
 
     OP_REQUIRES_OK(ctx, ctx->input("a_values", &a_values_t));
     OP_REQUIRES_OK(ctx, ctx->input("b_values", &b_values_t));
 
     OP_REQUIRES(ctx,
@@ -70,10 +75,17 @@ class SparseAddOp : public OpKernel {
                     TensorShapeUtils::IsVector(b_shape->shape()),
                 errors::InvalidArgument(
                     "Input shapes should be a vector but received shapes ",
                     a_shape->shape().DebugString(), " and ",
                     b_shape->shape().DebugString()));
+    OP_REQUIRES(
+        ctx, a_shape->NumElements() == num_dims,
+        errors::InvalidArgument("Second dimension of a_indices and length of "
+                                "a_shape must match, got ",
+                                num_dims, " and ", a_shape->NumElements()));
+    OP_REQUIRES(ctx, num_dims > 0,
+                errors::InvalidArgument("Tesors must not be empty"));
     OP_REQUIRES(
         ctx, a_shape->IsSameSize(*b_shape),
         errors::InvalidArgument(
             "Operands do not have the same ranks; got shapes: ",
             a_shape->SummarizeValue(10), " and ", b_shape->SummarizeValue(10)));
@@ -98,15 +110,10 @@ class SparseAddOp : public OpKernel {
     auto a_indices_mat = a_indices->matrix<int64>();
     auto b_indices_mat = b_indices->matrix<int64>();
     std::vector<std::pair<bool, int64>> entries_to_copy;  // from_a?, idx
     entries_to_copy.reserve(a_nnz + b_nnz);
     std::vector<T> out_values;
-    const int num_dims = a_shape->dim_size(0);
-
-    OP_REQUIRES(ctx, num_dims > 0,
-                errors::InvalidArgument("Invalid input_a shape. Received: ",
-                                        a_shape->DebugString()));
 
     // The input and output sparse tensors are assumed to be ordered along
     // increasing dimension number.
     int64 i = 0, j = 0;
     T s;

Vulnerable code:     const int num_dims = a_shape->dim_size(0);

    OP_REQUIRES(ctx, num_dims > 0,
                errors::InvalidArgument("Invalid input_a shape. Received: ",
                                        a_shape->DebugString()));
Fixed code:     const int num_dims = a_indices->dim_size(1);
    OP_REQUIRES(ctx, b_indices->dim_size(1) == num_dims,
                errors::InvalidArgument(
                    "Input indices must have the same dimension, got ",
                    num_dims, " and ", b_indices->dim_size(1)));
    OP_REQUIRES(
        ctx, a_shape->NumElements() == num_dims,
        errors::InvalidArgument("Second dimension of a_indices and length of "
                                "a_shape must match, got ",
                                num_dims, " and ", a_shape->NumElements()));
    OP_REQUIRES(ctx, num_dims > 0,
                errors::InvalidArgument("Tesors must not be empty"));