[compiler-rt] [asan][win] Fix CreateThread leak (PR #126738)

[via llvm-commits]

https://github.com/GkvJwa updated https://github.com/llvm/llvm-project/pull/126738

>From 2fcd13ef61c69e2b90590cfe9a39569c614bd80d Mon Sep 17 00:00:00 2001
From: GkvJwa <gkvjwa at gmail.com>
Date: Sun, 16 Feb 2025 23:54:01 +0800
Subject: [PATCH] [asan][win] Fix ExitThread leak

Use tls to store the memory created by `VirtualAlloc`, Then
intercept `ExitThread` and release the memory
---
 compiler-rt/lib/asan/asan_win.cpp | 36 ++++++++++++++++++++++++++++++-
 1 file changed, 35 insertions(+), 1 deletion(-)

diff --git a/compiler-rt/lib/asan/asan_win.cpp b/compiler-rt/lib/asan/asan_win.cpp
index 09a13b11cff1f..5ab6351c25219 100644
--- a/compiler-rt/lib/asan/asan_win.cpp
+++ b/compiler-rt/lib/asan/asan_win.cpp
@@ -136,7 +136,32 @@ struct ThreadStartParams {
   void *arg;
 };
 
+static atomic_uint32_t g_native_tls_key{TLS_OUT_OF_INDEXES};
+
+bool AllocTLS(DWORD *key) {
+  DWORD value = TlsAlloc();
+  if (value != TLS_OUT_OF_INDEXES) {
+    *key = value;
+    return true;
+  }
+  return false;
+}
+
 static thread_return_t THREAD_CALLING_CONV asan_thread_start(void *arg) {
+  DWORD key = atomic_load(&g_native_tls_key, memory_order_relaxed);
+  if (key == TLS_OUT_OF_INDEXES) {
+    // Leak global key
+    CHECK(AllocTLS(&key));
+    DWORD old_key = TLS_OUT_OF_INDEXES;
+    // Prevent another thread already store
+    if (!atomic_compare_exchange_strong(&g_native_tls_key, (u32 *)&old_key,
+                                        (u32)key, memory_order_relaxed)) {
+      TlsFree(key);
+      key = atomic_load(&g_native_tls_key, memory_order_relaxed);
+    }
+  }
+  CHECK(key != TLS_OUT_OF_INDEXES);
+  TlsSetValue(key, arg);
   AsanThread *t = (AsanThread *)arg;
   SetCurrentThread(t);
   t->ThreadStart(GetTid());
@@ -145,7 +170,6 @@ static thread_return_t THREAD_CALLING_CONV asan_thread_start(void *arg) {
   t->GetStartData(params);
 
   auto res = (*params.start_routine)(params.arg);
-  t->Destroy();  // POSIX calls this from TSD destructor.
   return res;
 }
 
@@ -166,6 +190,15 @@ INTERCEPTOR_WINAPI(HANDLE, CreateThread, LPSECURITY_ATTRIBUTES security,
                             thr_flags, tid);
 }
 
+INTERCEPTOR_WINAPI(void, ExitThread, DWORD dwExitCode) {
+  DWORD key = atomic_load(&g_native_tls_key, memory_order_relaxed);
+  AsanThread *t = (AsanThread *)TlsGetValue(key);
+  if (t) {
+    t->Destroy();
+  }
+  REAL(ExitThread)(dwExitCode);
+}
+
 // }}}
 
 namespace __asan {
@@ -181,6 +214,7 @@ void InitializePlatformInterceptors() {
       (LPCWSTR)&InitializePlatformInterceptors, &pinned));
 
   ASAN_INTERCEPT_FUNC(CreateThread);
+  ASAN_INTERCEPT_FUNC(ExitThread);
   ASAN_INTERCEPT_FUNC(SetUnhandledExceptionFilter);
 
 #ifdef _WIN64