[compiler-rt] Intercept weak functions in dll_thunks.

[mpividori]

Summary

In this diff, I update current implementation of the interception to consider the special case of weak functions.
First we check if the client has redefined the function in the main executable (for example: __sanitizer_cov_trace_pc_guard). It we can't find it, then we look for the default implementation (__sanitizer_cov_trace_pc_guard__dll). The default implementation is always available because the static runtime is linked to the main executable.

[zturner]

The code looks correct to me, rnk@ does the approach look sound to you?

compiler-rt/trunk/lib/asan/asan_win_dll_thunk.cc

 #include "asan_init_version.h"
 #include "interception/interception.h"
 #include "sanitizer_common/sanitizer_win_defs.h"
 #include "sanitizer_common/sanitizer_win_dll_thunk.h"
 #include "sanitizer_common/sanitizer_platform_interceptors.h"
 
 // ASan own interface functions.
 #define INTERFACE_FUNCTION(Name) INTERCEPT_SANITIZER_FUNCTION(Name)
-#define INTERFACE_WEAK_FUNCTION(Name)
+#define INTERFACE_WEAK_FUNCTION(Name) INTERCEPT_SANITIZER_WEAK_FUNCTION(Name)
 #include "asan_interface.inc"
 
 // Memory allocation functions.
 INTERCEPT_WRAP_V_W(free)
 INTERCEPT_WRAP_V_W(_free_base)
 INTERCEPT_WRAP_V_WW(_free_dbg)
 
 INTERCEPT_WRAP_W_W(malloc)

compiler-rt/trunk/lib/sanitizer_common/sanitizer_coverage_win_dll_thunk.cc

 // the DLLs that have instrumentation in order to delegate the calls to the
 // shared runtime that lives in the main binary.
 // See https://github.com/google/sanitizers/issues/209 for the details.
 //===----------------------------------------------------------------------===//
 #ifdef SANITIZER_DLL_THUNK
 #include "sanitizer_win_dll_thunk.h"
 // Sanitizer Coverage interface functions.
 #define INTERFACE_FUNCTION(Name) INTERCEPT_SANITIZER_FUNCTION(Name)
-#define INTERFACE_WEAK_FUNCTION(Name)
+#define INTERFACE_WEAK_FUNCTION(Name) INTERCEPT_SANITIZER_WEAK_FUNCTION(Name)
 #include "sanitizer_coverage_interface.inc"
 #endif // SANITIZER_DLL_THUNK

