[compiler-rt] [ASan][Windows] Synchronizing ASAN init on Windows (PR #71833)

[Zack Johnson via llvm-commits]

https://github.com/zacklj89 updated https://github.com/llvm/llvm-project/pull/71833

>From cd2beca62b91247f24ab930d01c1b23c5a11d899 Mon Sep 17 00:00:00 2001
From: Zachary Johnson <zajohnson at microsoft.com>
Date: Thu, 9 Nov 2023 15:13:24 -0500
Subject: [PATCH 1/2] [asan] Fixing initialization synchronization on Windows

---
 compiler-rt/lib/asan/asan_internal.h | 13 +++++++++++++
 compiler-rt/lib/asan/asan_thread.cpp | 23 +++++++++++++++++++++++
 2 files changed, 36 insertions(+)

diff --git a/compiler-rt/lib/asan/asan_internal.h b/compiler-rt/lib/asan/asan_internal.h
index 2944ebe213b5d5..213fc5c3984437 100644
--- a/compiler-rt/lib/asan/asan_internal.h
+++ b/compiler-rt/lib/asan/asan_internal.h
@@ -132,6 +132,19 @@ void InstallAtForkHandler();
   if (&__asan_on_error) \
   __asan_on_error()
 
+// Depending on the loading thread and when ASAN is loaded on Windows,
+// race conditions can appear causing incorrect states or internal check
+// failures.
+//
+// From a multithreaded managed environment, if an ASAN instrumented dll
+// is loading on a spawned thread, an intercepted function may be called on
+// multiple threads while ASAN is still in the process of initialization. This
+// can also cause the ASAN thread registry to create the "main" thread after
+// another thread, resulting in a TID != 0.
+//
+// Two threads can also race to initialize ASAN, resulting in either incorrect
+// state or internal check failures for init already running.
+//
 bool AsanInited();
 extern bool replace_intrin_cached;
 extern void (*death_callback)(void);
diff --git a/compiler-rt/lib/asan/asan_thread.cpp b/compiler-rt/lib/asan/asan_thread.cpp
index 8798968947e82e..88d526069352bc 100644
--- a/compiler-rt/lib/asan/asan_thread.cpp
+++ b/compiler-rt/lib/asan/asan_thread.cpp
@@ -27,6 +27,10 @@ namespace __asan {
 
 // AsanThreadContext implementation.
 
+#if SANITIZER_WINDOWS
+static atomic_uint8_t main_thread_created{0};
+#endif
+
 void AsanThreadContext::OnCreated(void *arg) {
   CreateThreadContextArgs *args = static_cast<CreateThreadContextArgs *>(arg);
   if (args->stack)
@@ -93,6 +97,11 @@ AsanThreadContext *GetThreadContextByTidLocked(u32 tid) {
 AsanThread *AsanThread::Create(const void *start_data, uptr data_size,
                                u32 parent_tid, StackTrace *stack,
                                bool detached) {
+#if SANITIZER_WINDOWS
+  while (atomic_load(&main_thread_created, memory_order_acquire) == 0) {
+    // If another thread is trying to be created before the main thread, wait.
+  }
+#endif
   uptr PageSize = GetPageSizeCached();
   uptr size = RoundUpTo(sizeof(AsanThread), PageSize);
   AsanThread *thread = (AsanThread *)MmapOrDie(size, __func__);
@@ -288,11 +297,25 @@ void AsanThread::ThreadStart(tid_t os_id) {
 }
 
 AsanThread *CreateMainThread() {
+// Depending on the loading thread, specifically in managed scenarios, the main
+// thread can be created after other threads on Windows. This ensures we start
+// the main thread before those threads.
+#  if SANITIZER_WINDOWS
+  uptr PageSize = GetPageSizeCached();
+  uptr size = RoundUpTo(sizeof(AsanThread), PageSize);
+  AsanThread *main_thread = (AsanThread *)MmapOrDie(size, __func__);
+  AsanThreadContext::CreateThreadContextArgs args = {main_thread, nullptr};
+  asanThreadRegistry().CreateThread(0, true, kMainTid, &args);
+  SetCurrentThread(main_thread);
+  main_thread->ThreadStart(internal_getpid());
+  atomic_store(&main_thread_created, 1, memory_order_release);
+#  else
   AsanThread *main_thread = AsanThread::Create(
       /* parent_tid */ kMainTid,
       /* stack */ nullptr, /* detached */ true);
   SetCurrentThread(main_thread);
   main_thread->ThreadStart(internal_getpid());
+#  endif
   return main_thread;
 }
 

>From 45b80b93fc0022173e77dd67d9a9cb6fac922693 Mon Sep 17 00:00:00 2001
From: Zachary Johnson <zajohnson at microsoft.com>
Date: Thu, 16 Nov 2023 10:14:11 -0500
Subject: [PATCH 2/2] adding yields

---
 compiler-rt/lib/asan/asan_thread.cpp | 1 +
 1 file changed, 1 insertion(+)

diff --git a/compiler-rt/lib/asan/asan_thread.cpp b/compiler-rt/lib/asan/asan_thread.cpp
index 88d526069352bc..dc0ad2caf3bbd1 100644
--- a/compiler-rt/lib/asan/asan_thread.cpp
+++ b/compiler-rt/lib/asan/asan_thread.cpp
@@ -100,6 +100,7 @@ AsanThread *AsanThread::Create(const void *start_data, uptr data_size,
 #if SANITIZER_WINDOWS
   while (atomic_load(&main_thread_created, memory_order_acquire) == 0) {
     // If another thread is trying to be created before the main thread, wait.
+    internal_sched_yield();
   }
 #endif
   uptr PageSize = GetPageSizeCached();