allocator: Add Windows _aligned_* shims

On Windows we don’t currently hook the _aligned_* allocation APIs, this
can cause issues because _aligned_realloc can call HeapSize and cause
GWP-ASan crashes similar to bug 909720. Unfortunately the
_aligned_realloc API is different enough that it can not be implemented
using the standard POSIX shims, in particular because _aligned_malloc
and _aligned_free don't return valid allocation addresses, they are
offsets into allocations.

I add new Windows platform-specific shims for _aligned_malloc,
_aligned_realloc, and _aligned_free and wire them in for all users of
the allocator shims. I implement these routines on top of the Windows
Heap* API and leave uncommon _aligned_* shims to crash to ensure that
any future uses immediately surface why their use fails.

Bug: 912500, 896019
Change-Id: Ieaa50b816ab277a6ad4b80ee8519027343fa9878
Reviewed-on: https://chromium-review.googlesource.com/c/1367485
Reviewed-by: danakj <[email protected]>
Reviewed-by: Vitaly Buka <[email protected]>
Reviewed-by: Alexei Filippov <[email protected]>
Reviewed-by: Erik Chen <[email protected]>
Reviewed-by: Will Harris <[email protected]>
Reviewed-by: Sigurður Ásgeirsson <[email protected]>
Commit-Queue: Vlad Tsyrklevich <[email protected]>
Cr-Commit-Position: refs/heads/master@{#616106}

base/BUILD.gn

@@ -2781,6 +2781,7 @@ test("base_unittests") {
   if (use_allocator_shim) {
     sources += [
       "allocator/allocator_shim_unittest.cc",
+      "allocator/winheap_stubs_win_unittest.cc",
       "sampling_heap_profiler/sampling_heap_profiler_unittest.cc",
     ]
   }

base/allocator/allocator_shim.cc

@@ -277,6 +277,38 @@ ALWAYS_INLINE void ShimFreeDefiniteSize(void* ptr, size_t size, void* context) {
                                                  context);
 }
 
+ALWAYS_INLINE void* ShimAlignedMalloc(size_t size,
+                                      size_t alignment,
+                                      void* context) {
+  const allocator::AllocatorDispatch* const chain_head = GetChainHead();
+  void* ptr = nullptr;
+  do {
+    ptr = chain_head->aligned_malloc_function(chain_head, size, alignment,
+                                              context);
+  } while (!ptr && g_call_new_handler_on_malloc_failure &&
+           CallNewHandler(size));
+  return ptr;
+}
+
+ALWAYS_INLINE void* ShimAlignedRealloc(void* address,
+                                       size_t size,
+                                       size_t alignment,
+                                       void* context) {
+  const allocator::AllocatorDispatch* const chain_head = GetChainHead();
+  void* ptr = nullptr;
+  do {
+    ptr = chain_head->aligned_realloc_function(chain_head, address, size,
+                                               alignment, context);
+  } while (!ptr && g_call_new_handler_on_malloc_failure &&
+           CallNewHandler(size));
+  return ptr;
+}
+
+ALWAYS_INLINE void ShimAlignedFree(void* address, void* context) {
+  const allocator::AllocatorDispatch* const chain_head = GetChainHead();
+  return chain_head->aligned_free_function(chain_head, address, context);
+}
+
 }  // extern "C"
 
 #if !defined(OS_WIN) && !defined(OS_MACOSX)

base/allocator/allocator_shim.h

