[sanitizer] make the allocator crash instead of returning 0 on huge size (controlled by the allocator_may_return_null flag)

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

8 files changed, 94 insertions, 23 deletions

diff --git a/lib/asan/asan_allocator2.cc b/lib/asan/asan_allocator2.cc
index 97c992970..5d13b91c3 100644
--- a/lib/asan/asan_allocator2.cc
+++ b/lib/asan/asan_allocator2.cc
@@ -345,7 +345,7 @@ static void *Allocate(uptr size, uptr alignment, StackTrace *stack,
   if (size > kMaxAllowedMallocSize || needed_size > kMaxAllowedMallocSize) {
     Report("WARNING: AddressSanitizer failed to allocate %p bytes\n",
            (void*)size);
-    return 0;
+    return AllocatorReturnNull();
   }
 
   AsanThread *t = GetCurrentThread();
diff --git a/lib/asan/lit_tests/TestCases/allocator_returns_null.cc b/lib/asan/lit_tests/TestCases/allocator_returns_null.cc
new file mode 100644
index 000000000..a459f24e1
--- /dev/null
+++ b/lib/asan/lit_tests/TestCases/allocator_returns_null.cc
@@ -0,0 +1,65 @@
+// Test the behavior of malloc/calloc/realloc when the allocation size is huge.
+// By default (allocator_may_return_null=0) the process shoudl crash.
+// With allocator_may_return_null=1 the allocator should return 0.
+//
+// RUN: %clangxx_asan -O0 %s -o %t
+// RUN: not %t malloc 2>&1 | FileCheck %s --check-prefix=CHECK-mCRASH
+// RUN: ASAN_OPTIONS=allocator_may_return_null=0 not %t malloc 2>&1 | FileCheck %s --check-prefix=CHECK-mCRASH
+// RUN: ASAN_OPTIONS=allocator_may_return_null=1     %t malloc 2>&1 | FileCheck %s --check-prefix=CHECK-mNULL
+// RUN: ASAN_OPTIONS=allocator_may_return_null=0 not %t calloc 2>&1 | FileCheck %s --check-prefix=CHECK-cCRASH
+// RUN: ASAN_OPTIONS=allocator_may_return_null=1     %t calloc 2>&1 | FileCheck %s --check-prefix=CHECK-cNULL
+// RUN: ASAN_OPTIONS=allocator_may_return_null=0 not %t realloc 2>&1 | FileCheck %s --check-prefix=CHECK-rCRASH
+// RUN: ASAN_OPTIONS=allocator_may_return_null=1     %t realloc 2>&1 | FileCheck %s --check-prefix=CHECK-rNULL
+// RUN: ASAN_OPTIONS=allocator_may_return_null=0 not %t realloc-after-malloc 2>&1 | FileCheck %s --check-prefix=CHECK-mrCRASH
+// RUN: ASAN_OPTIONS=allocator_may_return_null=1     %t realloc-after-malloc 2>&1 | FileCheck %s --check-prefix=CHECK-mrNULL
+
+#include <limits.h>
+#include <stdlib.h>
+#include <string.h>
+#include <stdio.h>
+#include <assert.h>
+#include <limits>
+int main(int argc, char **argv) {
+  volatile size_t size = std::numeric_limits<size_t>::max() - 10000;
+  assert(argc == 2);
+  char *x = 0;
+  if (!strcmp(argv[1], "malloc")) {
+    fprintf(stderr, "malloc:\n");
+    x = (char*)malloc(size);
+  }
+  if (!strcmp(argv[1], "calloc")) {
+    fprintf(stderr, "calloc:\n");
+    x = (char*)calloc(size / 4, 4);
+  }
+
+  if (!strcmp(argv[1], "realloc")) {
+    fprintf(stderr, "realloc:\n");
+    x = (char*)realloc(0, size);
+  }
+  if (!strcmp(argv[1], "realloc-after-malloc")) {
+    fprintf(stderr, "realloc-after-malloc:\n");
+    char *t = (char*)malloc(100);
+    *t = 42;
+    x = (char*)realloc(t, size);
+    assert(*t == 42);
+  }
+  fprintf(stderr, "x: %p\n", x);
+  return x != 0;
+}
+// CHECK-mCRASH: malloc:
+// CHECK-mCRASH: AddressSanitizer's allocator is terminating the process
+// CHECK-cCRASH: calloc:
+// CHECK-cCRASH: AddressSanitizer's allocator is terminating the process
+// CHECK-rCRASH: realloc:
+// CHECK-rCRASH: AddressSanitizer's allocator is terminating the process
+// CHECK-mrCRASH: realloc-after-malloc:
+// CHECK-mrCRASH: AddressSanitizer's allocator is terminating the process
+
+// CHECK-mNULL: malloc:
+// CHECK-mNULL: x: (nil)
+// CHECK-cNULL: calloc:
+// CHECK-cNULL: x: (nil)
+// CHECK-rNULL: realloc:
+// CHECK-rNULL: x: (nil)
+// CHECK-mrNULL: realloc-after-malloc:
+// CHECK-mrNULL: x: (nil)
diff --git a/lib/asan/tests/asan_test.cc b/lib/asan/tests/asan_test.cc
index c87e6da58..c216b6839 100644
--- a/lib/asan/tests/asan_test.cc
+++ b/lib/asan/tests/asan_test.cc
@@ -227,26 +227,6 @@ TEST(AddressSanitizer, BitFieldNegativeTest) {
   delete Ident(x);
 }
 
-static size_t kOOMAllocationSize =
-  SANITIZER_WORDSIZE == 64 ? (size_t)(1ULL << 48) : (0xf0000000);
-
-TEST(AddressSanitizer, OutOfMemoryTest) {
-  EXPECT_EQ(0, realloc(0, kOOMAllocationSize));
-  EXPECT_EQ(0, realloc(0, ~Ident(0)));
-  EXPECT_EQ(0, malloc(kOOMAllocationSize));
-  EXPECT_EQ(0, malloc(~Ident(0)));
-  EXPECT_EQ(0, calloc(1, kOOMAllocationSize));
-  EXPECT_EQ(0, calloc(1, ~Ident(0)));
-}
-
-TEST(AddressSanitizer, BadReallocTest) {
-  int *a = (int*)Ident(malloc(100));
-  a[0] = 42;
-  EXPECT_EQ(0, realloc(a, kOOMAllocationSize));
-  EXPECT_EQ(42, a[0]);
-  free(a);
-}
-
 #if ASAN_NEEDS_SEGV
 namespace {
 
diff --git a/lib/sanitizer_common/sanitizer_allocator.cc b/lib/sanitizer_common/sanitizer_allocator.cc
index 2975073cf..daaf7e1ce 100644
--- a/lib/sanitizer_common/sanitizer_allocator.cc
+++ b/lib/sanitizer_common/sanitizer_allocator.cc
@@ -14,6 +14,7 @@
 #include "sanitizer_allocator.h"
 #include "sanitizer_allocator_internal.h"
 #include "sanitizer_common.h"
+#include "sanitizer_flags.h"
 
 namespace __sanitizer {
 
@@ -139,4 +140,14 @@ bool CallocShouldReturnNullDueToOverflow(uptr size, uptr n) {
   return (max / size) < n;
 }
 
+void *AllocatorReturnNull() {
+  if (common_flags()->allocator_may_return_null)
+    return 0;
+  Report("%s's allocator is terminating the process instead of returning 0\n",
+         SanitizerToolName);
+  Report("If you don't like this behavior set allocator_may_return_null=1\n");
+  CHECK(0);
+  return 0;
+}
+
 }  // namespace __sanitizer
diff --git a/lib/sanitizer_common/sanitizer_allocator.h b/lib/sanitizer_common/sanitizer_allocator.h
index 72b76dcda..9fdc39ba3 100644
--- a/lib/sanitizer_common/sanitizer_allocator.h
+++ b/lib/sanitizer_common/sanitizer_allocator.h
@@ -23,6 +23,9 @@
 
 namespace __sanitizer {
 
+// Depending on allocator_may_return_null either return 0 or crash.
+void *AllocatorReturnNull();
+
 // SizeClassMap maps allocation sizes into size classes and back.
 // Class 0 corresponds to size 0.
 // Classes 1 - 16 correspond to sizes 16 to 256 (size = class_id * 16).
@@ -941,7 +944,7 @@ class LargeMmapAllocator {
     uptr map_size = RoundUpMapSize(size);
     if (alignment > page_size_)
       map_size += alignment;
-    if (map_size < size) return 0;  // Overflow.
+    if (map_size < size) return AllocatorReturnNull();  // Overflow.
     uptr map_beg = reinterpret_cast<uptr>(
         MmapOrDie(map_size, "LargeMmapAllocator"));
     MapUnmapCallback().OnMap(map_beg, map_size);
@@ -1176,7 +1179,7 @@ class CombinedAllocator {
     if (size == 0)
       size = 1;
     if (size + alignment < size)
-      return 0;
+      return AllocatorReturnNull();
     if (alignment > 8)
       size = RoundUpTo(size, alignment);
     void *res;
diff --git a/lib/sanitizer_common/sanitizer_flags.cc b/lib/sanitizer_common/sanitizer_flags.cc
index 6ca3247b1..734226167 100644
--- a/lib/sanitizer_common/sanitizer_flags.cc
+++ b/lib/sanitizer_common/sanitizer_flags.cc
@@ -31,6 +31,7 @@ void ParseCommonFlagsFromString(const char *str) {
   ParseFlag(str, &f->log_path, "log_path");
   ParseFlag(str, &f->detect_leaks, "detect_leaks");
   ParseFlag(str, &f->leak_check_at_exit, "leak_check_at_exit");
+  ParseFlag(str, &f->allocator_may_return_null, "allocator_may_return_null");
 }
 
 static bool GetFlagValue(const char *env, const char *name,
diff --git a/lib/sanitizer_common/sanitizer_flags.h b/lib/sanitizer_common/sanitizer_flags.h
index d1a2961b6..43632e65c 100644
--- a/lib/sanitizer_common/sanitizer_flags.h
+++ b/lib/sanitizer_common/sanitizer_flags.h
@@ -45,6 +45,8 @@ struct CommonFlags {
   // detect_leaks=false, or if __lsan_do_leak_check() is called before the
   // handler has a chance to run.
   bool leak_check_at_exit;
+  // If false, the allocator will crash instead of returning 0 on out-of-memory.
+  bool allocator_may_return_null;
 };
 
 extern CommonFlags common_flags_dont_use_directly;
diff --git a/lib/sanitizer_common/tests/sanitizer_allocator_test.cc b/lib/sanitizer_common/tests/sanitizer_allocator_test.cc
index 557420673..3541d439d 100644
--- a/lib/sanitizer_common/tests/sanitizer_allocator_test.cc
+++ b/lib/sanitizer_common/tests/sanitizer_allocator_test.cc
@@ -14,6 +14,7 @@
 #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"
 
@@ -416,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++) {