31 lines
1.2 KiB
Diff
31 lines
1.2 KiB
Diff
|
From a84358aa12f0b1518e606095ab9cfddbf597c121 Mon Sep 17 00:00:00 2001
|
||
|
|
From: Amit Patankar <amitpatankar@google.com>
|
||
|
|
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
|
||
|
|
|