[ASan/Win] Intercept memory allocation functions in the MD CRT

llvm-svn: 216382

4 files changed, 87 insertions, 22 deletions

diff --git a/compiler-rt/lib/asan/CMakeLists.txt b/compiler-rt/lib/asan/CMakeLists.txt
index 2de4f287d05..db8c9540812 100644
--- a/compiler-rt/lib/asan/CMakeLists.txt
+++ b/compiler-rt/lib/asan/CMakeLists.txt
@@ -54,6 +54,10 @@ endif()
 set(ASAN_DYNAMIC_DEFINITIONS
   ${ASAN_COMMON_DEFINITIONS} ASAN_DYNAMIC=1)
 
+if(WIN32)
+  list(APPEND ASAN_DYNAMIC_DEFINITIONS INTERCEPTION_DYNAMIC_CRT)
+endif()
+
 set(ASAN_DYNAMIC_CFLAGS ${ASAN_CFLAGS})
 append_if(COMPILER_RT_HAS_FTLS_MODEL_INITIAL_EXEC
   -ftls-model=initial-exec ASAN_DYNAMIC_CFLAGS)
diff --git a/compiler-rt/lib/asan/asan_malloc_win.cc b/compiler-rt/lib/asan/asan_malloc_win.cc
index 38dbaf5a4f5..8c1c448f229 100644
--- a/compiler-rt/lib/asan/asan_malloc_win.cc
+++ b/compiler-rt/lib/asan/asan_malloc_win.cc
@@ -23,18 +23,18 @@
 
 #include <stddef.h>
 
-// ---------------------- Replacement functions ---------------- {{{1
 using namespace __asan;  // NOLINT
 
-// FIXME: Simply defining functions with the same signature in *.obj
-// files overrides the standard functions in *.lib
-// This works well for simple helloworld-like tests but might need to be
-// revisited in the future.
+// MT: Simply defining functions with the same signature in *.obj
+// files overrides the standard functions in the CRT.
+// MD: Memory allocation functions are defined in the CRT .dll,
+// so we have to intercept them before they are called for the first time.
 
