62 lines
2.3 KiB
Diff
62 lines
2.3 KiB
Diff
From 532f5c5a547126c634fefd43bbad1dc6417678ac Mon Sep 17 00:00:00 2001
|
|
From: Mihai Maruseac <mihaimaruseac@google.com>
|
|
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<int64>()(),
|
|
- "' 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<int>();
|
|
|
|
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();
|