From 8ba6fa29cd8bf9cef9b718dc31c78c73081f5b31 Mon Sep 17 00:00:00 2001 From: Amit Patankar 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()(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]++; }