[libc-commits] [libc] [libc] fix tls teardown while being used (PR #108229)

[Schrodinger ZHU Yifan via libc-commits]

https://github.com/SchrodingerZhu created https://github.com/llvm/llvm-project/pull/108229

None

>From 1ba4d36b6542e0b83378d7f404786b087fb2651b Mon Sep 17 00:00:00 2001
From: Schrodinger ZHU Yifan <yifanzhu at rochester.edu>
Date: Wed, 11 Sep 2024 09:56:47 -0400
Subject: [PATCH] [libc] fix tls teardown while being used

---
 libc/src/stdlib/atexit.cpp      |  6 +++++-
 libc/src/stdlib/quick_exit.cpp  |  3 +++
 libc/startup/linux/do_start.cpp | 11 ++++++-----
 3 files changed, 14 insertions(+), 6 deletions(-)

diff --git a/libc/src/stdlib/atexit.cpp b/libc/src/stdlib/atexit.cpp
index c8a15dd3cfef2d..799aad136bda5c 100644
--- a/libc/src/stdlib/atexit.cpp
+++ b/libc/src/stdlib/atexit.cpp
@@ -16,6 +16,7 @@ namespace LIBC_NAMESPACE_DECL {
 
 constinit ExitCallbackList atexit_callbacks;
 Mutex handler_list_mtx(false, false, false, false);
+[[gnu::weak]] extern void teardown_main_tls();
 
 extern "C" {
 
@@ -24,8 +25,11 @@ int __cxa_atexit(AtExitCallback *callback, void *payload, void *) {
 }
 
 void __cxa_finalize(void *dso) {
-  if (!dso)
+  if (!dso) {
     call_exit_callbacks(atexit_callbacks);
+    if (teardown_main_tls)
+      teardown_main_tls();
+  }
 }
 
 } // extern "C"
diff --git a/libc/src/stdlib/quick_exit.cpp b/libc/src/stdlib/quick_exit.cpp
index a5abf3e05d1a13..29110b33afcf50 100644
--- a/libc/src/stdlib/quick_exit.cpp
+++ b/libc/src/stdlib/quick_exit.cpp
@@ -16,9 +16,12 @@
 namespace LIBC_NAMESPACE_DECL {
 
 extern ExitCallbackList at_quick_exit_callbacks;
+[[gnu::weak]] extern void teardown_main_tls();
 
 [[noreturn]] LLVM_LIBC_FUNCTION(void, quick_exit, (int status)) {
   call_exit_callbacks(at_quick_exit_callbacks);
+  if (teardown_main_tls)
+    teardown_main_tls();
   internal::exit(status);
 }
 
diff --git a/libc/startup/linux/do_start.cpp b/libc/startup/linux/do_start.cpp
index 72060b4adb2148..ff104c7f0d1d2f 100644
--- a/libc/startup/linux/do_start.cpp
+++ b/libc/startup/linux/do_start.cpp
@@ -6,6 +6,7 @@
 //
 //===----------------------------------------------------------------------===//
 #include "startup/linux/do_start.h"
+#include "config/linux/app.h"
 #include "include/llvm-libc-macros/link-macros.h"
 #include "src/__support/OSUtil/syscall.h"
 #include "src/__support/macros/config.h"
@@ -60,6 +61,10 @@ static void call_fini_array_callbacks() {
 }
 
 static ThreadAttributes main_thread_attrib;
+static TLSDescriptor tls;
+// We separate teardown_main_tls from callbacks as callback function themselves
+// may require TLS.
+void teardown_main_tls() { cleanup_tls(tls.addr, tls.size); }
 
 [[noreturn]] void do_start() {
   auto tid = syscall_impl<long>(SYS_gettid);
@@ -122,7 +127,6 @@ static ThreadAttributes main_thread_attrib;
 
   // This descriptor has to be static since its cleanup function cannot
   // capture the context.
-  static TLSDescriptor tls;
   init_tls(tls);
   if (tls.size != 0 && !set_thread_ptr(tls.tp))
     syscall_impl<long>(SYS_exit, 1);
@@ -130,10 +134,7 @@ static ThreadAttributes main_thread_attrib;
   self.attrib = &main_thread_attrib;
   main_thread_attrib.atexit_callback_mgr =
       internal::get_thread_atexit_callback_mgr();
-  // We register the cleanup_tls function to be the last atexit callback to be
-  // invoked. It will tear down the TLS. Other callbacks may depend on TLS (such
-  // as the stack protector canary).
-  atexit([]() { cleanup_tls(tls.addr, tls.size); });
+
   // We want the fini array callbacks to be run after other atexit
   // callbacks are run. So, we register them before running the init
   // array callbacks as they can potentially register their own atexit