Vulnerability: Incomplete validation in `SparseReshape` ### Impact
Incomplete validation in `SparseReshape` results in a denial of service based on a `CHECK`-failure.

```python
import tensorflow as tf

input_indices = tf.constant(41, shape=[1, 1], dtype=tf.int64)
input_shape = tf.zeros([11], dtype=tf.int64)
new_shape = tf.zeros([1], dtype=tf.int64)

tf.raw_ops.SparseReshape(input_indices=input_indices,
    input_shape=input_shape,
    new_shape=new_shape)
``` 
    
The [implementation](https://github.com/tensorflow/tensorflow/blob/e87b51ce05c3eb172065a6ea5f48415854223285/tensorflow/core/kernels/sparse_reshape_op.cc#L40) has no validation that the input arguments specify a valid sparse tensor.

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

The fix will be included in TensorFlow 2.5.0. We will also cherrypick this commit on TensorFlow 2.4.2 and TensorFlow 2.3.3, as these are the only affected versions.

### 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 Ying Wang and Yakun Zhang of Baidu X-Team.
Vulnerability ID: GHSA-9rpc-5v9q-5r7f
CWE: CWE-20
Severity score: 3.6
Project: https://github.com/tensorflow/tensorflow
Published: 2021-05-21
File: tensorflow/core/kernels/sparse_reshape_op.cc
Commit message: Fix heap-buffer-overflow issue with `tf.raw_ops.SparseReshape`.

PiperOrigin-RevId: 371218558
Change-Id: I6a6dc5bf15b50a1d05bdd95e9ba347cb39f40f45
Commit author: Amit Patankar

Vulnerable version (before fix):
 limitations under the License. #include "tensorflow/core/framework/types.h" #include "tensorflow/core/kernels/reshape_util.h" #include "tensorflow/core/lib/gtl/inlined_vector.h"  namespace tensorflow {  class SparseReshapeOp : public OpKernel {   explicit SparseReshapeOp(OpKernelConstruction* context) : OpKernel(context) {}    void Compute(OpKernelContext* context) override {     ReshapeSparseTensor<Device>(context, context->input(0), context->input(1),                                 context->input(2), 0 /* output indices index */,                                 1 /* output shape index */);

Fixed version (after fix):
 limitations under the License. #include "tensorflow/core/framework/types.h" #include "tensorflow/core/kernels/reshape_util.h" #include "tensorflow/core/lib/gtl/inlined_vector.h" #include "tensorflow/core/platform/errors.h"  namespace tensorflow {  class SparseReshapeOp : public OpKernel {   explicit SparseReshapeOp(OpKernelConstruction* context) : OpKernel(context) {}    void Compute(OpKernelContext* context) override {     const Tensor& input_indices_in = context->input(0);     const Tensor& input_shape_in = context->input(1);      OP_REQUIRES(context, TensorShapeUtils::IsMatrix(input_indices_in.shape()),                 errors::InvalidArgument("Input must be a matrix."));     OP_REQUIRES(context, TensorShapeUtils::IsVector(input_shape_in.shape()),                 errors::InvalidArgument("Input shape must be a vector."));     OP_REQUIRES(context,                 input_indices_in.dim_size(1) == input_shape_in.dim_size(0),                 errors::InvalidArgument(                     "Input tensor rank must match input shape length."));     ReshapeSparseTensor<Device>(context, context->input(0), context->input(1),                                 context->input(2), 0 /* output indices index */,                                 1 /* output shape index */);

Patch:
patch: @@ -26,6 +26,7 @@ limitations under the License.
 #include "tensorflow/core/framework/types.h"
 #include "tensorflow/core/kernels/reshape_util.h"
 #include "tensorflow/core/lib/gtl/inlined_vector.h"
+#include "tensorflow/core/platform/errors.h"
 
 namespace tensorflow {
 
@@ -38,6 +39,17 @@ class SparseReshapeOp : public OpKernel {
   explicit SparseReshapeOp(OpKernelConstruction* context) : OpKernel(context) {}
 
   void Compute(OpKernelContext* context) override {
+    const Tensor& input_indices_in = context->input(0);
+    const Tensor& input_shape_in = context->input(1);
+
+    OP_REQUIRES(context, TensorShapeUtils::IsMatrix(input_indices_in.shape()),
+                errors::InvalidArgument("Input must be a matrix."));
+    OP_REQUIRES(context, TensorShapeUtils::IsVector(input_shape_in.shape()),
+                errors::InvalidArgument("Input shape must be a vector."));
+    OP_REQUIRES(context,
+                input_indices_in.dim_size(1) == input_shape_in.dim_size(0),
+                errors::InvalidArgument(
+                    "Input tensor rank must match input shape length."));
     ReshapeSparseTensor<Device>(context, context->input(0), context->input(1),
                                 context->input(2), 0 /* output indices index */,
                                 1 /* output shape index */);
path: tensorflow/core/kernels/sparse_reshape_op.cc