[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 d1a37b8b2171131217595e62f334ebb245e15bce Mon Sep 17 00:00:00 2001
From: GkvJwa <gkvjwa at gmail.com>
Date: Tue, 11 Feb 2025 23:52:04 +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 | 29 ++++++++++++++++++++++++++++-
 1 file changed, 28 insertions(+), 1 deletion(-)

diff --git a/compiler-rt/lib/asan/asan_win.cpp b/compiler-rt/lib/asan/asan_win.cpp
index 09a13b11cff1f..1ecf116af644b 100644
--- a/compiler-rt/lib/asan/asan_win.cpp
+++ b/compiler-rt/lib/asan/asan_win.cpp
@@ -136,7 +136,25 @@ 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) {
+    CHECK(AllocTLS(&key));
+    atomic_store(&g_native_tls_key, key, memory_order_release);
+  }
+  CHECK(key != TLS_OUT_OF_INDEXES);
+  TlsSetValue(key, arg);
   AsanThread *t = (AsanThread *)arg;
   SetCurrentThread(t);
   t->ThreadStart(GetTid());
@@ -145,7 +163,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 +183,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();
+  }
+  return REAL(ExitThread)(dwExitCode);
+}
+
 // }}}
 
 namespace __asan {
@@ -181,6 +207,7 @@ void InitializePlatformInterceptors() {
       (LPCWSTR)&InitializePlatformInterceptors, &pinned));
 
   ASAN_INTERCEPT_FUNC(CreateThread);
+  ASAN_INTERCEPT_FUNC(ExitThread);
   ASAN_INTERCEPT_FUNC(SetUnhandledExceptionFilter);
 
 #ifdef _WIN64