Vulnerability: CHECK-fail in AddManySparseToTensorsMap ### Impact
An attacker can trigger a denial of service via a `CHECK`-fail in  `tf.raw_ops.AddManySparseToTensorsMap`:

```python
import tensorflow as tf
import numpy as np

sparse_indices = tf.constant(530, shape=[1, 1], dtype=tf.int64)
sparse_values = tf.ones([1], dtype=tf.int64)

shape = tf.Variable(tf.ones([55], dtype=tf.int64))
shape[:8].assign(np.array([855, 901, 429, 892, 892, 852, 93, 96], dtype=np.int64))

tf.raw_ops.AddManySparseToTensorsMap(sparse_indices=sparse_indices,
                    sparse_values=sparse_values,
                    sparse_shape=shape)
```

This is because the [implementation](https://github.com/tensorflow/tensorflow/blob/6f9896890c4c703ae0a0845394086e2e1e523299/tensorflow/core/kernels/sparse_tensors_map_ops.cc#L257) takes the values specified in `sparse_shape` as dimensions for the output shape: 

```cc
    TensorShape tensor_input_shape(input_shape->vec<int64>());
```

The [`TensorShape` constructor](https://github.com/tensorflow/tensorflow/blob/6f9896890c4c703ae0a0845394086e2e1e523299/tensorflow/core/framework/tensor_shape.cc#L183-L188) uses a `CHECK` operation which triggers when [`InitDims`](https://github.com/tensorflow/tensorflow/blob/6f9896890c4c703ae0a0845394086e2e1e523299/tensorflow/core/framework/tensor_shape.cc#L212-L296) returns a non-OK status.
  
```cc
template <class Shape>
TensorShapeBase<Shape>::TensorShapeBase(gtl::ArraySlice<int64> dim_sizes) {
  set_tag(REP16);
  set_data_type(DT_INVALID);
  TF_CHECK_OK(InitDims(dim_sizes));
}
```

In our scenario, this occurs when adding a dimension from the argument results in overflow:

```cc
template <class Shape>
Status TensorShapeBase<Shape>::InitDims(gtl::ArraySlice<int64> dim_sizes) {
  ...
  Status status = Status::OK();
  for (int64 s : dim_sizes) {
    status.Update(AddDimWithStatus(internal::SubtleMustCopy(s)));
    if (!status.ok()) {
      return status;
    }
  }
}

template <class Shape>
Status TensorShapeBase<Shape>::AddDimWithStatus(int64 size) {
  ...
  int64 new_num_elements;
  if (kIsPartial && (num_elements() < 0 || size < 0)) {
    new_num_elements = -1;
  } else {
    new_num_elements = MultiplyWithoutOverflow(num_elements(), size);
    if (TF_PREDICT_FALSE(new_num_elements < 0)) {
        return errors::Internal("Encountered overflow when multiplying ",
                                num_elements(), " with ", size,
                                ", result: ", new_num_elements);
      }
  }
  ...
}
```

This is a legacy implementation of the constructor and operations should use `BuildTensorShapeBase` or `AddDimWithStatus` to prevent `CHECK`-failures in the presence of overflows.

### Patches
We have patched the issue in GitHub commit [69c68ecbb24dff3fa0e46da0d16c821a2dd22d7c](https://github.com/tensorflow/tensorflow/commit/69c68ecbb24dff3fa0e46da0d16c821a2dd22d7c).

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.

### For more information
Please consult [our security guide](https://github.com/tensorflow/tensorflow/blob/master/SECURITY.md) for more information regarding the security model and how to contact us with issues and questions.

### Attribution
This vulnerability has been reported by Yakun Zhang and Ying Wang of Baidu X-Team.
Vulnerability ID: GHSA-2cpx-427x-q2c6
CWE: CWE-190
Severity score: 2.5
Project: https://github.com/tensorflow/tensorflow
Published: 2021-05-21
File: tensorflow/core/kernels/sparse_tensors_map_ops.cc
Commit message: Fix overflow CHECK issue with `tf.raw_ops.AddManySparseToTensorsMap`.

PiperOrigin-RevId: 369492969
Change-Id: I1d70d6c0c92e3d7a25bc3b3aa2a0c0ac9688bf81
Commit author: Amit Patankar

Vulnerable version (before fix):
 limitations under the License. #include <utility> #include <vector>  #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/register_types.h"  #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/register_types.h" #include "tensorflow/core/framework/resource_mgr.h" #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/framework/tensor_util.h" #include "tensorflow/core/framework/types.h" #include "tensorflow/core/lib/gtl/inlined_vector.h" #include "tensorflow/core/util/sparse/sparse_tensor.h"  namespace tensorflow { class AddManySparseToTensorsMapOp : public SparseTensorAccessingOp {         errors::InvalidArgument(             "Rank of input SparseTensor should be > 1, but saw rank: ", rank));      TensorShape tensor_input_shape(input_shape->vec<int64>());     gtl::InlinedVector<int64, 8> std_order(rank);     std::iota(std_order.begin(), std_order.end(), 0);     SparseTensor input_st;     OP_REQUIRES_OK(context, SparseTensor::Create(*input_indices, *input_values,                                                  tensor_input_shape, std_order,                                                  &input_st));      auto input_shape_t = input_shape->vec<int64>();     const int64 N = input_shape_t(0);      Tensor sparse_handles(DT_INT64, TensorShape({N}));     auto sparse_handles_t = sparse_handles.vec<int64>(); class AddManySparseToTensorsMapOp : public SparseTensorAccessingOp {     // minibatch entries.     TensorShape output_shape;     OP_REQUIRES_OK(context, TensorShapeUtils::MakeShape(                                 input_shape_t.data() + 1,                                 input_shape->NumElements() - 1, &output_shape));      // Get groups by minibatch dimension

Fixed version (after fix):
 limitations under the License. #include <utility> #include <vector>  #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/register_types.h" #include "tensorflow/core/framework/resource_mgr.h" #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/framework/tensor_util.h" #include "tensorflow/core/framework/types.h" #include "tensorflow/core/lib/gtl/inlined_vector.h" #include "tensorflow/core/util/overflow.h" #include "tensorflow/core/util/sparse/sparse_tensor.h"  namespace tensorflow { class AddManySparseToTensorsMapOp : public SparseTensorAccessingOp {         errors::InvalidArgument(             "Rank of input SparseTensor should be > 1, but saw rank: ", rank));      auto input_shape_vec = input_shape->vec<int64>();     int new_num_elements = 1;     bool overflow_ocurred = false;     for (int i = 0; i < input_shape_vec.size(); i++) {       new_num_elements =           MultiplyWithoutOverflow(new_num_elements, input_shape_vec(i));       if (new_num_elements < 0) {         overflow_ocurred = true;       }     }      OP_REQUIRES(         context, !overflow_ocurred,         errors::Internal("Encountered overflow from large input shape."));      TensorShape tensor_input_shape(input_shape_vec);     gtl::InlinedVector<int64, 8> std_order(rank);     std::iota(std_order.begin(), std_order.end(), 0);     SparseTensor input_st;     OP_REQUIRES_OK(context, SparseTensor::Create(*input_indices, *input_values,                                                  tensor_input_shape, std_order,                                                  &input_st));      const int64 N = input_shape_vec(0);      Tensor sparse_handles(DT_INT64, TensorShape({N}));     auto sparse_handles_t = sparse_handles.vec<int64>(); class AddManySparseToTensorsMapOp : public SparseTensorAccessingOp {     // minibatch entries.     TensorShape output_shape;     OP_REQUIRES_OK(context, TensorShapeUtils::MakeShape(                                 input_shape_vec.data() + 1,                                 input_shape->NumElements() - 1, &output_shape));      // Get groups by minibatch dimension

Patch:
patch: @@ -21,16 +21,14 @@ limitations under the License.
 #include <utility>
 #include <vector>
 
-#include "tensorflow/core/framework/op_kernel.h"
-#include "tensorflow/core/framework/register_types.h"
-
 #include "tensorflow/core/framework/op_kernel.h"
 #include "tensorflow/core/framework/register_types.h"
 #include "tensorflow/core/framework/resource_mgr.h"
 #include "tensorflow/core/framework/tensor.h"
 #include "tensorflow/core/framework/tensor_util.h"
 #include "tensorflow/core/framework/types.h"
 #include "tensorflow/core/lib/gtl/inlined_vector.h"
+#include "tensorflow/core/util/overflow.h"
 #include "tensorflow/core/util/sparse/sparse_tensor.h"
 
 namespace tensorflow {
@@ -254,16 +252,30 @@ class AddManySparseToTensorsMapOp : public SparseTensorAccessingOp {
         errors::InvalidArgument(
             "Rank of input SparseTensor should be > 1, but saw rank: ", rank));
 
-    TensorShape tensor_input_shape(input_shape->vec<int64>());
+    auto input_shape_vec = input_shape->vec<int64>();
+    int new_num_elements = 1;
+    bool overflow_ocurred = false;
+    for (int i = 0; i < input_shape_vec.size(); i++) {
+      new_num_elements =
+          MultiplyWithoutOverflow(new_num_elements, input_shape_vec(i));
+      if (new_num_elements < 0) {
+        overflow_ocurred = true;
+      }
+    }
+
+    OP_REQUIRES(
+        context, !overflow_ocurred,
+        errors::Internal("Encountered overflow from large input shape."));
+
+    TensorShape tensor_input_shape(input_shape_vec);
     gtl::InlinedVector<int64, 8> std_order(rank);
     std::iota(std_order.begin(), std_order.end(), 0);
     SparseTensor input_st;
     OP_REQUIRES_OK(context, SparseTensor::Create(*input_indices, *input_values,
                                                  tensor_input_shape, std_order,
                                                  &input_st));
 
-    auto input_shape_t = input_shape->vec<int64>();
-    const int64 N = input_shape_t(0);
+    const int64 N = input_shape_vec(0);
 
     Tensor sparse_handles(DT_INT64, TensorShape({N}));
     auto sparse_handles_t = sparse_handles.vec<int64>();
@@ -274,7 +286,7 @@ class AddManySparseToTensorsMapOp : public SparseTensorAccessingOp {
     // minibatch entries.
     TensorShape output_shape;
     OP_REQUIRES_OK(context, TensorShapeUtils::MakeShape(
-                                input_shape_t.data() + 1,
+                                input_shape_vec.data() + 1,
                                 input_shape->NumElements() - 1, &output_shape));
 
     // Get groups by minibatch dimension
path: tensorflow/core/kernels/sparse_tensors_map_ops.cc