@@ -41,7 +41,7 @@ namespace allocator {
 // It is possible to dynamically insert further AllocatorDispatch stages
 // to the front of the chain, for debugging / profiling purposes.
 //
-// All the functions must be thred safe. The shim does not enforce any
+// All the functions must be thread safe. The shim does not enforce any
 // serialization. This is to route to thread-aware allocators (e.g, tcmalloc)
 // wihout introducing unnecessary perf hits.
 
@@ -84,16 +84,35 @@ struct AllocatorDispatch {
                                   void* ptr,
                                   size_t size,
                                   void* context);
+  using AlignedMallocFn = void*(const AllocatorDispatch* self,
+                                size_t size,
+                                size_t alignment,
+                                void* context);
+  using AlignedReallocFn = void*(const AllocatorDispatch* self,
+                                 void* address,
+                                 size_t size,
+                                 size_t alignment,
+                                 void* context);
+  using AlignedFreeFn = void(const AllocatorDispatch* self,
+                             void* address,
+                             void* context);
 
   AllocFn* const alloc_function;
   AllocZeroInitializedFn* const alloc_zero_initialized_function;
   AllocAlignedFn* const alloc_aligned_function;
   ReallocFn* const realloc_function;
   FreeFn* const free_function;
   GetSizeEstimateFn* const get_size_estimate_function;
+  // batch_malloc, batch_free, and free_definite_size are specific to the OSX
+  // and iOS allocators.
   BatchMallocFn* const batch_malloc_function;
   BatchFreeFn* const batch_free_function;
   FreeDefiniteSizeFn* const free_definite_size_function;
+  // _aligned_malloc, _aligned_realloc, and _aligned_free are specific to the
+  // Windows allocator.
+  AlignedMallocFn* const aligned_malloc_function;
+  AlignedReallocFn* const aligned_realloc_function;
+  AlignedFreeFn* const aligned_free_function;
 
   const AllocatorDispatch* next;
 

base/allocator/allocator_shim_default_dispatch_to_glibc.cc

@@ -71,5 +71,8 @@ const AllocatorDispatch AllocatorDispatch::default_dispatch = {
     nullptr,               /* batch_malloc_function */
     nullptr,               /* batch_free_function */
     nullptr,               /* free_definite_size_function */
+    nullptr,               /* aligned_malloc_function */
+    nullptr,               /* aligned_realloc_function */
+    nullptr,               /* aligned_free_function */
     nullptr,               /* next */
 };

base/allocator/allocator_shim_default_dispatch_to_linker_wrapped_symbols.cc

@@ -109,5 +109,8 @@ const AllocatorDispatch AllocatorDispatch::default_dispatch = {
     nullptr,           /* batch_malloc_function */
     nullptr,           /* batch_free_function */
     nullptr,           /* free_definite_size_function */
+    nullptr,           /* aligned_malloc_function */
+    nullptr,           /* aligned_realloc_function */
+    nullptr,           /* aligned_free_function */
     nullptr,           /* next */
 };

base/allocator/allocator_shim_default_dispatch_to_mac_zoned_malloc.cc

@@ -102,6 +102,9 @@ const AllocatorDispatch AllocatorDispatch::default_dispatch = {
     &BatchMallocImpl,      /* batch_malloc_function */
     &BatchFreeImpl,        /* batch_free_function */
     &FreeDefiniteSizeImpl, /* free_definite_size_function */
+    nullptr,               /* aligned_malloc_function */
+    nullptr,               /* aligned_realloc_function */
+    nullptr,               /* aligned_free_function */
     nullptr,               /* next */
 };
 

base/allocator/allocator_shim_default_dispatch_to_tcmalloc.cc

@@ -62,6 +62,9 @@ const AllocatorDispatch AllocatorDispatch::default_dispatch = {
     nullptr,            /* batch_malloc_function */
     nullptr,            /* batch_free_function */
     nullptr,            /* free_definite_size_function */
+    nullptr,            /* aligned_malloc_function */
+    nullptr,            /* aligned_realloc_function */
+    nullptr,            /* aligned_free_function */
     nullptr,            /* next */
 };
 

base/allocator/allocator_shim_default_dispatch_to_winheap.cc

@@ -60,6 +60,27 @@ size_t DefaultWinHeapGetSizeEstimateImpl(const AllocatorDispatch*,
   return base::allocator::WinHeapGetSizeEstimate(address);
 }
 
+void* DefaultWinHeapAlignedMallocImpl(const AllocatorDispatch*,
+                                      size_t size,
+                                      size_t alignment,
+                                      void* context) {
+  return base::allocator::WinHeapAlignedMalloc(size, alignment);
+}
+
+void* DefaultWinHeapAlignedReallocImpl(const AllocatorDispatch*,
+                                       void* ptr,
+                                       size_t size,
+                                       size_t alignment,
+                                       void* context) {
+  return base::allocator::WinHeapAlignedRealloc(ptr, size, alignment);
+}
+
+void DefaultWinHeapAlignedFreeImpl(const AllocatorDispatch*,
+                                   void* ptr,
+                                   void* context) {
+  base::allocator::WinHeapAlignedFree(ptr);
+}
+
 }  // namespace
 
 // Guarantee that default_dispatch is compile-time initialized to avoid using
