[ASan] Honor allocator_may_return_null when set through user-function and fix large alloc edge case

compiler-rt/lib/asan/asan_flags.cpp

@@ -240,6 +240,13 @@ void InitializeFlags() {
 
         DisplayHelpMessages(&asan_parser);
         ProcessFlags();
+
+        // TODO: Update other globals and data structures that may need to change
+        // after initialization due to new flags potentially being set changing after
+        // `__asan_default_options` is registered.
+        // See GH issue 'https://github.com/llvm/llvm-project/issues/117925' for
+        // details.
+        SetAllocatorMayReturnNull(common_flags()->allocator_may_return_null);
       });
 
 #  if CAN_SANITIZE_UB

compiler-rt/lib/sanitizer_common/sanitizer_allocator.cpp

@@ -85,8 +85,12 @@ static void NORETURN ReportInternalAllocatorOutOfMemory(uptr requested_size) {
 
 void *InternalAlloc(uptr size, InternalAllocatorCache *cache, uptr alignment) {
   void *p = RawInternalAlloc(size, cache, alignment);
-  if (UNLIKELY(!p))
+  if (UNLIKELY(!p)) {
+    if (AllocatorMayReturnNull()) {
+      return nullptr;
+    }
     ReportInternalAllocatorOutOfMemory(size);
+  }
   return p;
 }
 

compiler-rt/lib/sanitizer_common/sanitizer_win.cpp

@@ -164,7 +164,24 @@ void UnmapOrDie(void *addr, uptr size, bool raw_report) {
 static void *ReturnNullptrOnOOMOrDie(uptr size, const char *mem_type,
                                      const char *mmap_type) {
   error_t last_error = GetLastError();
-  if (last_error == ERROR_NOT_ENOUGH_MEMORY)
+
+  // Assumption: VirtualAlloc is the last system call that was invoked before
+  //   this method.
+  // VirtualAlloc emits one of 2 error codes when running out of memory
+  // 1. ERROR_NOT_ENOUGH_MEMORY:
+  //  There's not enough memory to execute the command
+  // 2. ERROR_INVALID_PARAMETER:
+  //  VirtualAlloc will return this if the request would allocate memory at an
+  //  address exceeding or being very close to the maximum application address
+  //  (the `lpMaximumApplicationAddress` field within the `SystemInfo` struct).
+  //  This does not seem to be officially documented, but is corroborated here:
+  //  https://stackoverflow.com/questions/45833674/why-does-virtualalloc-fail-for-lpaddress-greater-than-0x6ffffffffff
+
+  // Note - It's possible that 'ERROR_COMMITMENT_LIMIT' needs to be handled here
+  // as well. It is currently not handled due to the lack of a reproducer that
+  // induces the error code.
+  if (last_error == ERROR_NOT_ENOUGH_MEMORY ||
+      last_error == ERROR_INVALID_PARAMETER)
     return nullptr;
   ReportMmapFailureAndDie(size, mem_type, mmap_type, last_error);
 }

compiler-rt/test/asan/TestCases/Windows/allocator_may_return_null_set_via_user_function.cpp

@@ -0,0 +1,27 @@
+// RUN: %clangxx_asan -O0 %s -o %t
+// RUN: %run %t 2>&1
+// CHECK: Success
+
+#include <cstdint>
+#include <cstdio>
+#include <cstdlib>
+#include <limits>
+
+// On Windows, flags configured through the user-defined function `__asan_default_options`
+// are suspected to not always be honored according to this GH issue:
+// https://github.com/llvm/llvm-project/issues/117925
+// This issue is resolved for the `allocator_may_return_null` flag, but not for all others.
+// This test ensures we do not regress on `allocator_may_return_null` specifically.
+extern "C" __declspec(dllexport) extern const char *__asan_default_options() {
+  return "allocator_may_return_null=1";
+}
+
+int main() {
+  // Attempt to allocate an excessive amount of memory, which should
+  // terminate the program unless `allocator_may_return_null` is set.
+  size_t max = std::numeric_limits<size_t>::max();
+
+  free(malloc(max));
+  printf("Success");
+  return 0;
+}

compiler-rt/test/asan/TestCases/allocator_may_return_null_limits.cpp

@@ -0,0 +1,23 @@
+// RUN: %clangxx_asan -O0 %s -o %t
+// RUN: %env_asan_opts=allocator_may_return_null=0 not %run %t 2>&1 | FileCheck %s --check-prefix=CHECK1
+// RUN: %env_asan_opts=allocator_may_return_null=1 %run %t 2>&1 | FileCheck %s --check-prefix=CHECK2
+
+// CHECK1: exceeds maximum supported size
+// CHECK1: ABORT
+
+// CHECK2: Success
+
+#include <cstdint>
+#include <cstdio>
+#include <cstdlib>
+#include <limits>
+
+int main() {
+  // Attempt to allocate an excessive amount of memory, which should
+  // terminate the program unless `allocator_may_return_null` is set.
+  size_t max = std::numeric_limits<size_t>::max();
+
+  free(malloc(max));
+  printf("Success");
+  return 0;
+}