[Sanitizer] move low-level (mmap-based) allocator to sanitizer_common

git-svn-id: https://llvm.org/svn/llvm-project/compiler-rt/trunk@162663 91177308-0d34-0410-b5e6-96231b3b80d8

6 files changed, 50 insertions, 31 deletions

diff --git a/lib/asan/asan_internal.h b/lib/asan/asan_internal.h
index d72a5569..a283c2b8 100644
--- a/lib/asan/asan_internal.h
+++ b/lib/asan/asan_internal.h
@@ -141,8 +141,6 @@ extern int asan_inited;
 extern bool asan_init_is_running;
 extern void (*death_callback)(void);
 
-enum LinkerInitialized { LINKER_INITIALIZED = 0 };
-
 #define ASAN_ARRAY_SIZE(a) (sizeof(a)/sizeof((a)[0]))
 
 #if !defined(_WIN32) || defined(__clang__)
@@ -177,19 +175,6 @@ const int kAsanInternalHeapMagic = 0xfe;
 static const uptr kCurrentStackFrameMagic = 0x41B58AB3;
 static const uptr kRetiredStackFrameMagic = 0x45E0360E;
 
-// -------------------------- LowLevelAllocator ----- {{{1
-// A simple low-level memory allocator for internal use.
-class LowLevelAllocator {
- public:
-  explicit LowLevelAllocator(LinkerInitialized) {}
-  // 'size' must be a power of two.
-  // Requires an external lock.
-  void *Allocate(uptr size);
- private:
-  char *allocated_end_;
-  char *allocated_current_;
-};
-
 }  // namespace __asan
 
 #endif  // ASAN_INTERNAL_H
diff --git a/lib/asan/asan_rtl.cc b/lib/asan/asan_rtl.cc
index e00cc5f4..853a844b 100644
--- a/lib/asan/asan_rtl.cc
+++ b/lib/asan/asan_rtl.cc
@@ -172,21 +172,9 @@ static void ReserveShadowMemoryRange(uptr beg, uptr end) {
   CHECK(res == (void*)beg && "ReserveShadowMemoryRange failed");
 }
 
-// ---------------------- LowLevelAllocator ------------- {{{1
-void *LowLevelAllocator::Allocate(uptr size) {
-  CHECK((size & (size - 1)) == 0 && "size must be a power of two");
-  if (allocated_end_ - allocated_current_ < (sptr)size) {
-    uptr size_to_allocate = Max(size, kPageSize);
-    allocated_current_ =
-        (char*)MmapOrDie(size_to_allocate, __FUNCTION__);
-    allocated_end_ = allocated_current_ + size_to_allocate;
-    PoisonShadow((uptr)allocated_current_, size_to_allocate,
-                 kAsanInternalHeapMagic);
-  }
-  CHECK(allocated_end_ - allocated_current_ >= (sptr)size);
-  void *res = allocated_current_;
-  allocated_current_ += size;
-  return res;
+// --------------- LowLevelAllocateCallbac ---------- {{{1
+static void OnLowLevelAllocate(uptr ptr, uptr size) {
+  PoisonShadow(ptr, size, kAsanInternalHeapMagic);
 }
 
 // -------------------------- Run-time entry ------------------- {{{1
@@ -290,8 +278,12 @@ void NOINLINE __asan_set_death_callback(void (*callback)(void)) {
 
 void __asan_init() {
   if (asan_inited) return;
+  CHECK(!asan_init_is_running && "ASan init calls itself!");
   asan_init_is_running = true;
 
+  // Setup internal allocator callback.
+  SetLowLevelAllocateCallback(OnLowLevelAllocate);
+
   // Make sure we are not statically linked.
   AsanDoesNotSupportStaticLinkage();
 
diff --git a/lib/asan/asan_thread_registry.cc b/lib/asan/asan_thread_registry.cc
index 4540d589..7c9747bd 100644
--- a/lib/asan/asan_thread_registry.cc
+++ b/lib/asan/asan_thread_registry.cc
@@ -20,7 +20,7 @@
 
 namespace __asan {
 
-static AsanThreadRegistry asan_thread_registry(__asan::LINKER_INITIALIZED);
+static AsanThreadRegistry asan_thread_registry(LINKER_INITIALIZED);
 
 AsanThreadRegistry &asanThreadRegistry() {
   return asan_thread_registry;
diff --git a/lib/sanitizer_common/sanitizer_allocator.cc b/lib/sanitizer_common/sanitizer_allocator.cc
index 816fddf1..b08434ad 100644
--- a/lib/sanitizer_common/sanitizer_allocator.cc
+++ b/lib/sanitizer_common/sanitizer_allocator.cc
@@ -56,4 +56,29 @@ void *InternalAllocBlock(void *p) {
   return pp + 1;
 }
 
+// LowLevelAllocator
+static LowLevelAllocateCallback low_level_alloc_callback;
+
+void *LowLevelAllocator::Allocate(uptr size) {
+  CHECK((size & (size - 1)) == 0 && "size must be a power of two");
+  if (allocated_end_ - allocated_current_ < (sptr)size) {
+    uptr size_to_allocate = Max(size, kPageSize);
+    allocated_current_ =
+        (char*)MmapOrDie(size_to_allocate, __FUNCTION__);
+    allocated_end_ = allocated_current_ + size_to_allocate;
+    if (low_level_alloc_callback) {
+      low_level_alloc_callback((uptr)allocated_current_,
+                               size_to_allocate);
+    }
+  }
+  CHECK(allocated_end_ - allocated_current_ >= (sptr)size);
+  void *res = allocated_current_;
+  allocated_current_ += size;
+  return res;
+}
+
+void SetLowLevelAllocateCallback(LowLevelAllocateCallback callback) {
+  low_level_alloc_callback = callback;
+}
+
 }  // namespace __sanitizer
diff --git a/lib/sanitizer_common/sanitizer_common.h b/lib/sanitizer_common/sanitizer_common.h
index 0e9371fa..d5b8f623 100644
--- a/lib/sanitizer_common/sanitizer_common.h
+++ b/lib/sanitizer_common/sanitizer_common.h
@@ -78,6 +78,21 @@ class InternalScopedBuffer {
   void operator=(const InternalScopedBuffer&);
 };
 
+// Simple low-level (mmap-based) allocator for internal use.
+class LowLevelAllocator {
+ public:
+  explicit LowLevelAllocator(LinkerInitialized) {}
+  // 'size' must be a power of two. Requires an external lock.
+  void *Allocate(uptr size);
+ private:
+  char *allocated_end_;
+  char *allocated_current_;
+};
+typedef void (*LowLevelAllocateCallback)(uptr ptr, uptr size);
+// Allows to register tool-specific callbacks for LowLevelAllocator.
+// Passing NULL removes the callback.
+void SetLowLevelAllocateCallback(LowLevelAllocateCallback callback);
+
 // IO
 void RawWrite(const char *buffer);
 void Printf(const char *format, ...);
diff --git a/lib/sanitizer_common/sanitizer_internal_defs.h b/lib/sanitizer_common/sanitizer_internal_defs.h
index b8cf61fa..a68cab25 100644
--- a/lib/sanitizer_common/sanitizer_internal_defs.h
+++ b/lib/sanitizer_common/sanitizer_internal_defs.h
@@ -160,4 +160,6 @@ void NORETURN CheckFailed(const char *file, int line, const char *cond,
 #undef UINT64_MAX
 #define UINT64_MAX             (__UINT64_C(18446744073709551615))
 
+enum LinkerInitialized { LINKER_INITIALIZED = 0 };
+
 #endif  // SANITIZER_DEFS_H