[ASan/Win] Work around PR22545 - unregister globals when using the MD runtime

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

3 files changed, 141 insertions, 9 deletions

diff --git a/lib/asan/asan_globals.cc b/lib/asan/asan_globals.cc
index e5d4417aa..06140bbb3 100644
--- a/lib/asan/asan_globals.cc
+++ b/lib/asan/asan_globals.cc
@@ -210,6 +210,20 @@ void StopInitOrderChecking() {
   }
 }
 
+#if SANITIZER_WINDOWS  // Should only be called on Windows.
+SANITIZER_INTERFACE_ATTRIBUTE
+void UnregisterGlobalsInRange(void *beg, void *end) {
+  if (!flags()->report_globals)
+    return;
+  BlockingMutexLock lock(&mu_for_globals);
+  for (ListOfGlobals *l = list_of_all_globals; l; l = l->next) {
+    void *address = (void *)l->g->beg;
+    if (beg <= address && address < end)
+      UnregisterGlobal(l->g);
+  }
+}
+#endif
+
 }  // namespace __asan
 
 // ---------------------- Interface ---------------- {{{1
diff --git a/lib/asan/asan_win_dynamic_runtime_thunk.cc b/lib/asan/asan_win_dynamic_runtime_thunk.cc
index 3a4de7dbf..19456141c 100644
--- a/lib/asan/asan_win_dynamic_runtime_thunk.cc
+++ b/lib/asan/asan_win_dynamic_runtime_thunk.cc
@@ -23,10 +23,11 @@
 // Using #ifdef rather than relying on Makefiles etc.
 // simplifies the build procedure.
 #ifdef ASAN_DYNAMIC_RUNTIME_THUNK
-extern "C" {
-__declspec(dllimport) int __asan_set_seh_filter();
-__declspec(dllimport) int __asan_should_detect_stack_use_after_return();
+#include <windows.h>
+#include <psapi.h>
 
+extern "C" {
+////////////////////////////////////////////////////////////////////////////////
 // Define a copy of __asan_option_detect_stack_use_after_return that should be
 // used when linking an MD runtime with a set of object files on Windows.
 //
@@ -37,16 +38,82 @@ __declspec(dllimport) int __asan_should_detect_stack_use_after_return();
 // with a MT or MD runtime and we don't want to use ugly __imp_ names on Windows
 // just to work around this issue, let's clone the a variable that is
 // constant after initialization anyways.
+__declspec(dllimport) int __asan_should_detect_stack_use_after_return();
 int __asan_option_detect_stack_use_after_return =
     __asan_should_detect_stack_use_after_return();
+}
+
+////////////////////////////////////////////////////////////////////////////////
+// For some reason, the MD CRT doesn't call the C/C++ terminators as MT does.
+// To work around this, for each DLL we schedule a call to
+// UnregisterGlobalsInRange atexit() specifying the address range of the DLL
+// image to unregister globals in that range.   We don't do the same
+// for the main module (.exe) as the asan_globals.cc allocator is destroyed
+// by the time UnregisterGlobalsInRange is executed.
+// See PR22545 for the details.
+namespace __asan {
+__declspec(dllimport)
+void UnregisterGlobalsInRange(void *beg, void *end);
+}
+
+namespace {
+void *this_module_base, *this_module_end;
+
+void UnregisterGlobals() {
+  __asan::UnregisterGlobalsInRange(this_module_base, this_module_end);
+}
+
+int ScheduleUnregisterGlobals() {
+  HMODULE this_module = 0;
+  // Increments the reference counter of the DLL module, so need to call
+  // FreeLibrary later.
+  if (!GetModuleHandleEx(GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS,
+                         (LPCTSTR)&UnregisterGlobals, &this_module))
+    return 1;
+
+  // Skip the main module.
+  if (this_module == GetModuleHandle(0))
+    return 0;
+
+  MODULEINFO mi;
+  bool success =
+      GetModuleInformation(GetCurrentProcess(), this_module, &mi, sizeof(mi));
+  if (!FreeLibrary(this_module))
+    return 2;
+  if (!success)
+    return 3;
 
-// Set the ASan-specific SEH handler at the end of CRT initialization of each
-// module (see asan_win.cc for the details).
+  this_module_base = mi.lpBaseOfDll;
+  this_module_end = (char*)mi.lpBaseOfDll + mi.SizeOfImage;
+
+  return atexit(UnregisterGlobals);
+}
+}  // namespace
+
+///////////////////////////////////////////////////////////////////////////////
+// ASan SEH handling.
+extern "C" __declspec(dllimport) int __asan_set_seh_filter();
+static int SetSEHFilter() { return __asan_set_seh_filter(); }
+
+///////////////////////////////////////////////////////////////////////////////
+// We schedule some work at start-up by placing callbacks to our code to the
+// list of CRT C initializers.
+//
+// First, declare sections we'll be using:
+#pragma section(".CRT$XID", long, read)  // NOLINT
+#pragma section(".CRT$XIZ", long, read)  // NOLINT
+
+// We need to call 'atexit(UnregisterGlobals);' after atexit() is initialized
+// (.CRT$XIC) but before the C++ constructors (.CRT$XCA).
+__declspec(allocate(".CRT$XID"))
+static int (*__asan_schedule_unregister_globals)() = ScheduleUnregisterGlobals;
+
+// We need to set the ASan-specific SEH handler at the end of CRT initialization
+// of each module (see also asan_win.cc).
 //
 // Unfortunately, putting a pointer to __asan_set_seh_filter into
 // __asan_intercept_seh gets optimized out, so we have to use an extra function.
-static int SetSEHFilter() { return __asan_set_seh_filter(); }
-#pragma section(".CRT$XIZ", long, read)  // NOLINT
-__declspec(allocate(".CRT$XIZ")) int (*__asan_seh_interceptor)() = SetSEHFilter;
-}
+extern "C" __declspec(allocate(".CRT$XIZ"))
+int (*__asan_seh_interceptor)() = SetSEHFilter;
+
 #endif // ASAN_DYNAMIC_RUNTIME_THUNK
