46 lines
1.9 KiB
Diff
46 lines
1.9 KiB
Diff
|
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
|
||
|
|
|