From 532f5c5a547126c634fefd43bbad1dc6417678ac Mon Sep 17 00:00:00 2001 From: Mihai Maruseac Date: Thu, 29 Jul 2021 22:24:39 -0700 Subject: [PATCH] Prevent nullptr deref in validation of indexes in map ops. PiperOrigin-RevId: 387738023 Change-Id: I83d18d36a7b82ffd2a40b5124a4e5b4c72238f27 --- tensorflow/core/kernels/map_stage_op.cc | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/tensorflow/core/kernels/map_stage_op.cc b/tensorflow/core/kernels/map_stage_op.cc index 9411792762baa..1a2f5a8fa2247 100644 --- a/tensorflow/core/kernels/map_stage_op.cc +++ b/tensorflow/core/kernels/map_stage_op.cc @@ -210,9 +210,9 @@ class StagingMap : public ResourceBase { const OptionalTuple& tuple) TF_EXCLUSIVE_LOCKS_REQUIRED(mu_) { if (tuple[index].has_value()) { - return Status(errors::InvalidArgument( + return errors::InvalidArgument( "The tensor for index '", index, "' for key '", key.scalar()(), - "' was already initialized '", dtypes_.size(), "'.")); + "' was already initialized '", dtypes_.size(), "'."); } return Status::OK(); @@ -220,6 +220,10 @@ class StagingMap : public ResourceBase { // Check that the indices are strictly ordered Status check_index_ordering(const Tensor& indices) { + if (indices.NumElements() == 0) { + return errors::InvalidArgument("Indices are empty"); + } + auto findices = indices.flat(); for (std::size_t i = 0; i < findices.dimension(0) - 1; ++i) { @@ -227,8 +231,7 @@ class StagingMap : public ResourceBase { continue; } - return Status( - errors::InvalidArgument("Indices are not strictly ordered")); + return errors::InvalidArgument("Indices are not strictly ordered"); } return Status::OK(); @@ -238,10 +241,10 @@ class StagingMap : public ResourceBase { Status check_memory_limit(std::size_t bytes) TF_EXCLUSIVE_LOCKS_REQUIRED(mu_) { if (has_memory_limit() && bytes > memory_limit_) { - return Status(errors::ResourceExhausted( + return errors::ResourceExhausted( "Attempted to insert tensors with combined size of '", bytes, "' bytes into Staging Area with a memory limit of '", memory_limit_, - "'.")); + "'."); } return Status::OK();