diff --git a/test/asan/TestCases/Windows/globals_multiple_dlls.cc b/test/asan/TestCases/Windows/globals_multiple_dlls.cc
new file mode 100644
index 000000000..634e57827
--- /dev/null
+++ b/test/asan/TestCases/Windows/globals_multiple_dlls.cc
@@ -0,0 +1,51 @@
+// Make sure everything works even if the main module doesn't have any stack
+// variables, thus doesn't explicitly reference any symbol exported by the
+// runtime thunk.
+//
+// RUN: %clang_cl_asan -LD -O0 -DDLL %s -Fe%t.dll
+// RUN: %clang_cl_asan -O0 -DEXE %s -Fe%te.exe
+// RUN: env ASAN_OPTIONS=report_globals=2 %run %te.exe %t.dll 2>&1 | FileCheck %s
+
+#include <windows.h>
+#include <stdio.h>
+#include <string.h>
+
+extern "C" {
+#if defined(EXE)
+int main(int argc, char **argv) {
+  if (argc != 2) {
+    printf("Usage: %s [client].dll\n", argv[0]);
+    return 101;
+  }
+  const char *dll_name = argv[1];
+
+// CHECK: time to load DLL
+  printf("time to load DLL\n");
+  fflush(0);
+
+// On DLL load, the "in DLL\n" string is registered:
+// CHECK: Added Global{{.*}} size=19
+// CHECK: in DLL(reason=1)
+  HMODULE dll = LoadLibrary(dll_name);
+  if (dll == NULL)
+    return 3;
+
+// CHECK: in DLL(reason=0)
+// CHECK-NEXT: Removed Global{{.*}} size=19
+  if (!FreeLibrary(dll))
+    return 4;
+
+// CHECK: bye!
+  printf("bye!\n");
+  fflush(0);
+}
+#elif defined(DLL)
+BOOL WINAPI DllMain(HMODULE, DWORD reason, LPVOID) {
+  printf("in DLL(reason=%d)\n", (int)reason);
+  fflush(0);
+  return TRUE;
+}
+#else
+# error oops!
+#endif
+}