From a84358aa12f0b1518e606095ab9cfddbf597c121 Mon Sep 17 00:00:00 2001 From: Amit Patankar Date: Tue, 4 May 2021 13:45:57 -0700 Subject: [PATCH] Fix heap-buffer-overflow issue with `tf.raw_ops.RaggedTensorToTensor`. --- tensorflow/core/kernels/ragged_tensor_to_tensor_op.cc | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tensorflow/core/kernels/ragged_tensor_to_tensor_op.cc b/tensorflow/core/kernels/ragged_tensor_to_tensor_op.cc index 988a3333..38cb4257 100644 --- a/tensorflow/core/kernels/ragged_tensor_to_tensor_op.cc +++ b/tensorflow/core/kernels/ragged_tensor_to_tensor_op.cc @@ -313,6 +313,12 @@ class RaggedTensorToTensorBaseOp : public OpKernel { output_index_multiplier, output_size, result); return tensorflow::Status::OK(); case RowPartitionType::ROW_SPLITS: + if (row_partition_tensor.size() - 1 > parent_output_index.size()) { + return errors::InvalidArgument( + "Row partition size is greater than output size: ", + row_partition_tensor.size() - 1, " > ", + parent_output_index.size()); + } CalculateOutputIndexRowSplit( context, row_partition_tensor, parent_output_index, output_index_multiplier, output_size, result); -- 2.23.0