From 1d04d7d93f4ed3854abf75d6b712d72c3f70d6b6 Mon Sep 17 00:00:00 2001
From: Amit Patankar <amitpatankar@google.com>
Date: Thu, 29 Apr 2021 15:30:30 -0700
Subject: [PATCH] Fix heap-buffer-overflow issue with
 `tf.raw_ops.SparseReshape`.

PiperOrigin-RevId: 371218558
Change-Id: I6a6dc5bf15b50a1d05bdd95e9ba347cb39f40f45
---
 tensorflow/core/kernels/sparse_reshape_op.cc | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/tensorflow/core/kernels/sparse_reshape_op.cc b/tensorflow/core/kernels/sparse_reshape_op.cc
index 6eb5f0af..3896c959 100644
--- a/tensorflow/core/kernels/sparse_reshape_op.cc
+++ b/tensorflow/core/kernels/sparse_reshape_op.cc
@@ -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 {
 
@@ -34,6 +35,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(context, context->input(0), context->input(1),
                         context->input(2), 0 /* output indices index */,
                         1 /* output shape index */);
-- 
2.27.0