CVE-2021-29607: Improper Check for Unusual or Exceptional Conditions (CWE-754)
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_sparse_binary_op_shared.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 does not check or incorrectly checks for unusual or exceptional conditions that are not expected to occur frequently during day to day operation of the product.
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_sparse_binary_op_shared.cc
Commit date: 2021-05-11 18:37:23 -0700
Commit message: Validate that a and b are proper sparse tensors

PiperOrigin-RevId: 373274848
Change-Id: I3a665ac3a29dee9fb69bdf408a939330cb93ea75

Diff: diff --git a/tensorflow/core/kernels/sparse_sparse_binary_op_shared.cc b/tensorflow/core/kernels/sparse_sparse_binary_op_shared.cc
index 9fe42e05d879e..eb993a5965043 100644
--- a/tensorflow/core/kernels/sparse_sparse_binary_op_shared.cc
+++ b/tensorflow/core/kernels/sparse_sparse_binary_op_shared.cc
@@ -148,10 +148,11 @@ class SparseSparseBinaryOpShared : public OpKernel {
                     a_values_t->shape().DebugString(), " and ",
                     b_values_t->shape().DebugString()));
 
     const int64 a_nnz = a_indices_t->dim_size(0);
     const int64 b_nnz = b_indices_t->dim_size(0);
+
     const auto a_values = a_values_t->vec<T>();
     const auto b_values = b_values_t->vec<T>();
 
     OP_REQUIRES(
         ctx, a_values.size() == a_nnz && b_values.size() == b_nnz,
@@ -164,10 +165,18 @@ class SparseSparseBinaryOpShared : public OpKernel {
                     TensorShapeUtils::IsVector(b_shape_t->shape()),
                 errors::InvalidArgument(
                     "Input shapes should be a vector but received shapes ",
                     a_shape_t->shape().DebugString(), " and ",
                     b_shape_t->shape().DebugString()));
+    const int num_dims = a_indices_t->dim_size(1);
+    OP_REQUIRES(
+        ctx, a_shape_t->NumElements() == num_dims,
+        errors::InvalidArgument("Second dimension of a_indices and length of "
+                                "a_shape must match, got ",
+                                num_dims, " and ", a_shape_t->NumElements()));
+    OP_REQUIRES(ctx, num_dims > 0,
+                errors::InvalidArgument("Tensors must not be empty"));
     OP_REQUIRES(ctx, a_shape_t->IsSameSize(*b_shape_t),
                 errors::InvalidArgument(
                     "Operands do not have the same ranks; got shapes: ",
                     a_shape_t->SummarizeValue(10), " and ",
                     b_shape_t->SummarizeValue(10)));
@@ -178,16 +187,10 @@ class SparseSparseBinaryOpShared : public OpKernel {
                   errors::InvalidArgument("Operands' shapes do not match: got ",
                                           a_shape(i), " and ", b_shape(i),
                                           " for dimension ", i));
     }
 
-    OP_REQUIRES(
-        ctx, a_indices_t->dim_size(1) == b_indices_t->dim_size(1),
-        errors::InvalidArgument(
-            "Indices' dimensions do not match: got ", a_indices_t->dim_size(1),
-            " and ", b_indices_t->dim_size(1), " for the second dimension."));
-    const int num_dims = a_indices_t->dim_size(1);
     const auto a_indices_mat = a_indices_t->matrix<int64>();
     const auto b_indices_mat = b_indices_t->matrix<int64>();
     std::vector<T> a_augmented_values, b_augmented_values;
     std::vector<std::pair<bool, int64>> entries_to_copy;  // from_a?, idx
     UnionSparseIndicesAndValues(a_indices_mat, a_values, a_nnz, b_indices_mat,

Vulnerable code:     OP_REQUIRES(
        ctx, a_indices_t->dim_size(1) == b_indices_t->dim_size(1),
        errors::InvalidArgument(
            "Indices' dimensions do not match: got ", a_indices_t->dim_size(1),
            " and ", b_indices_t->dim_size(1), " for the second dimension."));
    const int num_dims = a_indices_t->dim_size(1);
Fixed code: 
    const int num_dims = a_indices_t->dim_size(1);
    OP_REQUIRES(
        ctx, a_shape_t->NumElements() == num_dims,
        errors::InvalidArgument("Second dimension of a_indices and length of "
                                "a_shape must match, got ",
                                num_dims, " and ", a_shape_t->NumElements()));
    OP_REQUIRES(ctx, num_dims > 0,
                errors::InvalidArgument("Tensors must not be empty"));