[asan] added internal flag mmap_limit_mb

llvm-svn: 206178

7 files changed, 60 insertions, 0 deletions

diff --git a/compiler-rt/lib/asan/asan_rtl.cc b/compiler-rt/lib/asan/asan_rtl.cc
index 0ac805d2ea0..298baadb05f 100644
--- a/compiler-rt/lib/asan/asan_rtl.cc
+++ b/compiler-rt/lib/asan/asan_rtl.cc
@@ -328,6 +328,7 @@ static void ReserveShadowMemoryRange(uptr beg, uptr end) {
   CHECK_EQ((beg % GetPageSizeCached()), 0);
   CHECK_EQ(((end + 1) % GetPageSizeCached()), 0);
   uptr size = end - beg + 1;
+  DecreaseTotalMmap(size);  // Don't count the shadow against mmap_limit_mb.
   void *res = MmapFixedNoReserve(beg, size);
   if (res != (void*)beg) {
     Report("ReserveShadowMemoryRange failed while trying to map 0x%zx bytes. "
diff --git a/compiler-rt/lib/sanitizer_common/sanitizer_common.cc b/compiler-rt/lib/sanitizer_common/sanitizer_common.cc
index 2dcfd0ccc00..b59f1b823e4 100644
--- a/compiler-rt/lib/sanitizer_common/sanitizer_common.cc
+++ b/compiler-rt/lib/sanitizer_common/sanitizer_common.cc
@@ -244,6 +244,24 @@ char *StripModuleName(const char *module) {
   return internal_strdup(short_module_name);
 }
 
+static atomic_uintptr_t g_total_mmaped;
+
+void IncreaseTotalMmap(uptr size) {
+  if (!common_flags()->mmap_limit_mb) return;
+  uptr total_mmaped =
+      atomic_fetch_add(&g_total_mmaped, size, memory_order_relaxed) + size;
+  if ((total_mmaped >> 20) > common_flags()->mmap_limit_mb) {
+    // Since for now mmap_limit_mb is not a user-facing flag, just CHECK.
+    common_flags()->mmap_limit_mb = 0;  // Allow mmap in CHECK.
+    CHECK_LT(total_mmaped >> 20, common_flags()->mmap_limit_mb);
+  }
+}
+
+void DecreaseTotalMmap(uptr size) {
+  if (!common_flags()->mmap_limit_mb) return;
+  atomic_fetch_sub(&g_total_mmaped, size, memory_order_relaxed);
+}
+
 }  // namespace __sanitizer
 
 using namespace __sanitizer;  // NOLINT
diff --git a/compiler-rt/lib/sanitizer_common/sanitizer_common.h b/compiler-rt/lib/sanitizer_common/sanitizer_common.h
index fe79342acaa..00c686ead64 100644
--- a/compiler-rt/lib/sanitizer_common/sanitizer_common.h
+++ b/compiler-rt/lib/sanitizer_common/sanitizer_common.h
@@ -60,6 +60,8 @@ void *MmapAlignedOrDie(uptr size, uptr alignment, const char *mem_type);
 // Used to check if we can map shadow memory to a fixed location.
 bool MemoryRangeIsAvailable(uptr range_start, uptr range_end);
 void FlushUnneededShadowMemory(uptr addr, uptr size);
+void IncreaseTotalMmap(uptr size);
+void DecreaseTotalMmap(uptr size);
 
 // InternalScopedBuffer can be used instead of large stack arrays to
 // keep frame size low.
diff --git a/compiler-rt/lib/sanitizer_common/sanitizer_flags.cc b/compiler-rt/lib/sanitizer_common/sanitizer_flags.cc
index e44d04f5577..78b9747d808 100644
--- a/compiler-rt/lib/sanitizer_common/sanitizer_flags.cc
+++ b/compiler-rt/lib/sanitizer_common/sanitizer_flags.cc
@@ -119,6 +119,9 @@ void ParseCommonFlagsFromString(CommonFlags *f, const char *str) {
   ParseFlag(str, &f->intercept_tls_get_addr, "intercept_tls_get_addr",
             "Intercept __tls_get_addr.");
   ParseFlag(str, &f->help, "help", "Print the flag descriptions.");
+  ParseFlag(str, &f->mmap_limit_mb, "mmap_limit_mb",
+            "Limit the amount of mmap-ed memory (excluding shadow) in Mb; "
+            "not a user-facing flag, used mosly for testing the tools");
 
   // Do a sanity check for certain flags.
   if (f->malloc_context_size < 1)
diff --git a/compiler-rt/lib/sanitizer_common/sanitizer_flags.h b/compiler-rt/lib/sanitizer_common/sanitizer_flags.h
index 25dd887a801..a322eadebbc 100644
--- a/compiler-rt/lib/sanitizer_common/sanitizer_flags.h
+++ b/compiler-rt/lib/sanitizer_common/sanitizer_flags.h
@@ -52,6 +52,7 @@ struct CommonFlags {
   bool legacy_pthread_cond;
   bool intercept_tls_get_addr;
   bool help;
+  uptr mmap_limit_mb;
 };
 
 inline CommonFlags *common_flags() {
diff --git a/compiler-rt/lib/sanitizer_common/sanitizer_posix.cc b/compiler-rt/lib/sanitizer_common/sanitizer_posix.cc
index f3e8f385539..736701d3ad1 100644
--- a/compiler-rt/lib/sanitizer_common/sanitizer_posix.cc
+++ b/compiler-rt/lib/sanitizer_common/sanitizer_posix.cc
@@ -91,6 +91,7 @@ void *MmapOrDie(uptr size, const char *mem_type) {
     DumpProcessMap();
     CHECK("unable to mmap" && 0);
   }
+  IncreaseTotalMmap(size);
   return (void *)res;
 }
 
@@ -102,6 +103,7 @@ void UnmapOrDie(void *addr, uptr size) {
            SanitizerToolName, size, size, addr);
     CHECK("unable to unmap" && 0);
   }
+  DecreaseTotalMmap(size);
 }
 
 void *MmapNoReserveOrDie(uptr size, const char *mem_type) {
@@ -118,6 +120,7 @@ void *MmapNoReserveOrDie(uptr size, const char *mem_type) {
            SanitizerToolName, size, size, mem_type, reserrno);
     CHECK("unable to mmap" && 0);
   }
+  IncreaseTotalMmap(size);
   return (void *)p;
 }
 
