[asan] Intercept all Heap* related imports from ucrtbase.dll

ucrtbase.dll appears to be built with some kind of cross-module
inlining, because there are calls to imported Heap* routines sprinkled
throughout the code. This inlining defeats our attempts to hotpatch
malloc, _malloc_base, and related functions. Failing to intercept an
allocation or deallocation results in a crash when the program attempts
to deallocate or reallocate memory with the wrong allocator.

This change patches the IAT of ucrtbase.dll to replace the addresses of
the imported Heap* functions with implementations provided by ASan.  We
don't globally intercept the win32 Heap* functions because they are
typically used by system DLLs that run before ASan initializes.
Eventually, we may want to intercept them, but for now I think this is
the minimal change that will keep ASan stable.

Reviewers: samsonov

Differential Revision: http://reviews.llvm.org/D18413

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

3 files changed, 85 insertions, 15 deletions

diff --git a/lib/asan/asan_allocator.cc b/lib/asan/asan_allocator.cc
index bbca721d2..6493c7947 100644
--- a/lib/asan/asan_allocator.cc
+++ b/lib/asan/asan_allocator.cc
@@ -762,7 +762,7 @@ int asan_posix_memalign(void **memptr, uptr alignment, uptr size,
   return 0;
 }
 
-uptr asan_malloc_usable_size(void *ptr, uptr pc, uptr bp) {
+uptr asan_malloc_usable_size(const void *ptr, uptr pc, uptr bp) {
   if (!ptr) return 0;
   uptr usable_size = instance.AllocationSize(reinterpret_cast<uptr>(ptr));
   if (flags()->check_malloc_usable_size && (usable_size == 0)) {
diff --git a/lib/asan/asan_allocator.h b/lib/asan/asan_allocator.h
index e3d53330c..5ddaeb325 100644
--- a/lib/asan/asan_allocator.h
+++ b/lib/asan/asan_allocator.h
@@ -171,7 +171,7 @@ void *asan_pvalloc(uptr size, BufferedStackTrace *stack);
 
 int asan_posix_memalign(void **memptr, uptr alignment, uptr size,
                         BufferedStackTrace *stack);
-uptr asan_malloc_usable_size(void *ptr, uptr pc, uptr bp);
+uptr asan_malloc_usable_size(const void *ptr, uptr pc, uptr bp);
 
 uptr asan_mz_size(const void *ptr);
 void asan_mz_force_lock();
diff --git a/lib/asan/asan_malloc_win.cc b/lib/asan/asan_malloc_win.cc
index a525f717f..4a233dfe9 100644
--- a/lib/asan/asan_malloc_win.cc
+++ b/lib/asan/asan_malloc_win.cc
@@ -14,6 +14,8 @@
 
 #include "sanitizer_common/sanitizer_platform.h"
 #if SANITIZER_WINDOWS
+#define WIN32_LEAN_AND_MEAN
+#include <windows.h>
 
 #include "asan_allocator.h"
 #include "asan_interceptors.h"
@@ -123,7 +125,7 @@ void *_recalloc(void *p, size_t n, size_t elem_size) {
 }
 
 ALLOCATION_FUNCTION_ATTRIBUTE
-size_t _msize(void *ptr) {
+size_t _msize(const void *ptr) {
   GET_CURRENT_PC_BP_SP;
   (void)sp;
   return asan_malloc_usable_size(ptr, pc, bp);
@@ -159,21 +161,89 @@ int _CrtSetReportMode(int, int) {
 }
 }  // extern "C"
 
+INTERCEPTOR_WINAPI(LPVOID, HeapAlloc, HANDLE hHeap, DWORD dwFlags,
+                   SIZE_T dwBytes) {
+  GET_STACK_TRACE_MALLOC;
+  void *p = asan_malloc(dwBytes, &stack);
+  // Reading MSDN suggests that the *entire* usable allocation is zeroed out.
+  // Otherwise it is difficult to HeapReAlloc with HEAP_ZERO_MEMORY.
+  // https://blogs.msdn.microsoft.com/oldnewthing/20120316-00/?p=8083
+  if (dwFlags == HEAP_ZERO_MEMORY)
+    internal_memset(p, 0, asan_mz_size(p));
+  else
+    CHECK(dwFlags == 0 && "unsupported heap flags");
+  return p;
+}
+
+INTERCEPTOR_WINAPI(BOOL, HeapFree, HANDLE hHeap, DWORD dwFlags, LPVOID lpMem) {
+  CHECK(dwFlags == 0 && "unsupported heap flags");
+  GET_STACK_TRACE_FREE;
+  asan_free(lpMem, &stack, FROM_MALLOC);
+  return true;
+}
+
+INTERCEPTOR_WINAPI(LPVOID, HeapReAlloc, HANDLE hHeap, DWORD dwFlags,
+                   LPVOID lpMem, SIZE_T dwBytes) {
+  GET_STACK_TRACE_MALLOC;
+  // Realloc should never reallocate in place.
+  if (dwFlags & HEAP_REALLOC_IN_PLACE_ONLY)
+    return nullptr;
+  CHECK(dwFlags == 0 && "unsupported heap flags");
+  return asan_realloc(lpMem, dwBytes, &stack);
+}
+
+INTERCEPTOR_WINAPI(SIZE_T, HeapSize, HANDLE hHeap, DWORD dwFlags,
+                   LPCVOID lpMem) {
+  CHECK(dwFlags == 0 && "unsupported heap flags");
+  GET_CURRENT_PC_BP_SP;
+  (void)sp;
+  return asan_malloc_usable_size(lpMem, pc, bp);
+}
+
 namespace __asan {
+
+static void TryToOverrideFunction(const char *fname, uptr new_func) {
+  // Failure here is not fatal. The CRT may not be present, and different CRT
+  // versions use different symbols.
+  if (!__interception::OverrideFunction(fname, new_func))
+    VPrintf(2, "Failed to override function %s\n", fname);
+}
+
 void ReplaceSystemMalloc() {
 #if defined(ASAN_DYNAMIC)
-  // We don't check the result because CRT might not be used in the process.
-  __interception::OverrideFunction("free", (uptr)free);
-  __interception::OverrideFunction("malloc", (uptr)malloc);
-  __interception::OverrideFunction("_malloc_crt", (uptr)malloc);
-  __interception::OverrideFunction("calloc", (uptr)calloc);
-  __interception::OverrideFunction("_calloc_crt", (uptr)calloc);
-  __interception::OverrideFunction("realloc", (uptr)realloc);
-  __interception::OverrideFunction("_realloc_crt", (uptr)realloc);
-  __interception::OverrideFunction("_recalloc", (uptr)_recalloc);
-  __interception::OverrideFunction("_recalloc_crt", (uptr)_recalloc);
-  __interception::OverrideFunction("_msize", (uptr)_msize);
-  __interception::OverrideFunction("_expand", (uptr)_expand);
+  TryToOverrideFunction("free", (uptr)free);
+  TryToOverrideFunction("_free_base", (uptr)free);
+  TryToOverrideFunction("malloc", (uptr)malloc);
+  TryToOverrideFunction("_malloc_base", (uptr)malloc);
+  TryToOverrideFunction("_malloc_crt", (uptr)malloc);
+  TryToOverrideFunction("calloc", (uptr)calloc);
+  TryToOverrideFunction("_calloc_base", (uptr)calloc);
+  TryToOverrideFunction("_calloc_crt", (uptr)calloc);
+  TryToOverrideFunction("realloc", (uptr)realloc);
+  TryToOverrideFunction("_realloc_base", (uptr)realloc);
+  TryToOverrideFunction("_realloc_crt", (uptr)realloc);
+  TryToOverrideFunction("_recalloc", (uptr)_recalloc);
+  TryToOverrideFunction("_recalloc_crt", (uptr)_recalloc);
+  TryToOverrideFunction("_msize", (uptr)_msize);
+  TryToOverrideFunction("_expand", (uptr)_expand);
+  TryToOverrideFunction("_expand_base", (uptr)_expand);
+
+  // Recent versions of ucrtbase.dll appear to be built with PGO and LTCG, which
+  // enable cross-module inlining. This means our _malloc_base hook won't catch
+  // all CRT allocations. This code here patches the import table of
+  // ucrtbase.dll so that all attempts to use the lower-level win32 heap
+  // allocation API will be directed to ASan's heap. We don't currently
+  // intercept all calls to HeapAlloc. If we did, we would have to check on
+  // HeapFree whether the pointer came from ASan of from the system.
+#define INTERCEPT_UCRT_FUNCTION(func)                                         \
+  if (!INTERCEPT_FUNCTION_DLLIMPORT("ucrtbase.dll",                           \
+                                    "api-ms-win-core-heap-l1-1-0.dll", func)) \
+    VPrintf(2, "Failed to intercept ucrtbase.dll import %s\n", #func);
+  INTERCEPT_UCRT_FUNCTION(HeapAlloc);
+  INTERCEPT_UCRT_FUNCTION(HeapFree);
+  INTERCEPT_UCRT_FUNCTION(HeapReAlloc);
+  INTERCEPT_UCRT_FUNCTION(HeapSize);
+#undef INTERCEPT_UCRT_FUNCTION
 #endif
 }
 }  // namespace __asan