From 1d04d7d93f4ed3854abf75d6b712d72c3f70d6b6 Mon Sep 17 00:00:00 2001 From: Amit Patankar 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