Add pointer-stability validation in btree.

When we insert a new element, when ASan is enabled, we replace the node that the new element is on in order to try to detect cases of code depending on pointer/reference-stability.

PiperOrigin-RevId: 528826645
Change-Id: Ie5c15c13016a8aa831a0d1edc3ad33c1f5ca4d65

1 files changed, 37 insertions, 5 deletions

diff --git a/absl/container/btree_test.cc b/absl/container/btree_test.cc
index 1b6a3f47..15335c8c 100644
--- a/absl/container/btree_test.cc
+++ b/absl/container/btree_test.cc
@@ -18,6 +18,7 @@
 #include <array>
 #include <cstdint>
 #include <functional>
+#include <iostream>
 #include <iterator>
 #include <limits>
 #include <map>
@@ -1480,9 +1481,17 @@ void ExpectOperationCounts(const int expected_moves,
   tracker->ResetCopiesMovesSwaps();
 }
 
+#ifdef ABSL_HAVE_ADDRESS_SANITIZER
+constexpr bool kAsan = true;
+#else
+constexpr bool kAsan = false;
+#endif
+
 // Note: when the values in this test change, it is expected to have an impact
 // on performance.
 TEST(Btree, MovesComparisonsCopiesSwapsTracking) {
+  if (kAsan) GTEST_SKIP() << "We do extra operations in ASan mode.";
+
   InstanceTracker tracker;
   // Note: this is minimum number of values per node.
   SizedBtreeSet<MovableOnlyInstance, /*TargetValuesPerNode=*/4> set4;
@@ -1534,6 +1543,8 @@ struct MovableOnlyInstanceThreeWayCompare {
 // Note: when the values in this test change, it is expected to have an impact
 // on performance.
 TEST(Btree, MovesComparisonsCopiesSwapsTrackingThreeWayCompare) {
+  if (kAsan) GTEST_SKIP() << "We do extra operations in ASan mode.";
+
   InstanceTracker tracker;
   // Note: this is minimum number of values per node.
   SizedBtreeSet<MovableOnlyInstance, /*TargetValuesPerNode=*/4,
@@ -3218,19 +3229,24 @@ TEST(Btree, InvalidIteratorUse) {
   if (!BtreeNodePeer::UsesGenerations<absl::btree_set<int>>())
     GTEST_SKIP() << "Generation validation for iterators is disabled.";
 
+  // Invalid memory use can trigger heap-use-after-free in ASan or invalidated
+  // iterator assertions.
+  constexpr const char *kInvalidMemoryDeathMessage =
+      "heap-use-after-free|invalidated iterator";
+
   {
     absl::btree_set<int> set;
     for (int i = 0; i < 10; ++i) set.insert(i);
     auto it = set.begin();
     set.erase(it++);
-    EXPECT_DEATH(set.erase(it++), "invalidated iterator");
+    EXPECT_DEATH(set.erase(it++), kInvalidMemoryDeathMessage);
   }
   {
     absl::btree_set<int> set;
     for (int i = 0; i < 10; ++i) set.insert(i);
     auto it = set.insert(20).first;
     set.insert(30);
-    EXPECT_DEATH(*it, "invalidated iterator");
+    EXPECT_DEATH(*it, kInvalidMemoryDeathMessage);
   }
   {
     absl::btree_set<int> set;
@@ -3238,15 +3254,15 @@ TEST(Btree, InvalidIteratorUse) {
     auto it = set.find(5000);
     ASSERT_NE(it, set.end());
     set.erase(1);
-    EXPECT_DEATH(*it, "invalidated iterator");
+    EXPECT_DEATH(*it, kInvalidMemoryDeathMessage);
   }
   {
     absl::btree_set<int> set;
     for (int i = 0; i < 10; ++i) set.insert(i);
     auto it = set.insert(20).first;
     set.insert(30);
-    EXPECT_DEATH(void(it == set.begin()), "invalidated iterator");
-    EXPECT_DEATH(void(set.begin() == it), "invalidated iterator");
+    EXPECT_DEATH(void(it == set.begin()), kInvalidMemoryDeathMessage);
+    EXPECT_DEATH(void(set.begin() == it), kInvalidMemoryDeathMessage);
   }
 }
 #endif
@@ -3537,6 +3553,22 @@ TEST(Btree, InvalidIteratorComparison) {
   EXPECT_DEATH(void(iter2 == iter1), kDifferentContainerDeathMessage);
 }
 
+TEST(Btree, InvalidPointerUse) {
+  if (!kAsan)
+    GTEST_SKIP() << "We only detect invalid pointer use in ASan mode.";
+
+  absl::btree_set<int> set;
+  set.insert(0);
+  const int *ptr = &*set.begin();
+  set.insert(1);
+  EXPECT_DEATH(std::cout << *ptr, "heap-use-after-free");
+  size_t slots_per_node = BtreeNodePeer::GetNumSlotsPerNode<decltype(set)>();
+  for (int i = 2; i < slots_per_node - 1; ++i) set.insert(i);
+  ptr = &*set.begin();
+  set.insert(static_cast<int>(slots_per_node));
+  EXPECT_DEATH(std::cout << *ptr, "heap-use-after-free");
+}
+
 }  // namespace
 }  // namespace container_internal
 ABSL_NAMESPACE_END