[asan] Install `pthread_atfork` (#75290)

This prevents deadlocks in forked process
if parent had more then one running threads.

6 files changed, 32 insertions, 1 deletions

diff --git a/compiler-rt/lib/asan/asan_fuchsia.cpp b/compiler-rt/lib/asan/asan_fuchsia.cpp
index 2b15504..12625e9 100644
--- a/compiler-rt/lib/asan/asan_fuchsia.cpp
+++ b/compiler-rt/lib/asan/asan_fuchsia.cpp
@@ -240,6 +240,8 @@ void FlushUnneededASanShadowMemory(uptr p, uptr size) {
 // So this doesn't install any atexit hook like on other platforms.
 void InstallAtExitCheckLeaks() {}
 
+void InstallAtForkHandler() {}
+
 }  // namespace __asan
 
 namespace __lsan {
diff --git a/compiler-rt/lib/asan/asan_internal.h b/compiler-rt/lib/asan/asan_internal.h
index 5b97e77..2944ebe 100644
--- a/compiler-rt/lib/asan/asan_internal.h
+++ b/compiler-rt/lib/asan/asan_internal.h
@@ -126,6 +126,7 @@ void *AsanDlSymNext(const char *sym);
 bool HandleDlopenInit();
 
 void InstallAtExitCheckLeaks();
+void InstallAtForkHandler();
 
 #define ASAN_ON_ERROR() \
   if (&__asan_on_error) \
diff --git a/compiler-rt/lib/asan/asan_posix.cpp b/compiler-rt/lib/asan/asan_posix.cpp
index e1f6664..206551b 100644
--- a/compiler-rt/lib/asan/asan_posix.cpp
+++ b/compiler-rt/lib/asan/asan_posix.cpp
@@ -148,6 +148,30 @@ void PlatformTSDDtor(void *tsd) {
 }
 #endif
 
+void InstallAtForkHandler() {
+  auto before = []() {
+    if (CAN_SANITIZE_LEAKS) {
+      __lsan::LockGlobal();
+    }
+    // `_lsan` functions defined regardless of `CAN_SANITIZE_LEAKS` and lock the
+    // stuff we need.
+    __lsan::LockThreads();
+    __lsan::LockAllocator();
+    StackDepotLockAll();
+  };
+  auto after = []() {
+    StackDepotUnlockAll();
+    // `_lsan` functions defined regardless of `CAN_SANITIZE_LEAKS` and unlock
+    // the stuff we need.
+    __lsan::UnlockAllocator();
+    __lsan::UnlockThreads();
+    if (CAN_SANITIZE_LEAKS) {
+      __lsan::UnlockGlobal();
+    }
+  };
+  pthread_atfork(before, after, after);
+}
+
 void InstallAtExitCheckLeaks() {
   if (CAN_SANITIZE_LEAKS) {
     if (common_flags()->detect_leaks && common_flags()->leak_check_at_exit) {
diff --git a/compiler-rt/lib/asan/asan_rtl.cpp b/compiler-rt/lib/asan/asan_rtl.cpp
index b28f9f1..a61deed 100644
--- a/compiler-rt/lib/asan/asan_rtl.cpp
+++ b/compiler-rt/lib/asan/asan_rtl.cpp
@@ -493,6 +493,8 @@ static bool AsanInitInternal() {
     InstallAtExitCheckLeaks();
   }
 
+  InstallAtForkHandler();
+
 #if CAN_SANITIZE_UB
   __ubsan::InitAsPlugin();
 #endif
diff --git a/compiler-rt/lib/asan/asan_win.cpp b/compiler-rt/lib/asan/asan_win.cpp
index d5a30f4..f16ce67 100644
--- a/compiler-rt/lib/asan/asan_win.cpp
+++ b/compiler-rt/lib/asan/asan_win.cpp
@@ -203,6 +203,8 @@ void InitializePlatformInterceptors() {
 
 void InstallAtExitCheckLeaks() {}
 
+void InstallAtForkHandler() {}
+
 void AsanApplyToGlobals(globals_op_fptr op, const void *needle) {
   UNIMPLEMENTED();
 }
diff --git a/compiler-rt/test/sanitizer_common/TestCases/Posix/fork_threaded.c b/compiler-rt/test/sanitizer_common/TestCases/Posix/fork_threaded.c
index 7592269..673d353 100644
--- a/compiler-rt/test/sanitizer_common/TestCases/Posix/fork_threaded.c
+++ b/compiler-rt/test/sanitizer_common/TestCases/Posix/fork_threaded.c
@@ -1,6 +1,6 @@
 // RUN: %clang -O0 %s -o %t && %env_tool_opts=die_after_fork=0 %run %t
 
-// UNSUPPORTED: asan, hwasan
+// UNSUPPORTED: hwasan
 
 // The test uses pthread barriers which are not available on Darwin.
 // UNSUPPORTED: darwin