27 lines
1.1 KiB
Diff
27 lines
1.1 KiB
Diff
From 8ba6fa29cd8bf9cef9b718dc31c78c73081f5b31 Mon Sep 17 00:00:00 2001
|
|
From: Amit Patankar <amitpatankar@google.com>
|
|
Date: Thu, 29 Apr 2021 17:58:08 -0700
|
|
Subject: [PATCH] Fix heap-buffer-overflow issue with `tf.raw_ops.SparseSplit`.
|
|
|
|
PiperOrigin-RevId: 371242872
|
|
Change-Id: I482bb3d12602c7c3cc9446f97fb9f584bb98e9a4
|
|
---
|
|
tensorflow/core/util/sparse/sparse_tensor.h | 4 ++++
|
|
1 file changed, 4 insertions(+)
|
|
|
|
diff --git a/tensorflow/core/util/sparse/sparse_tensor.h b/tensorflow/core/util/sparse/sparse_tensor.h
|
|
index 062226d7699bc..341290dbbc698 100644
|
|
--- a/tensorflow/core/util/sparse/sparse_tensor.h
|
|
+++ b/tensorflow/core/util/sparse/sparse_tensor.h
|
|
@@ -527,6 +527,10 @@ inline Status SparseTensor::Split(const SparseTensor& input_tensor,
|
|
for (int i = 0; i < input_tensor.indices().dim_size(0); ++i) {
|
|
const int dim = input_tensor.indices().matrix<int64>()(i, split_dim);
|
|
int slice_index = GetSliceIndex(dim, split_size, residual);
|
|
+ if (slice_index >= num_values.size()) {
|
|
+ return errors::InvalidArgument("Slice index ", slice_index,
|
|
+ " is larger than num_split.");
|
|
+ }
|
|
num_values[slice_index]++;
|
|
}
|
|
|