@@ -133,6 +136,7 @@ void *MmapFixedNoReserve(uptr fixed_addr, uptr size) {
     Report("ERROR: %s failed to "
            "allocate 0x%zx (%zd) bytes at address %zu (errno: %d)\n",
            SanitizerToolName, size, size, fixed_addr, reserrno);
+  IncreaseTotalMmap(size);
   return (void *)p;
 }
 
@@ -150,6 +154,7 @@ void *MmapFixedOrDie(uptr fixed_addr, uptr size) {
            SanitizerToolName, size, size, fixed_addr, reserrno);
     CHECK("unable to mmap" && 0);
   }
+  IncreaseTotalMmap(size);
   return (void *)p;
 }
 
diff --git a/compiler-rt/test/asan/TestCases/mmap_limit_mb.cc b/compiler-rt/test/asan/TestCases/mmap_limit_mb.cc
new file mode 100644
index 00000000000..cd8783b5017
--- /dev/null
+++ b/compiler-rt/test/asan/TestCases/mmap_limit_mb.cc
@@ -0,0 +1,30 @@
+// Test the mmap_limit_mb flag.
+//
+// RUN: %clangxx_asan -std=c++11 -O2 %s -o %t
+// RUN: %t 100 16
+// RUN: %t 100 1000000
+// RUN: ASAN_OPTIONS=mmap_limit_mb=500 %t 100 16
+// RUN: ASAN_OPTIONS=mmap_limit_mb=500 %t 100 1000000
+// RUN: ASAN_OPTIONS=mmap_limit_mb=500 not %t 500 16 2>&1 | FileCheck %s
+// RUN: ASAN_OPTIONS=mmap_limit_mb=500 not %t 500 1000000 2>&1 | FileCheck %s
+
+#include <assert.h>
+#include <stdlib.h>
+#include <stdio.h>
+
+#include <algorithm>
+#include <vector>
+
+int main(int argc, char **argv) {
+  assert(argc == 3);
+  long total_mb = atoi(argv[1]);
+  long allocation_size = atoi(argv[2]);
+  std::vector<char *> v;
+  for (long total = total_mb << 20; total > 0; total -= allocation_size)
+    v.push_back(new char[allocation_size]);
+  std::for_each(v.begin(), v.end(),
+                [](std::vector<char *>::reference ref) { delete[] ref; });
+  printf("PASS\n");
+  // CHECK: AddressSanitizer CHECK failed{{.*}}total_mmaped{{.*}}mmap_limit_mb
+  return 0;
+}