@@ -75,5 +96,8 @@ constexpr AllocatorDispatch AllocatorDispatch::default_dispatch = {
     nullptr, /* batch_malloc_function */
     nullptr, /* batch_free_function */
     nullptr, /* free_definite_size_function */
+    &DefaultWinHeapAlignedMallocImpl,
+    &DefaultWinHeapAlignedReallocImpl,
+    &DefaultWinHeapAlignedFreeImpl,
     nullptr, /* next */
 };

base/allocator/allocator_shim_override_ucrt_symbols_win.h

@@ -70,6 +70,60 @@ size_t _msize(void* memblock) {
   return ShimGetSizeEstimate(memblock, nullptr);
 }
 
+__declspec(restrict) void* _aligned_malloc(size_t size, size_t alignment) {
+  return ShimAlignedMalloc(size, alignment, nullptr);
+}
+
+__declspec(restrict) void* _aligned_realloc(void* address,
+                                            size_t size,
+                                            size_t alignment) {
+  return ShimAlignedRealloc(address, size, alignment, nullptr);
+}
+
+void _aligned_free(void* address) {
+  ShimAlignedFree(address, nullptr);
+}
+
+// The following uncommon _aligned_* routines are not used in Chromium and have
+// been shimmed to immediately crash to ensure that implementations are added if
+// uses are introduced.
+__declspec(restrict) void* _aligned_recalloc(void* address,
+                                             size_t num,
+                                             size_t size,
+                                             size_t alignment) {
+  CHECK(false) << "This routine has not been implemented";
+  __builtin_unreachable();
+}
+
+size_t _aligned_msize(void* address, size_t alignment, size_t offset) {
+  CHECK(false) << "This routine has not been implemented";
+  __builtin_unreachable();
+}
+
+__declspec(restrict) void* _aligned_offset_malloc(size_t size,
+                                                  size_t alignment,
+                                                  size_t offset) {
+  CHECK(false) << "This routine has not been implemented";
+  __builtin_unreachable();
+}
+
+__declspec(restrict) void* _aligned_offset_realloc(void* address,
+                                                   size_t size,
+                                                   size_t alignment,
+                                                   size_t offset) {
+  CHECK(false) << "This routine has not been implemented";
+  __builtin_unreachable();
+}
+
+__declspec(restrict) void* _aligned_offset_recalloc(void* address,
+                                                    size_t num,
+                                                    size_t size,
+                                                    size_t alignment,
+                                                    size_t offset) {
+  CHECK(false) << "This routine has not been implemented";
+  __builtin_unreachable();
+}
+
 // The symbols
 //   * __acrt_heap
 //   * __acrt_initialize_heap

base/allocator/allocator_shim_unittest.cc

@@ -179,6 +179,40 @@ class AllocatorShimTest : public testing::Test {
     self->next->free_definite_size_function(self->next, ptr, size, context);
   }
 
