In sanitizer mode, detect when references become invalidated after reserved growth runs out.
PiperOrigin-RevId: 502625638
Change-Id: I1c06b2162dbdaaa6a36cea503ac6d07cd157b2e2
diff --git a/absl/container/internal/raw_hash_set.h b/absl/container/internal/raw_hash_set.h
index 0baca2a..61ef196 100644
--- a/absl/container/internal/raw_hash_set.h
+++ b/absl/container/internal/raw_hash_set.h
@@ -750,6 +750,13 @@
#endif
class CommonFieldsGenerationInfoEnabled {
+ // A sentinel value for reserved_growth_ indicating that we just ran out of
+ // reserved growth on the last insertion. When reserve is called and then
+ // insertions take place, reserved_growth_'s state machine is N, ..., 1,
+ // kReservedGrowthJustRanOut, 0.
+ static constexpr size_t kReservedGrowthJustRanOut =
+ (std::numeric_limits<size_t>::max)();
+
public:
CommonFieldsGenerationInfoEnabled() = default;
CommonFieldsGenerationInfoEnabled(CommonFieldsGenerationInfoEnabled&& that)
@@ -760,9 +767,19 @@
CommonFieldsGenerationInfoEnabled& operator=(
CommonFieldsGenerationInfoEnabled&&) = default;
+ // Whether we should rehash on insert in order to detect bugs of using invalid
+ // references. We rehash on the first insertion after reserved_growth_ reaches
+ // 0 after a call to reserve.
+ // TODO(b/254649633): we could potentially do a rehash with low probability
+ // whenever reserved_growth_ is zero.
+ bool should_rehash_for_bug_detection_on_insert() const {
+ return reserved_growth_ == kReservedGrowthJustRanOut;
+ }
void maybe_increment_generation_on_insert() {
+ if (reserved_growth_ == kReservedGrowthJustRanOut) reserved_growth_ = 0;
+
if (reserved_growth_ > 0) {
- --reserved_growth_;
+ if (--reserved_growth_ == 0) reserved_growth_ = kReservedGrowthJustRanOut;
} else {
++*generation_;
}
@@ -782,12 +799,6 @@
// a prior call to reserve. Note: we store reserved growth rather than
// reservation size because calls to erase() decrease size_ but don't decrease
// reserved growth.
- // TODO(b/254649633): we can use reserved_growth_ to find more bugs by doing
- // extra rehashes in sanitizer mode when reserved_growth_ is 0. We could
- // potentially do a rehash with low probability whenever reserved_growth_ is
- // zero, but also add a deterministic rehash the first insert after
- // reserved_growth_ is zero after a call to reserve. This would detect cases
- // of invalid references (as opposed to invalid iterators).
size_t reserved_growth_ = 0;
// Pointer to the generation counter, which is used to validate iterators and
// is stored in the backing array between the control bytes and the slots.
@@ -809,6 +820,7 @@
CommonFieldsGenerationInfoDisabled& operator=(
CommonFieldsGenerationInfoDisabled&&) = default;
+ bool should_rehash_for_bug_detection_on_insert() const { return false; }
void maybe_increment_generation_on_insert() {}
void reset_reserved_growth(size_t, size_t) {}
size_t reserved_growth() const { return 0; }
@@ -938,6 +950,12 @@
// A valid capacity is a non-zero integer `2^m - 1`.
inline bool IsValidCapacity(size_t n) { return ((n + 1) & n) == 0 && n > 0; }
+// Returns the next valid capacity after `n`.
+inline size_t NextCapacity(size_t n) {
+ assert(IsValidCapacity(n) || n == 0);
+ return n * 2 + 1;
+}
+
// Applies the following mapping to every byte in the control array:
// * kDeleted -> kEmpty
// * kEmpty -> kEmpty
@@ -2385,7 +2403,7 @@
drop_deletes_without_resize();
} else {
// Otherwise grow the container.
- resize(cap * 2 + 1);
+ resize(NextCapacity(cap));
}
}
@@ -2449,8 +2467,16 @@
//
// REQUIRES: At least one non-full slot available.
size_t prepare_insert(size_t hash) ABSL_ATTRIBUTE_NOINLINE {
+ const bool rehash_for_bug_detection =
+ common().should_rehash_for_bug_detection_on_insert();
+ if (rehash_for_bug_detection) {
+ // Move to a different heap allocation in order to detect bugs.
+ const size_t cap = capacity();
+ resize(growth_left() > 0 ? cap : NextCapacity(cap));
+ }
auto target = find_first_non_full(common(), hash);
- if (ABSL_PREDICT_FALSE(growth_left() == 0 &&
+ if (!rehash_for_bug_detection &&
+ ABSL_PREDICT_FALSE(growth_left() == 0 &&
!IsDeleted(control()[target.offset]))) {
rehash_and_grow_if_necessary();
target = find_first_non_full(common(), hash);
diff --git a/absl/container/internal/raw_hash_set_test.cc b/absl/container/internal/raw_hash_set_test.cc
index 801b758..3d3b089 100644
--- a/absl/container/internal/raw_hash_set_test.cc
+++ b/absl/container/internal/raw_hash_set_test.cc
@@ -20,6 +20,7 @@
#include <cstdint>
#include <deque>
#include <functional>
+#include <iostream>
#include <iterator>
#include <list>
#include <map>
@@ -2286,14 +2287,20 @@
t.insert(i);
EXPECT_EQ(*it, 0);
}
+ // ptr will become invalidated on rehash.
+ const int64_t* ptr = &*it;
+
// erase decreases size but does not decrease reserved growth so the next
// insertion still invalidates iterators.
t.erase(0);
- // Unreserved growth can rehash.
+ // The first insert after reserved growth is 0 is guaranteed to rehash when
+ // generations are enabled.
t.insert(10);
EXPECT_DEATH_IF_SUPPORTED(*it, kInvalidIteratorDeathMessage);
EXPECT_DEATH_IF_SUPPORTED(void(it == t.begin()),
kInvalidIteratorDeathMessage);
+ EXPECT_DEATH_IF_SUPPORTED(std::cout << *ptr,
+ "heap-use-after-free|use-of-uninitialized-value");
}
TEST(Table, ReservedGrowthUpdatesWhenTableDoesntGrow) {
