1 files changed, 71 insertions, 20 deletions

diff --git a/lib/sanitizer_common/tests/sanitizer_allocator_test.cc b/lib/sanitizer_common/tests/sanitizer_allocator_test.cc
index de949ca7defe..d92a07fe4c2d 100644
--- a/lib/sanitizer_common/tests/sanitizer_allocator_test.cc
+++ b/lib/sanitizer_common/tests/sanitizer_allocator_test.cc
@@ -12,7 +12,9 @@
 //
 //===----------------------------------------------------------------------===//
 #include "sanitizer_common/sanitizer_allocator.h"
+#include "sanitizer_common/sanitizer_allocator_internal.h"
 #include "sanitizer_common/sanitizer_common.h"
+#include "sanitizer_common/sanitizer_flags.h"
 
 #include "sanitizer_test_utils.h"
 
@@ -67,6 +69,10 @@ TEST(SanitizerCommon, CompactSizeClassMap) {
   TestSizeClassMap<CompactSizeClassMap>();
 }
 
+TEST(SanitizerCommon, InternalSizeClassMap) {
+  TestSizeClassMap<InternalSizeClassMap>();
+}
+
 template <class Allocator>
 void TestSizeClassAllocator() {
   Allocator *a = new Allocator;
@@ -411,12 +417,20 @@ void TestCombinedAllocator() {
   memset(&cache, 0, sizeof(cache));
   a->InitCache(&cache);
 
+  bool allocator_may_return_null = common_flags()->allocator_may_return_null;
+  common_flags()->allocator_may_return_null = true;
   EXPECT_EQ(a->Allocate(&cache, -1, 1), (void*)0);
   EXPECT_EQ(a->Allocate(&cache, -1, 1024), (void*)0);
   EXPECT_EQ(a->Allocate(&cache, (uptr)-1 - 1024, 1), (void*)0);
   EXPECT_EQ(a->Allocate(&cache, (uptr)-1 - 1024, 1024), (void*)0);
   EXPECT_EQ(a->Allocate(&cache, (uptr)-1 - 1023, 1024), (void*)0);
 
+  common_flags()->allocator_may_return_null = false;
+  EXPECT_DEATH(a->Allocate(&cache, -1, 1),
+               "allocator is terminating the process");
+  // Restore the original value.
+  common_flags()->allocator_may_return_null = allocator_may_return_null;
+
   const uptr kNumAllocs = 100000;
   const uptr kNumIter = 10;
   for (uptr iter = 0; iter < kNumIter; iter++) {
@@ -611,6 +625,11 @@ TEST(Allocator, Stress) {
   }
 }
 
+TEST(Allocator, InternalAllocFailure) {
+  EXPECT_DEATH(Ident(InternalAlloc(10 << 20)),
+               "Unexpected mmap in InternalAllocator!");
+}
+
 TEST(Allocator, ScopedBuffer) {
   const int kSize = 512;
   {
@@ -625,16 +644,9 @@ TEST(Allocator, ScopedBuffer) {
   }
 }
 
-class IterationTestCallback {
- public:
-  explicit IterationTestCallback(std::set<void *> *chunks)
-    : chunks_(chunks) {}
-  void operator()(void *chunk) const {
-    chunks_->insert(chunk);
-  }
- private:
-  std::set<void *> *chunks_;
-};
+void IterationTestCallback(uptr chunk, void *arg) {
+  reinterpret_cast<std::set<uptr> *>(arg)->insert(chunk);
+}
 
 template <class Allocator>
 void TestSizeClassAllocatorIteration() {
@@ -663,15 +675,15 @@ void TestSizeClassAllocatorIteration() {
     }
   }
 
-  std::set<void *> reported_chunks;
-  IterationTestCallback callback(&reported_chunks);
+  std::set<uptr> reported_chunks;
   a->ForceLock();
-  a->ForEachChunk(callback);
+  a->ForEachChunk(IterationTestCallback, &reported_chunks);
   a->ForceUnlock();
 
   for (uptr i = 0; i < allocated.size(); i++) {
     // Don't use EXPECT_NE. Reporting the first mismatch is enough.
-    ASSERT_NE(reported_chunks.find(allocated[i]), reported_chunks.end());
+    ASSERT_NE(reported_chunks.find(reinterpret_cast<uptr>(allocated[i])),
+              reported_chunks.end());
   }
 
   a->TestOnlyUnmap();
@@ -698,22 +710,61 @@ TEST(SanitizerCommon, LargeMmapAllocatorIteration) {
   char *allocated[kNumAllocs];
   static const uptr size = 40;
   // Allocate some.
-  for (uptr i = 0; i < kNumAllocs; i++) {
+  for (uptr i = 0; i < kNumAllocs; i++)
     allocated[i] = (char *)a.Allocate(&stats, size, 1);
-  }
 
-  std::set<void *> reported_chunks;
-  IterationTestCallback callback(&reported_chunks);
+  std::set<uptr> reported_chunks;
   a.ForceLock();
-  a.ForEachChunk(callback);
+  a.ForEachChunk(IterationTestCallback, &reported_chunks);
   a.ForceUnlock();
 
   for (uptr i = 0; i < kNumAllocs; i++) {
     // Don't use EXPECT_NE. Reporting the first mismatch is enough.
-    ASSERT_NE(reported_chunks.find(allocated[i]), reported_chunks.end());
+    ASSERT_NE(reported_chunks.find(reinterpret_cast<uptr>(allocated[i])),
+              reported_chunks.end());
+  }
+  for (uptr i = 0; i < kNumAllocs; i++)
+    a.Deallocate(&stats, allocated[i]);
+}
+
+TEST(SanitizerCommon, LargeMmapAllocatorBlockBegin) {
+  LargeMmapAllocator<> a;
+  a.Init();
+  AllocatorStats stats;
+  stats.Init();
+
+  static const uptr kNumAllocs = 1024;
+  static const uptr kNumExpectedFalseLookups = 10000000;
+  char *allocated[kNumAllocs];
+  static const uptr size = 4096;
+  // Allocate some.
+  for (uptr i = 0; i < kNumAllocs; i++) {
+    allocated[i] = (char *)a.Allocate(&stats, size, 1);
   }
+
+  a.ForceLock();
+  for (uptr i = 0; i < kNumAllocs  * kNumAllocs; i++) {
+    // if ((i & (i - 1)) == 0) fprintf(stderr, "[%zd]\n", i);
+    char *p1 = allocated[i % kNumAllocs];
+    EXPECT_EQ(p1, a.GetBlockBeginFastLocked(p1));
+    EXPECT_EQ(p1, a.GetBlockBeginFastLocked(p1 + size / 2));
+    EXPECT_EQ(p1, a.GetBlockBeginFastLocked(p1 + size - 1));
+    EXPECT_EQ(p1, a.GetBlockBeginFastLocked(p1 - 100));
+  }
+
+  for (uptr i = 0; i < kNumExpectedFalseLookups; i++) {
+    void *p = reinterpret_cast<void *>(i % 1024);
+    EXPECT_EQ((void *)0, a.GetBlockBeginFastLocked(p));
+    p = reinterpret_cast<void *>(~0L - (i % 1024));
+    EXPECT_EQ((void *)0, a.GetBlockBeginFastLocked(p));
+  }
+  a.ForceUnlock();
+
+  for (uptr i = 0; i < kNumAllocs; i++)
+    a.Deallocate(&stats, allocated[i]);
 }
 
+
 #if SANITIZER_WORDSIZE == 64
 // Regression test for out-of-memory condition in PopulateFreeList().
 TEST(SanitizerCommon, SizeClassAllocator64PopulateFreeListOOM) {