+  static void* MockAlignedMalloc(const AllocatorDispatch* self,
+                                 size_t size,
+                                 size_t alignment,
+                                 void* context) {
+    if (instance_ && size < kMaxSizeTracked) {
+      ++instance_->aligned_mallocs_intercepted_by_size[size];
+    }
+    return self->next->aligned_malloc_function(self->next, size, alignment,
+                                               context);
+  }
+
+  static void* MockAlignedRealloc(const AllocatorDispatch* self,
+                                  void* address,
+                                  size_t size,
+                                  size_t alignment,
+                                  void* context) {
+    if (instance_) {
+      if (size < kMaxSizeTracked)
+        ++instance_->aligned_reallocs_intercepted_by_size[size];
+      ++instance_->aligned_reallocs_intercepted_by_addr[Hash(address)];
+    }
+    return self->next->aligned_realloc_function(self->next, address, size,
+                                                alignment, context);
+  }
+
+  static void MockAlignedFree(const AllocatorDispatch* self,
+                              void* address,
+                              void* context) {
+    if (instance_) {
+      ++instance_->aligned_frees_intercepted_by_addr[Hash(address)];
+    }
+    self->next->aligned_free_function(self->next, address, context);
+  }
+
   static void NewHandler() {
     if (!instance_)
       return;
@@ -196,10 +230,15 @@ class AllocatorShimTest : public testing::Test {
     memset(&aligned_allocs_intercepted_by_size, 0, array_size);
     memset(&aligned_allocs_intercepted_by_alignment, 0, array_size);
     memset(&reallocs_intercepted_by_size, 0, array_size);
+    memset(&reallocs_intercepted_by_addr, 0, array_size);
     memset(&frees_intercepted_by_addr, 0, array_size);
     memset(&batch_mallocs_intercepted_by_size, 0, array_size);
     memset(&batch_frees_intercepted_by_addr, 0, array_size);
     memset(&free_definite_sizes_intercepted_by_size, 0, array_size);
+    memset(&aligned_mallocs_intercepted_by_size, 0, array_size);
+    memset(&aligned_reallocs_intercepted_by_size, 0, array_size);
+    memset(&aligned_reallocs_intercepted_by_addr, 0, array_size);
+    memset(&aligned_frees_intercepted_by_addr, 0, array_size);
     did_fail_realloc_0xfeed_once.reset(new ThreadLocalBoolean());
     subtle::Release_Store(&num_new_handler_calls, 0);
     instance_ = this;
@@ -227,6 +266,10 @@ class AllocatorShimTest : public testing::Test {
   size_t batch_mallocs_intercepted_by_size[kMaxSizeTracked];
   size_t batch_frees_intercepted_by_addr[kMaxSizeTracked];
   size_t free_definite_sizes_intercepted_by_size[kMaxSizeTracked];
+  size_t aligned_mallocs_intercepted_by_size[kMaxSizeTracked];
+  size_t aligned_reallocs_intercepted_by_size[kMaxSizeTracked];
+  size_t aligned_reallocs_intercepted_by_addr[kMaxSizeTracked];
+  size_t aligned_frees_intercepted_by_addr[kMaxSizeTracked];
   std::unique_ptr<ThreadLocalBoolean> did_fail_realloc_0xfeed_once;
   subtle::Atomic32 num_new_handler_calls;
 
@@ -271,6 +314,9 @@ AllocatorDispatch g_mock_dispatch = {
     &AllocatorShimTest::MockBatchMalloc,      /* batch_malloc_function */
     &AllocatorShimTest::MockBatchFree,        /* batch_free_function */
     &AllocatorShimTest::MockFreeDefiniteSize, /* free_definite_size_function */
+    &AllocatorShimTest::MockAlignedMalloc,    /* aligned_malloc_function */
+    &AllocatorShimTest::MockAlignedRealloc,   /* aligned_realloc_function */
+    &AllocatorShimTest::MockAlignedFree,      /* aligned_free_function */
     nullptr,                                  /* next */
 };
 
@@ -399,6 +445,25 @@ TEST_F(AllocatorShimTest, InterceptLibcSymbolsFreeDefiniteSize) {
 }
 #endif  // defined(OS_MACOSX)
 
+#if defined(OS_WIN)
+TEST_F(AllocatorShimTest, InterceptUcrtAlignedAllocationSymbols) {
+  InsertAllocatorDispatch(&g_mock_dispatch);
+
+  constexpr size_t kAlignment = 32;
+  void* alloc_ptr = _aligned_malloc(123, kAlignment);
+  EXPECT_GE(aligned_mallocs_intercepted_by_size[123], 1u);
+
+  void* new_alloc_ptr = _aligned_realloc(alloc_ptr, 1234, kAlignment);
+  EXPECT_GE(aligned_reallocs_intercepted_by_size[1234], 1u);
+  EXPECT_GE(aligned_reallocs_intercepted_by_addr[Hash(alloc_ptr)], 1u);
+
+  _aligned_free(new_alloc_ptr);
+  EXPECT_GE(aligned_frees_intercepted_by_addr[Hash(new_alloc_ptr)], 1u);
+
+  RemoveAllocatorDispatchForTesting(&g_mock_dispatch);
+}
+#endif
+
 TEST_F(AllocatorShimTest, InterceptCppSymbols) {
   InsertAllocatorDispatch(&g_mock_dispatch);