Vulnerability: Integer overflow leading to crash in Tensorflow ### Impact 
The [implementation of `SparseCountSparseOutput`](https://github.com/tensorflow/tensorflow/blob/5100e359aef5c8021f2e71c7b986420b85ce7b3d/tensorflow/core/kernels/count_ops.cc#L168-L273) can be made to crash a TensorFlow process by an integer overflow whose result is then used in a memory allocation:

```python
import tensorflow as tf
import numpy as np
    
tf.raw_ops.SparseCountSparseOutput(
  indices=[[1,1]],
  values=[2],
  dense_shape=[2 ** 31, 2 ** 32],
  weights=[1],
  binary_output=True,
  minlength=-1,
  maxlength=-1,
  name=None)
```

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

The fix will be included in TensorFlow 2.8.0. We will also cherrypick this commit on TensorFlow 2.7.1, TensorFlow 2.6.3, and TensorFlow 2.5.3, 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 Faysal Hossain Shezan from University of Virginia.
Vulnerability ID: GHSA-x4qx-4fjv-hmw6
CWE: CWE-190
Severity score: 6.5
Project: https://github.com/tensorflow/tensorflow
Published: 2022-02-09
File: tensorflow/core/kernels/count_ops.cc
Commit message: Prevent crash due to integer overflow followed by allocating negative sized array.

PiperOrigin-RevId: 414891322
Change-Id: I5df390e0dc1d9f115209293708950cdf9306931c
Commit author: Mihai Maruseac

Vulnerable version (before fix):
 See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/  #include "absl/container/flat_hash_map.h" #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/op_requires.h" limitations under the License.  namespace tensorflow {  template <class T> using BatchedMap = std::vector<absl::flat_hash_map<int64_t, T>>;  class SparseCount : public OpKernel {      bool is_1d = shape.NumElements() == 1;     int num_batches = is_1d ? 1 : shape_vector(0);      const auto values_values = values.flat<T>();     const auto weight_values = weights.flat<W>();

Fixed version (after fix):
 See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/  #include <limits>  #include "absl/container/flat_hash_map.h" #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/op_requires.h" limitations under the License.  namespace tensorflow {  // Don't allocate too large `BatchedMap<T>` objects static int kMaxBatches = std::numeric_limits<int>::max();  template <class T> using BatchedMap = std::vector<absl::flat_hash_map<int64_t, T>>;  class SparseCount : public OpKernel {      bool is_1d = shape.NumElements() == 1;     int num_batches = is_1d ? 1 : shape_vector(0);     OP_REQUIRES(         context, 0 < num_batches && num_batches < kMaxBatches,         errors::InvalidArgument("Cannot allocate ", num_batches,                                 " batches, is the dense shape too wide?"));      const auto values_values = values.flat<T>();     const auto weight_values = weights.flat<W>();

Patch:
patch: @@ -13,6 +13,8 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include <limits>
+
 #include "absl/container/flat_hash_map.h"
 #include "tensorflow/core/framework/op_kernel.h"
 #include "tensorflow/core/framework/op_requires.h"
@@ -23,6 +25,9 @@ limitations under the License.
 
 namespace tensorflow {
 
+// Don't allocate too large `BatchedMap<T>` objects
+static int kMaxBatches = std::numeric_limits<int>::max();
+
 template <class T>
 using BatchedMap = std::vector<absl::flat_hash_map<int64_t, T>>;
 
@@ -235,6 +240,10 @@ class SparseCount : public OpKernel {
 
     bool is_1d = shape.NumElements() == 1;
     int num_batches = is_1d ? 1 : shape_vector(0);
+    OP_REQUIRES(
+        context, 0 < num_batches && num_batches < kMaxBatches,
+        errors::InvalidArgument("Cannot allocate ", num_batches,
+                                " batches, is the dense shape too wide?"));
 
     const auto values_values = values.flat<T>();
     const auto weight_values = weights.flat<W>();
path: tensorflow/core/kernels/count_ops.cc