compiler-rt/trunk/lib/sanitizer_common/sanitizer_win_dll_thunk.h

 #ifndef SANITIZER_WIN_DLL_THUNK_H
 #define SANITIZER_WIN_DLL_THUNK_H
 #include "sanitizer_internal_defs.h"
 
 namespace __sanitizer {
 uptr dllThunkGetRealAddrOrDie(const char *name);
 
 int dllThunkIntercept(const char* main_function, uptr dll_function);
+
+int dllThunkInterceptWhenPossible(const char* main_function,
+    const char* default_function, uptr dll_function);
 }
 
 extern "C" int __dll_thunk_init();
 
 // ----------------- Function interception helper macros -------------------- //
 // Override dll_function with main_function from main executable.
 #define INTERCEPT_OR_DIE(main_function, dll_function)                          \
   static int intercept_##dll_function() {                                      \
     return __sanitizer::dllThunkIntercept(main_function, (__sanitizer::uptr)   \
         dll_function);                                                         \
   }                                                                            \
   __pragma(section(".DLLTH$M", long, read))                                    \
   __declspec(allocate(".DLLTH$M")) int (*__dll_thunk_##dll_function)() =       \
     intercept_##dll_function;
 
+// Try to override dll_function with main_function from main executable.
+// If main_function is not present, override dll_function with default_function.
+#define INTERCEPT_WHEN_POSSIBLE(main_function, default_function, dll_function) \
+  static int intercept_##dll_function() {                                      \
+    return __sanitizer::dllThunkInterceptWhenPossible(main_function,           \
+        default_function, (__sanitizer::uptr)dll_function);                    \
+  }                                                                            \
+  __pragma(section(".DLLTH$M", long, read))                                    \
+  __declspec(allocate(".DLLTH$M")) int (*__dll_thunk_##dll_function)() =       \
+    intercept_##dll_function;
+
 // -------------------- Function interception macros ------------------------ //
 // Special case of hooks -- ASan own interface functions.  Those are only called
 // after __asan_init, thus an empty implementation is sufficient.
 #define INTERCEPT_SANITIZER_FUNCTION(name)                                     \
   extern "C" __declspec(noinline) void name() {                                \
     volatile int prevent_icf = (__LINE__ << 8); (void)prevent_icf;             \
     __debugbreak();                                                            \
   }                                                                            \
   INTERCEPT_OR_DIE(#name, name)
 
+// Special case of hooks -- Weak functions, could be redefined in the main
+// executable, but that is not necessary, so we shouldn't die if we can not find
+// a reference. Instead, when the function is not present in the main executable
+// we consider the default impl provided by asan library.
+#define INTERCEPT_SANITIZER_WEAK_FUNCTION(name)                                \
+  extern "C" __declspec(noinline) void name() {                                \
+    volatile int prevent_icf = (__LINE__ << 8); (void)prevent_icf;             \
+    __debugbreak();                                                            \
+  }                                                                            \
+  INTERCEPT_WHEN_POSSIBLE(#name, STRINGIFY(WEAK_EXPORT_NAME(name)), name)
+
 // We can't define our own version of strlen etc. because that would lead to
 // link-time or even type mismatch errors.  Instead, we can declare a function
 // just to be able to get its address.  Me may miss the first few calls to the
 // functions since it can be called before __dll_thunk_init, but that would lead
 // to false negatives in the startup code before user's global initializers,
 // which isn't a big deal.
 #define INTERCEPT_LIBRARY_FUNCTION(name)                                       \
   extern "C" void name();                                                      \

compiler-rt/trunk/lib/sanitizer_common/sanitizer_win_dll_thunk.cc

 }
 
 int dllThunkIntercept(const char* main_function, uptr dll_function) {
   uptr wrapper = dllThunkGetRealAddrOrDie(main_function);
   if (!__interception::OverrideFunction(dll_function, wrapper, 0))
     abort();
   return 0;
 }
+
+int dllThunkInterceptWhenPossible(const char* main_function,
+    const char* default_function, uptr dll_function) {
+  uptr wrapper = __interception::InternalGetProcAddress(
+    (void *)GetModuleHandleA(0), main_function);
+  if (!wrapper)
+    wrapper = dllThunkGetRealAddrOrDie(default_function);
+  if (!__interception::OverrideFunction(dll_function, wrapper, 0))
+    abort();
+  return 0;
+}
 } // namespace __sanitizer
 
 // Include Sanitizer Common interface.
 #define INTERFACE_FUNCTION(Name) INTERCEPT_SANITIZER_FUNCTION(Name)
-#define INTERFACE_WEAK_FUNCTION(Name)
+#define INTERFACE_WEAK_FUNCTION(Name) INTERCEPT_SANITIZER_WEAK_FUNCTION(Name)
 #include "sanitizer_common_interface.inc"
 
 #pragma section(".DLLTH$A", read)  // NOLINT
 #pragma section(".DLLTH$Z", read)  // NOLINT
 
 typedef void (*DllThunkCB)();
 extern "C" {
 __declspec(allocate(".DLLTH$A")) DllThunkCB __start_dll_thunk;

compiler-rt/trunk/lib/ubsan/ubsan_win_dll_thunk.cc

 // the DLLs that have instrumentation in order to delegate the calls to the
 // shared runtime that lives in the main binary.
 // See https://github.com/google/sanitizers/issues/209 for the details.
 //===----------------------------------------------------------------------===//
 #ifdef SANITIZER_DLL_THUNK
 #include "sanitizer_common/sanitizer_win_dll_thunk.h"
 // Ubsan interface functions.
 #define INTERFACE_FUNCTION(Name) INTERCEPT_SANITIZER_FUNCTION(Name)
-#define INTERFACE_WEAK_FUNCTION(Name)
+#define INTERFACE_WEAK_FUNCTION(Name) INTERCEPT_SANITIZER_WEAK_FUNCTION(Name)
 #include "ubsan_interface.inc"
 #endif // SANITIZER_DLL_THUNK