[sanitizer_common] Create max_allocation_size_mb flag.

Summary:
The flag allows the user to specify a maximum allocation size that the
sanitizers will honor.  Any larger allocations will return nullptr or
crash depending on allocator_may_return_null.

Reviewers: kcc, eugenis

Reviewed By: kcc, eugenis

Subscribers: #sanitizers, llvm-commits

Tags: #sanitizers, #llvm

Differential Revision: https://reviews.llvm.org/D69576

compiler-rt/lib/asan/asan_allocator.cpp

@@ -246,6 +246,7 @@ struct Allocator {
   AllocatorCache fallback_allocator_cache;
   QuarantineCache fallback_quarantine_cache;
 
+  uptr max_user_defined_malloc_size;
   atomic_uint8_t rss_limit_exceeded;
 
   // ------------------- Options --------------------------
@@ -280,6 +281,10 @@ struct Allocator {
     SetAllocatorMayReturnNull(options.may_return_null);
     allocator.InitLinkerInitialized(options.release_to_os_interval_ms);
     SharedInitCode(options);
+    max_user_defined_malloc_size = common_flags()->max_allocation_size_mb
+                                       ? common_flags()->max_allocation_size_mb
+                                             << 20
+                                       : kMaxAllowedMallocSize;
   }
 
   bool RssLimitExceeded() {
@@ -435,14 +440,16 @@ struct Allocator {
       using_primary_allocator = false;
     }
     CHECK(IsAligned(needed_size, min_alignment));
-    if (size > kMaxAllowedMallocSize || needed_size > kMaxAllowedMallocSize) {
+    if (size > kMaxAllowedMallocSize || needed_size > kMaxAllowedMallocSize ||
+        size > max_user_defined_malloc_size) {
       if (AllocatorMayReturnNull()) {
         Report("WARNING: AddressSanitizer failed to allocate 0x%zx bytes\n",
                (void*)size);
         return nullptr;
       }
-      ReportAllocationSizeTooBig(size, needed_size, kMaxAllowedMallocSize,
-                                 stack);
+      uptr malloc_limit =
+          Min(kMaxAllowedMallocSize, max_user_defined_malloc_size);
+      ReportAllocationSizeTooBig(size, needed_size, malloc_limit, stack);
     }
 
     AsanThread *t = GetCurrentThread();

compiler-rt/lib/lsan/lsan_allocator.cpp

@@ -36,10 +36,17 @@ static const uptr kMaxAllowedMallocSize = 8UL << 30;
 
 static Allocator allocator;
 
+static uptr max_malloc_size;
+
 void InitializeAllocator() {
   SetAllocatorMayReturnNull(common_flags()->allocator_may_return_null);
   allocator.InitLinkerInitialized(
       common_flags()->allocator_release_to_os_interval_ms);
+  if (common_flags()->max_allocation_size_mb)
+    max_malloc_size = Min(common_flags()->max_allocation_size_mb << 20,
+                          kMaxAllowedMallocSize);
+  else
+    max_malloc_size = kMaxAllowedMallocSize;
 }
 
 void AllocatorThreadFinish() {
@@ -72,14 +79,14 @@ static void *ReportAllocationSizeTooBig(uptr size, const StackTrace &stack) {
     Report("WARNING: LeakSanitizer failed to allocate 0x%zx bytes\n", size);
     return nullptr;
   }
-  ReportAllocationSizeTooBig(size, kMaxAllowedMallocSize, &stack);
+  ReportAllocationSizeTooBig(size, max_malloc_size, &stack);
 }
 
 void *Allocate(const StackTrace &stack, uptr size, uptr alignment,
                bool cleared) {
   if (size == 0)
     size = 1;
-  if (size > kMaxAllowedMallocSize)
+  if (size > max_malloc_size)
     return ReportAllocationSizeTooBig(size, stack);
   void *p = allocator.Allocate(GetAllocatorCache(), size, alignment);
   if (UNLIKELY(!p)) {
@@ -117,7 +124,7 @@ void Deallocate(void *p) {
 void *Reallocate(const StackTrace &stack, void *p, uptr new_size,
                  uptr alignment) {
   RegisterDeallocation(p);
-  if (new_size > kMaxAllowedMallocSize) {
+  if (new_size > max_malloc_size) {
     allocator.Deallocate(GetAllocatorCache(), p);
     return ReportAllocationSizeTooBig(new_size, stack);
   }

compiler-rt/lib/msan/msan_allocator.cpp

@@ -115,9 +115,16 @@ static Allocator allocator;
 static AllocatorCache fallback_allocator_cache;
 static StaticSpinMutex fallback_mutex;
 
+static uptr max_malloc_size;
+
 void MsanAllocatorInit() {
   SetAllocatorMayReturnNull(common_flags()->allocator_may_return_null);
   allocator.Init(common_flags()->allocator_release_to_os_interval_ms);
+  if (common_flags()->max_allocation_size_mb)
+    max_malloc_size = Min(common_flags()->max_allocation_size_mb << 20,
+                          kMaxAllowedMallocSize);
+  else
+    max_malloc_size = kMaxAllowedMallocSize;
 }
 
 AllocatorCache *GetAllocatorCache(MsanThreadLocalMallocStorage *ms) {
@@ -132,12 +139,12 @@ void MsanThreadLocalMallocStorage::CommitBack() {
 
 static void *MsanAllocate(StackTrace *stack, uptr size, uptr alignment,
                           bool zeroise) {
-  if (size > kMaxAllowedMallocSize) {
+  if (size > max_malloc_size) {
     if (AllocatorMayReturnNull()) {
       Report("WARNING: MemorySanitizer failed to allocate 0x%zx bytes\n", size);
       return nullptr;
     }
-    ReportAllocationSizeTooBig(size, kMaxAllowedMallocSize, stack);
+    ReportAllocationSizeTooBig(size, max_malloc_size, stack);
   }
   MsanThread *t = GetCurrentThread();
   void *allocated;

compiler-rt/lib/sanitizer_common/sanitizer_flags.inc

@@ -132,6 +132,9 @@ COMMON_FLAG(uptr, soft_rss_limit_mb, 0,
             " until the RSS goes below the soft limit."
             " This limit does not affect memory allocations other than"
             " malloc/new.")
+COMMON_FLAG(uptr, max_allocation_size_mb, 0,
+            "If non-zero, malloc/new calls larger than this size will return "
+            "nullptr (or crash if allocator_may_return_null=false).")
 COMMON_FLAG(bool, heap_profile, false, "Experimental heap profiler, asan-only")
 COMMON_FLAG(s32, allocator_release_to_os_interval_ms,
             ((bool)SANITIZER_FUCHSIA || (bool)SANITIZER_WINDOWS) ? -1 : 5000,

compiler-rt/lib/tsan/rtl/tsan_mman.cpp

@@ -113,9 +113,16 @@ ScopedGlobalProcessor::~ScopedGlobalProcessor() {
   gp->mtx.Unlock();
 }
 
+static constexpr uptr kMaxAllowedMallocSize = 1ull << 40;
+static uptr max_user_defined_malloc_size;
+
 void InitializeAllocator() {
   SetAllocatorMayReturnNull(common_flags()->allocator_may_return_null);
   allocator()->Init(common_flags()->allocator_release_to_os_interval_ms);
+  max_user_defined_malloc_size = common_flags()->max_allocation_size_mb
+                                     ? common_flags()->max_allocation_size_mb
+                                           << 20
+                                     : kMaxAllowedMallocSize;
 }
 
 void InitializeAllocatorLate() {
@@ -150,15 +157,17 @@ static void SignalUnsafeCall(ThreadState *thr, uptr pc) {
   OutputReport(thr, rep);
 }
 
-static constexpr uptr kMaxAllowedMallocSize = 1ull << 40;
 
 void *user_alloc_internal(ThreadState *thr, uptr pc, uptr sz, uptr align,
                           bool signal) {
-  if (sz >= kMaxAllowedMallocSize || align >= kMaxAllowedMallocSize) {
+  if (sz >= kMaxAllowedMallocSize || align >= kMaxAllowedMallocSize ||
+      sz > max_user_defined_malloc_size) {
     if (AllocatorMayReturnNull())
       return nullptr;
+    uptr malloc_limit =
+        Min(kMaxAllowedMallocSize, max_user_defined_malloc_size);
     GET_STACK_TRACE_FATAL(thr, pc);
-    ReportAllocationSizeTooBig(sz, kMaxAllowedMallocSize, &stack);
+    ReportAllocationSizeTooBig(sz, malloc_limit, &stack);
   }
   void *p = allocator()->Allocate(&thr->proc()->alloc_cache, sz, align);
   if (UNLIKELY(!p)) {

compiler-rt/test/sanitizer_common/TestCases/max_allocation_size.cpp

@@ -0,0 +1,127 @@
+// Test the behavior of malloc/calloc/realloc/new when the allocation size
+// exceeds the configured max_allocation_size_mb flag.
+// By default (allocator_may_return_null=0) the process should crash. With
+// allocator_may_return_null=1 the allocator should return nullptr and set errno
+// to the appropriate error code.
+//
+// RUN: %clangxx -O0 %s -o %t
+// RUN: %run %t malloc 2>&1 | FileCheck %s --check-prefix=CHECK-NOTNULL
+// RUN: %env_tool_opts=max_allocation_size_mb=3 %run %t malloc 2>&1 \
+// RUN:   | FileCheck %s --check-prefix=CHECK-NOTNULL
+// RUN: %env_tool_opts=max_allocation_size_mb=2:allocator_may_return_null=0 \
+// RUN:   not %run %t malloc 2>&1 | FileCheck %s --check-prefix=CHECK-mCRASH
+// RUN: %env_tool_opts=max_allocation_size_mb=2:allocator_may_return_null=1 \
+// RUN:   %run %t malloc 2>&1 | FileCheck %s --check-prefix=CHECK-NULL
+// RUN: %env_tool_opts=max_allocation_size_mb=2:allocator_may_return_null=0 \
+// RUN:   not %run %t calloc 2>&1 | FileCheck %s --check-prefix=CHECK-cCRASH
+// RUN: %env_tool_opts=max_allocation_size_mb=2:allocator_may_return_null=1 \
+// RUN:   %run %t calloc 2>&1 | FileCheck %s --check-prefix=CHECK-NULL
+// RUN: %env_tool_opts=max_allocation_size_mb=2:allocator_may_return_null=0 \
+// RUN:   not %run %t realloc 2>&1 | FileCheck %s --check-prefix=CHECK-rCRASH
+// RUN: %env_tool_opts=max_allocation_size_mb=2:allocator_may_return_null=1 \
+// RUN:   %run %t realloc 2>&1 | FileCheck %s --check-prefix=CHECK-NULL
+// RUN: %env_tool_opts=max_allocation_size_mb=2:allocator_may_return_null=0 \
+// RUN:   not %run %t realloc-after-malloc 2>&1 \
+// RUN:   | FileCheck %s --check-prefix=CHECK-mrCRASH
+// RUN: %env_tool_opts=max_allocation_size_mb=2:allocator_may_return_null=1 \
+// RUN:   %run %t realloc-after-malloc 2>&1 \
+// RUN:   | FileCheck %s --check-prefix=CHECK-NULL
+// RUN: %env_tool_opts=max_allocation_size_mb=2:allocator_may_return_null=0 \
+// RUN:   not %run %t new 2>&1 | FileCheck %s --check-prefix=CHECK-nCRASH
+// RUN: %env_tool_opts=max_allocation_size_mb=2:allocator_may_return_null=1 \
+// RUN:   not %run %t new 2>&1 | FileCheck %s --check-prefix=CHECK-nCRASH-OOM
+// RUN: %env_tool_opts=max_allocation_size_mb=2:allocator_may_return_null=0 \
+// RUN:   not %run %t new-nothrow 2>&1 \
+// RUN:   | FileCheck %s --check-prefix=CHECK-nnCRASH
+// RUN: %env_tool_opts=max_allocation_size_mb=2:allocator_may_return_null=1 \
+// RUN:   %run %t new-nothrow 2>&1 | FileCheck %s --check-prefix=CHECK-NULL
+
+// win32 is disabled due to failing errno tests.
+// UNSUPPORTED: ubsan, windows-msvc
+
+#include <assert.h>
+#include <errno.h>
+#include <limits>
+#include <new>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+
+static void *allocate(const char *Action, size_t Size) {
+  if (!strcmp(Action, "malloc"))
+    return malloc(Size);
+  if (!strcmp(Action, "calloc"))
+    return calloc((Size + 3) / 4, 4);
+  if (!strcmp(Action, "realloc"))
+    return realloc(nullptr, Size);
+  if (!strcmp(Action, "realloc-after-malloc")) {
+    void *P = malloc(100);
+    if (void *Ret = realloc(P, Size))
+      return Ret;
+    free(P);
+    return nullptr;
+  }
+  if (!strcmp(Action, "new"))
+    return ::operator new(Size);
+  if (!strcmp(Action, "new-nothrow"))
+    return ::operator new(Size, std::nothrow);
+  assert(0);
+}
+
+static void deallocate(const char *Action, void *Ptr) {
+  if (!strcmp(Action, "malloc") || !strcmp(Action, "calloc") ||
+      !strcmp(Action, "realloc") || !strcmp(Action, "realloc-after-malloc"))
+    return free(Ptr);
+  if (!strcmp(Action, "new"))
+    return ::operator delete(Ptr);
+  if (!strcmp(Action, "new-nothrow"))
+    return ::operator delete(Ptr, std::nothrow);
+  assert(0);
+}
+
+int main(int Argc, char **Argv) {
+  assert(Argc == 2);
+  const char *Action = Argv[1];
+  fprintf(stderr, "%s:\n", Action);
+
+  constexpr size_t MaxAllocationSize = size_t{2} << 20;
+
+  // Should succeed when max_allocation_size_mb is set.
+  void *volatile P = allocate(Action, MaxAllocationSize);
+  assert(P);
+  deallocate(Action, P);
+
+  // Should fail when max_allocation_size_mb is set.
+  P = allocate(Action, MaxAllocationSize + 1);
+  // The NULL pointer is printed differently on different systems, while (long)0
+  // is always the same.
+  fprintf(stderr, "errno: %d, P: %lx\n", errno, (long)P);
+  deallocate(Action, P);
+
+  // Should succeed when max_allocation_size_mb is set.
+  P = allocate(Action, MaxAllocationSize);
+  assert(P);
+  deallocate(Action, P);
+
+  return 0;
+}
+
+// CHECK-mCRASH: malloc:
+// CHECK-mCRASH: {{SUMMARY: .*Sanitizer: allocation-size-too-big}}
+// CHECK-cCRASH: calloc:
+// CHECK-cCRASH: {{SUMMARY: .*Sanitizer: allocation-size-too-big}}
+// CHECK-rCRASH: realloc:
+// CHECK-rCRASH: {{SUMMARY: .*Sanitizer: allocation-size-too-big}}
+// CHECK-mrCRASH: realloc-after-malloc:
+// CHECK-mrCRASH: {{SUMMARY: .*Sanitizer: allocation-size-too-big}}
+// CHECK-nCRASH: new:
+// CHECK-nCRASH: {{SUMMARY: .*Sanitizer: allocation-size-too-big}}
+// CHECK-nCRASH-OOM: new:
+// CHECK-nCRASH-OOM: {{SUMMARY: .*Sanitizer: out-of-memory}}
+// CHECK-nnCRASH: new-nothrow:
+// CHECK-nnCRASH: {{SUMMARY: .*Sanitizer: allocation-size-too-big}}
+
+// CHECK-NULL: {{malloc|calloc|calloc-overflow|realloc|realloc-after-malloc|new-nothrow}}
+// CHECK-NULL: errno: 12, P: 0
+//
+// CHECK-NOTNULL-NOT: P: 0