[hwasan] Improve support of forking with threads

[vitalybuka]

Lock Lsan and Thread related date at_fork.

Clean shadow before thread starts, forked process may reuse already
mapped stack of 'lost' parent thread for new threads.

[llvmbot]

@llvm/pr-subscribers-compiler-rt-sanitizer

Author: Vitaly Buka (vitalybuka)

Changes

Lock Lsan and Thread related date at_fork.

Clean shadow before thread starts, forked process may reuse already
mapped stack of 'lost' parent thread for new threads.

---

Full diff: https://github.com/llvm/llvm-project/pull/75291.diff

3 Files Affected:

• (modified) compiler-rt/lib/hwasan/hwasan_linux.cpp (+14-2)
• (modified) compiler-rt/lib/hwasan/hwasan_thread.cpp (+1)
• (modified) compiler-rt/test/sanitizer_common/TestCases/Posix/fork_threaded.c (+4-2)

diff --git a/compiler-rt/lib/hwasan/hwasan_linux.cpp b/compiler-rt/lib/hwasan/hwasan_linux.cpp
index f01fa42764134..e0db24b9519f7 100644
--- a/compiler-rt/lib/hwasan/hwasan_linux.cpp
+++ b/compiler-rt/lib/hwasan/hwasan_linux.cpp
@@ -523,12 +523,24 @@ uptr TagMemoryAligned(uptr p, uptr size, tag_t tag) {
 
 void HwasanInstallAtForkHandler() {
   auto before = []() {
-    HwasanAllocatorLock();
+    if (CAN_SANITIZE_LEAKS) {
+      __lsan::LockGlobal();
+    }
+    // `_lsan` functions defined regardless of `CAN_SANITIZE_LEAKS` and do the
+    // job.
+    __lsan::LockThreads();
+    __lsan::LockAllocator();
     StackDepotLockAll();
   };
   auto after = []() {
     StackDepotUnlockAll();
-    HwasanAllocatorUnlock();
+    // `_lsan` functions defined regardless of `CAN_SANITIZE_LEAKS` and do the
+    // job.
+    __lsan::UnlockAllocator();
+    __lsan::UnlockThreads();
+    if (CAN_SANITIZE_LEAKS) {
+      __lsan::UnlockGlobal();
+    }
   };
   pthread_atfork(before, after, after);
 }
diff --git a/compiler-rt/lib/hwasan/hwasan_thread.cpp b/compiler-rt/lib/hwasan/hwasan_thread.cpp
index ce36547580e6e..3e14a718513d7 100644
--- a/compiler-rt/lib/hwasan/hwasan_thread.cpp
+++ b/compiler-rt/lib/hwasan/hwasan_thread.cpp
@@ -68,6 +68,7 @@ void Thread::Init(uptr stack_buffer_start, uptr stack_buffer_size,
     }
     Print("Creating  : ");
   }
+  ClearShadowForThreadStackAndTLS();
 }
 
 void Thread::InitStackRingBuffer(uptr stack_buffer_start,
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 1a52702c5de8c..f27859f815ff9 100644
--- a/compiler-rt/test/sanitizer_common/TestCases/Posix/fork_threaded.c
+++ b/compiler-rt/test/sanitizer_common/TestCases/Posix/fork_threaded.c
@@ -1,7 +1,5 @@
 // RUN: %clang -O0 %s -o %t && %env_tool_opts=die_after_fork=0 %run %t
 
-// UNSUPPORTED: hwasan
-
 // Forking in multithread environment is unsupported. However we already have
 // some workarounds, and will add more, so this is the test.
 // The test try to check two things:
@@ -58,6 +56,10 @@ NOSAN static void *inchild(void *arg) {
 }
 
 int main(void) {
+#if __has_feature(hwaddress_sanitizer)
+  __hwasan_enable_allocator_tagging();
+#endif
+
   pid_t pid;
 
   pthread_barrier_init(&bar, NULL, 2);