-// The function attributes will be different for -MD CRT.
-// Just introduce an extra macro for now, it will get a different value under
-// ASAN_DYNAMIC soon.
-#define ALLOCATION_FUNCTION_ATTRIBUTE SANITIZER_INTERFACE_ATTRIBUTE
+#if ASAN_DYNAMIC
+# define ALLOCATION_FUNCTION_ATTRIBUTE
+#else
+# define ALLOCATION_FUNCTION_ATTRIBUTE SANITIZER_INTERFACE_ATTRIBUTE
+#endif
 
 extern "C" {
 ALLOCATION_FUNCTION_ATTRIBUTE
@@ -141,8 +141,36 @@ int _CrtSetReportMode(int, int) {
 
 namespace __asan {
 void ReplaceSystemMalloc() {
-#if defined(_DLL)
-# error MD CRT is not yet supported, see PR20214.
+#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);
+
+  // Override different versions of 'operator new' and 'operator delete'.
+  // No need to override the nothrow versions as they just wrap the throw
+  // versions.
+  // FIXME: Unfortunately, MSVC miscompiles the statements that take the
+  // addresses of the array versions of these operators,
+  // see https://connect.microsoft.com/VisualStudio/feedbackdetail/view/946992
+  // We might want to try to work around this by [inline] assembly or compiling
+  // parts of the RTL with Clang.
+  void *(*op_new)(size_t sz) = operator new;
+  void (*op_delete)(void *p) = operator delete;
+  void *(*op_array_new)(size_t sz) = operator new[];
+  void (*op_array_delete)(void *p) = operator delete[];
+  __interception::OverrideFunction("??2@YAPAXI@Z", (uptr)op_new);
+  __interception::OverrideFunction("??3@YAXPAX@Z", (uptr)op_delete);
+  __interception::OverrideFunction("??_U@YAPAXI@Z", (uptr)op_array_new);
+  __interception::OverrideFunction("??_V@YAXPAX@Z", (uptr)op_array_delete);
 #endif
 }
 }  // namespace __asan
diff --git a/compiler-rt/lib/interception/interception_win.cc b/compiler-rt/lib/interception/interception_win.cc
index 9a5fb67fe25..ecd5d2efa78 100644
--- a/compiler-rt/lib/interception/interception_win.cc
+++ b/compiler-rt/lib/interception/interception_win.cc
@@ -181,6 +181,33 @@ bool OverrideFunction(uptr old_func, uptr new_func, uptr *orig_old_func) {
   return true;
 }
 
+static const void **InterestingDLLsAvailable() {
+  const char *InterestingDLLs[] = { "kernel32.dll", "msvcr120.dll", NULL };
+  static void *result[ARRAY_SIZE(InterestingDLLs)] = { 0 };
+  if (!result[0]) {
+    for (size_t i = 0, j = 0; InterestingDLLs[i]; ++i) {
+      if (HMODULE h = GetModuleHandleA(InterestingDLLs[i]))
+        result[j++] = (void *)h;
+    }
+  }
+  return (const void **)&result[0];
+}
+
+static bool GetFunctionAddressInDLLs(const char *func_name, uptr *func_addr) {
+  *func_addr = 0;
+  const void **DLLs = InterestingDLLsAvailable();
+  for (size_t i = 0; *func_addr == 0 && DLLs[i]; ++i)
+    *func_addr = (uptr)GetProcAddress((HMODULE)DLLs[i], func_name);
+  return (*func_addr != 0);
+}
+
+bool OverrideFunction(const char *name, uptr new_func, uptr *orig_old_func) {
+  uptr orig_func;
+  if (!GetFunctionAddressInDLLs(name, &orig_func))
+    return false;
+  return OverrideFunction(orig_func, new_func, orig_old_func);
+}
+
 }  // namespace __interception
 
 #endif  // _WIN32
diff --git a/compiler-rt/lib/interception/interception_win.h b/compiler-rt/lib/interception/interception_win.h
index 1f866a1e0e2..ba768a7233f 100644
--- a/compiler-rt/lib/interception/interception_win.h
+++ b/compiler-rt/lib/interception/interception_win.h
@@ -22,23 +22,29 @@
 #define INTERCEPTION_WIN_H
 
 namespace __interception {
+// All the functions in the OverrideFunction() family return true on success,
+// false on failure (including "couldn't find the function").
 
-// returns true if the old function existed, false on failure.
-bool OverrideFunction(uptr old_func, uptr new_func, uptr *orig_old_func);
+// Overrides a function by its address.
+bool OverrideFunction(uptr old_func, uptr new_func, uptr *orig_old_func = 0);
+
+// Overrides a function in a system DLL or DLL CRT by its exported name.
+bool OverrideFunction(const char *name, uptr new_func, uptr *orig_old_func = 0);
 }  // namespace __interception
 
-#if defined(_DLL)
-# error Not implemented yet
+#if defined(INTERCEPTION_DYNAMIC_CRT)
+#define INTERCEPT_FUNCTION_WIN(func)                                           \
+  ::__interception::OverrideFunction(#func,                                    \
+                                     (::__interception::uptr)WRAP(func),       \
+                                     (::__interception::uptr *)&REAL(func))
 #else
-# define INTERCEPT_FUNCTION_WIN(func) \
-    ::__interception::OverrideFunction( \
-        (::__interception::uptr)func, \
-        (::__interception::uptr)WRAP(func), \
-        (::__interception::uptr*)&REAL(func))
+#define INTERCEPT_FUNCTION_WIN(func)                                           \
+  ::__interception::OverrideFunction((::__interception::uptr)func,             \
+                                     (::__interception::uptr)WRAP(func),       \
+                                     (::__interception::uptr *)&REAL(func))
 #endif
 
-#define INTERCEPT_FUNCTION_VER_WIN(func, symver) \
-    INTERCEPT_FUNCTION_WIN(func)
+#define INTERCEPT_FUNCTION_VER_WIN(func, symver) INTERCEPT_FUNCTION_WIN(func)
 
 #endif  // INTERCEPTION_WIN_H
 #endif  // _WIN32