Skip to content
Permalink
Browse files Browse the repository at this point in the history
Prevent nullptr deref in validation of indexes in map ops.
PiperOrigin-RevId: 387738023
Change-Id: I83d18d36a7b82ffd2a40b5124a4e5b4c72238f27
  • Loading branch information
mihaimaruseac authored and tensorflower-gardener committed Jul 30, 2021
1 parent a4e1386 commit 532f5c5
Showing 1 changed file with 9 additions and 6 deletions.
15 changes: 9 additions & 6 deletions tensorflow/core/kernels/map_stage_op.cc
Expand Up @@ -210,25 +210,28 @@ 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<int64>()(),
"' was already initialized '", dtypes_.size(), "'."));
"' was already initialized '", dtypes_.size(), "'.");
}

return Status::OK();
}

// 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<int>();

for (std::size_t i = 0; i < findices.dimension(0) - 1; ++i) {
if (findices(i) < findices(i + 1)) {
continue;
}

return Status(
errors::InvalidArgument("Indices are not strictly ordered"));
return errors::InvalidArgument("Indices are not strictly ordered");
}

return Status::OK();
Expand All @@ -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();
Expand Down

0 comments on commit 532f5c5

Please sign in to comment.