[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

llvm-svn: 264327

compiler-rt/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)) {

compiler-rt/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();

compiler-rt/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

compiler-rt/lib/interception/interception_win.cc

@@ -197,11 +197,11 @@ static void **InterestingDLLsAvailable() {
       "kernel32.dll",
       "msvcr110.dll",      // VS2012
       "msvcr120.dll",      // VS2013
+      "vcruntime140.dll",  // VS2015
       "ucrtbase.dll",      // Universal CRT
-    // NTDLL should go last as it exports some functions that we should override
-    // in the CRT [presumably only used internally].
-    "ntdll.dll", NULL
-  };
+      // NTDLL should go last as it exports some functions that we should
+      // override in the CRT [presumably only used internally].
+      "ntdll.dll", NULL};
   static void *result[ARRAY_SIZE(InterestingDLLs)] = { 0 };
   if (!result[0]) {
     for (size_t i = 0, j = 0; InterestingDLLs[i]; ++i) {
@@ -278,6 +278,71 @@ bool OverrideFunction(const char *name, uptr new_func, uptr *orig_old_func) {
   return OverrideFunction(orig_func, new_func, orig_old_func);
 }
 
+bool OverrideImportedFunction(const char *module_to_patch,
+                              const char *imported_module,
+                              const char *function_name, uptr new_function,
+                              uptr *orig_old_func) {
+  HMODULE module = GetModuleHandleA(module_to_patch);
+  if (!module)
+    return false;
+
+  // Check that the module header is full and present.
+  RVAPtr<IMAGE_DOS_HEADER> dos_stub(module, 0);
+  RVAPtr<IMAGE_NT_HEADERS> headers(module, dos_stub->e_lfanew);
+  if (!module || dos_stub->e_magic != IMAGE_DOS_SIGNATURE || // "MZ"
+      headers->Signature != IMAGE_NT_SIGNATURE ||           // "PE\0\0"
+      headers->FileHeader.SizeOfOptionalHeader <
+          sizeof(IMAGE_OPTIONAL_HEADER)) {
+    return false;
+  }
+
+  IMAGE_DATA_DIRECTORY *import_directory =
+      &headers->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT];
+
+  // Iterate the list of imported DLLs. FirstThunk will be null for the last
+  // entry.
+  RVAPtr<IMAGE_IMPORT_DESCRIPTOR> imports(module,
+                                          import_directory->VirtualAddress);
+  for (; imports->FirstThunk != 0; ++imports) {
+    RVAPtr<const char> modname(module, imports->Name);
+    if (_stricmp(&*modname, imported_module) == 0)
+      break;
+  }
+  if (imports->FirstThunk == 0)
+    return false;
+
+  // We have two parallel arrays: the import address table (IAT) and the table
+  // of names. They start out containing the same data, but the loader rewrites
+  // the IAT to hold imported addresses and leaves the name table in
+  // OriginalFirstThunk alone.
+  RVAPtr<IMAGE_THUNK_DATA> name_table(module, imports->OriginalFirstThunk);
+  RVAPtr<IMAGE_THUNK_DATA> iat(module, imports->FirstThunk);
+  for (; name_table->u1.Ordinal != 0; ++name_table, ++iat) {
+    if (!IMAGE_SNAP_BY_ORDINAL(name_table->u1.Ordinal)) {
+      RVAPtr<IMAGE_IMPORT_BY_NAME> import_by_name(
+          module, name_table->u1.ForwarderString);
+      const char *funcname = &import_by_name->Name[0];
+      if (strcmp(funcname, function_name) == 0)
+        break;
+    }
+  }
+  if (name_table->u1.Ordinal == 0)
+    return false;
+
+  // Now we have the correct IAT entry. Do the swap. We have to make the page
+  // read/write first.
+  if (orig_old_func)
+    *orig_old_func = iat->u1.AddressOfData;
+  DWORD old_prot, unused_prot;
+  if (!VirtualProtect(&iat->u1.AddressOfData, 4, PAGE_EXECUTE_READWRITE,
+                      &old_prot))
+    return false;
+  iat->u1.AddressOfData = new_function;
+  if (!VirtualProtect(&iat->u1.AddressOfData, 4, old_prot, &unused_prot))
+    return false;  // Not clear if this failure bothers us.
+  return true;
+}
+
 }  // namespace __interception
 
 #endif  // _WIN32

compiler-rt/lib/interception/interception_win.h

@@ -34,6 +34,14 @@ bool OverrideFunction(const char *name, uptr new_func, uptr *orig_old_func = 0);
 // Windows-only replacement for GetProcAddress. Useful for some sanitizers.
 uptr InternalGetProcAddress(void *module, const char *func_name);
 
+// Overrides a function only when it is called from a specific DLL. For example,
+// this is used to override calls to HeapAlloc/HeapFree from ucrtbase without
+// affecting other third party libraries.
+bool OverrideImportedFunction(const char *module_to_patch,
+                              const char *imported_module,
+                              const char *function_name, uptr new_function,
+                              uptr *orig_old_func);
+
 }  // namespace __interception
 
 #if defined(INTERCEPTION_DYNAMIC_CRT)
@@ -50,5 +58,10 @@ uptr InternalGetProcAddress(void *module, const char *func_name);
 
 #define INTERCEPT_FUNCTION_VER_WIN(func, symver) INTERCEPT_FUNCTION_WIN(func)
 
+#define INTERCEPT_FUNCTION_DLLIMPORT(user_dll, provider_dll, func)       \
+  ::__interception::OverrideImportedFunction(                            \
+      user_dll, provider_dll, #func, (::__interception::uptr)WRAP(func), \
+      (::__interception::uptr *)&REAL(func))
+
 #endif  // INTERCEPTION_WIN_H
 #endif  // _WIN32