[compiler-rt] [asan] Plumb may_return_null through asan_memalign (PR #196413)

[Vitaly Buka via llvm-commits]

================
@@ -120,47 +120,47 @@ INTERCEPTOR(void*, calloc, uptr nmemb, uptr size) {
 // TODO: AIX needs a method to ensure 16-byte alignment if the incoming
 // pointer was allocated with a 16-byte alignment requirement (or perhaps
 // merely if it happens to have 16-byte alignment).
-INTERCEPTOR(void*, realloc, void *ptr, uptr size) {
+INTERCEPTOR(void*, realloc, void* ptr, uptr size) {
   if (DlsymAlloc::Use() || DlsymAlloc::PointerIsMine(ptr))
     return DlsymAlloc::Realloc(ptr, size);
   GET_STACK_TRACE_MALLOC;
   return asan_realloc(ptr, size, &stack);
 }
 
-#if SANITIZER_INTERCEPT_REALLOCARRAY
-INTERCEPTOR(void*, reallocarray, void *ptr, uptr nmemb, uptr size) {
+#  if SANITIZER_INTERCEPT_REALLOCARRAY
+INTERCEPTOR(void*, reallocarray, void* ptr, uptr nmemb, uptr size) {
   AsanInitFromRtl();
   GET_STACK_TRACE_MALLOC;
   return asan_reallocarray(ptr, nmemb, size, &stack);
 }
-#endif  // SANITIZER_INTERCEPT_REALLOCARRAY
+#  endif  // SANITIZER_INTERCEPT_REALLOCARRAY
 
-#if SANITIZER_INTERCEPT_MEMALIGN
+#  if SANITIZER_INTERCEPT_MEMALIGN
 INTERCEPTOR(void*, memalign, uptr boundary, uptr size) {
   GET_STACK_TRACE_MALLOC;
-  return asan_memalign(boundary, size, &stack);
+  return asan_memalign(boundary, size, &stack, AllocatorMayReturnNull());
----------------
vitalybuka wrote:

Seems like just this is needed?

```
diff --git a/compiler-rt/lib/asan/asan_allocator.cpp b/compiler-rt/lib/asan/asan_allocator.cpp
index ce8311c05c5d..497c81113eb4 100644
--- a/compiler-rt/lib/asan/asan_allocator.cpp
+++ b/compiler-rt/lib/asan/asan_allocator.cpp
@@ -536,12 +536,15 @@ struct Allocator {
   }
 
   // -------------------- Allocation/Deallocation routines ---------------
-  void* Allocate(uptr size, uptr alignment, BufferedStackTrace* stack,
-                 AllocType alloc_type, bool can_fill) {
+  void* AllocateImpl(uptr size, uptr alignment, BufferedStackTrace* stack,
+                     AllocType alloc_type, bool can_fill,
+                     bool recover_oom) {
     if (UNLIKELY(!AsanInited()))
       AsanInitFromRtl();
     if (UNLIKELY(IsRssLimitExceeded())) {
-      if (AllocatorMayReturnNull())
+      if (recover_oom || AllocatorMayReturnNull()
         return nullptr;
       ReportRssLimitExceeded(stack);
     }
@@ -578,7 +581,7 @@ struct Allocator {
     CHECK(IsAligned(needed_size, min_alignment));
     if (size > kMaxAllowedMallocSize || needed_size > kMaxAllowedMallocSize ||
         size > max_user_defined_malloc_size) {
-      if (AllocatorMayReturnNull()) {
+      if (recover_oom || AllocatorMayReturnNull()
         Report("WARNING: AddressSanitizer failed to allocate 0x%zx bytes\n",
                size);
         return nullptr;
@@ -600,7 +603,7 @@ struct Allocator {
     }
     if (UNLIKELY(!allocated)) {
       SetAllocatorOutOfMemory();
-      if (AllocatorMayReturnNull())
+      if (recover_oom || AllocatorMayReturnNull()
         return nullptr;
       ReportOutOfMemory(size, stack);
     }
@@ -679,6 +682,13 @@ struct Allocator {
     return res;
   }
 
+  // Defer to the global, flag controlled, OOM policy.
+  void* Allocate(uptr size, uptr alignment, BufferedStackTrace* stack,
+                 AllocType alloc_type, bool can_fill) {
+    return AllocateImpl(size, alignment, stack, alloc_type, can_fill,
+                        /*recover_oom=*/false);
+  }
+
   // Set quarantine flag if chunk is allocated, issue ASan error report on
   // available and quarantined chunks. Return true on success, false otherwise.
   bool AtomicallySetQuarantineFlagIfAllocated(AsanChunk* m, void* ptr,
```

https://github.com/llvm/llvm-project/pull/196413