[hwasan] Support malloc in atfork.

Before this change we were locking the StackDepot in the fork()
interceptor. This results in a deadlock when allocator functions are
used in a pthread_atfork() callback.

Instead, set up a pthread_atfork() callback at init that locks/unlocks
both StackDepot and the allocator. Since our callback is set up very
early, the pre-fork callback is executed late, and both post-fork ones
are executed early, which works perfect for us.

Differential Revision: https://reviews.llvm.org/D108063

compiler-rt/lib/hwasan/hwasan.cpp

@@ -360,6 +360,7 @@ __attribute__((constructor(0))) void __hwasan_init() {
   HwasanTSDThreadInit();
 
   HwasanAllocatorInit();
+  HwasanInstallAtForkHandler();
 
 #if HWASAN_CONTAINS_UBSAN
   __ubsan::InitAsPlugin();

compiler-rt/lib/hwasan/hwasan.h

@@ -107,6 +107,8 @@ void InitThreads();
 void InitializeInterceptors();
 
 void HwasanAllocatorInit();
+void HwasanAllocatorLock();
+void HwasanAllocatorUnlock();
 
 void *hwasan_malloc(uptr size, StackTrace *stack);
 void *hwasan_calloc(uptr nmemb, uptr size, StackTrace *stack);
@@ -140,6 +142,8 @@ void HwasanAtExit();
 
 void HwasanOnDeadlySignal(int signo, void *info, void *context);
 
+void HwasanInstallAtForkHandler();
+
 void UpdateMemoryUsage();
 
 void AppendToErrorMessageBuffer(const char *buffer);

compiler-rt/lib/hwasan/hwasan_allocator.cpp

@@ -107,6 +107,10 @@ void HwasanAllocatorInit() {
     tail_magic[i] = GetCurrentThread()->GenerateRandomTag();
 }
 
+void HwasanAllocatorLock() { allocator.ForceLock(); }
+
+void HwasanAllocatorUnlock() { allocator.ForceUnlock(); }
+
 void AllocatorSwallowThreadLocalCache(AllocatorCache *cache) {
   allocator.SwallowCache(cache);
 }

compiler-rt/lib/hwasan/hwasan_fuchsia.cpp

@@ -180,6 +180,8 @@ void HwasanTSDThreadInit() {}
 // function is unneeded.
 void InstallAtExitHandler() {}
 
+void HwasanInstallAtForkHandler() {}
+
 // TODO(fxbug.dev/81499): Once we finalize the tagged pointer ABI in zircon, we should come back
 // here and implement the appropriate check that TBI is enabled.
 void InitializeOsSupport() {}

compiler-rt/lib/hwasan/hwasan_interceptors.cpp

@@ -125,22 +125,6 @@ INTERCEPTOR(void, longjmp, __hw_jmp_buf env, int val) {
 
 #endif // HWASAN_WITH_INTERCEPTORS && __aarch64__
 
-static void BeforeFork() {
-  StackDepotLockAll();
-}
-
-static void AfterFork() {
-  StackDepotUnlockAll();
-}
-
-INTERCEPTOR(int, fork, void) {
-  ENSURE_HWASAN_INITED();
-  BeforeFork();
-  int pid = REAL(fork)();
-  AfterFork();
-  return pid;
-}
-
 namespace __hwasan {
 
 int OnExit() {
@@ -156,8 +140,6 @@ void InitializeInterceptors() {
   static int inited = 0;
   CHECK_EQ(inited, 0);
 
-  INTERCEPT_FUNCTION(fork);
-
 #if HWASAN_WITH_INTERCEPTORS
 #if defined(__linux__)
   INTERCEPT_FUNCTION(vfork);

compiler-rt/lib/hwasan/hwasan_linux.cpp

@@ -38,6 +38,7 @@
 #  include "hwasan_thread_list.h"
 #  include "sanitizer_common/sanitizer_common.h"
 #  include "sanitizer_common/sanitizer_procmaps.h"
+#  include "sanitizer_common/sanitizer_stackdepot.h"
 
 // Configurations of HWASAN_WITH_INTERCEPTORS and SANITIZER_ANDROID.
 //
@@ -428,6 +429,18 @@ uptr TagMemoryAligned(uptr p, uptr size, tag_t tag) {
   return AddTagToPointer(p, tag);
 }
 
+void HwasanInstallAtForkHandler() {
+  auto before = []() {
+    HwasanAllocatorLock();
+    StackDepotLockAll();
+  };
+  auto after = []() {
+    StackDepotUnlockAll();
+    HwasanAllocatorUnlock();
+  };
+  pthread_atfork(before, after, after);
+}
+
 }  // namespace __hwasan
 
 #endif  // SANITIZER_FREEBSD || SANITIZER_LINUX || SANITIZER_NETBSD

compiler-rt/test/hwasan/TestCases/Linux/atfork.cpp

@@ -0,0 +1,34 @@
+// RUN: %clang_hwasan -O0 %s -o %t && %run %t 2>&1
+
+// REQUIRES: aarch64-target-arch || x86_64-target-arch
+// REQUIRES: pointer-tagging
+
+#include <assert.h>
+#include <pthread.h>
+#include <sanitizer/hwasan_interface.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <sys/types.h>
+#include <sys/wait.h>
+#include <unistd.h>
+
+void *volatile sink;
+
+int main(int argc, char **argv) {
+  pthread_atfork(nullptr, nullptr, []() {
+    alarm(5);
+    sink = malloc(10);
+  });
+  int pid = fork();
+  if (pid) {
+    int wstatus;
+    do {
+      waitpid(pid, &wstatus, 0);
+    } while (!WIFEXITED(wstatus) && !WIFSIGNALED(wstatus));
+    if (!WIFEXITED(wstatus) || WEXITSTATUS(wstatus)) {
+      fprintf(stderr, "abnormal exit\n");
+      return 1;
+    }
+  }
+  return 0;
+}