[hwasan] Support malloc in atfork.

[eugenis]

Summary

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.

[eugenis]

(+llvm-commits)

[pcc]

LGTM

[eugenis]

rebase (-formatting changes), remove the fork interceptor completely

compiler-rt/lib/hwasan/hwasan.h

 extern int hwasan_report_count;
 
 bool InitShadow();
 void InitializeOsSupport();
 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);
 void *hwasan_realloc(void *ptr, uptr size, StackTrace *stack);
 void *hwasan_reallocarray(void *ptr, uptr nmemb, uptr size, StackTrace *stack);
 void *hwasan_valloc(uptr size, StackTrace *stack);
 void *hwasan_pvalloc(uptr size, StackTrace *stack);
 void *hwasan_aligned_alloc(uptr alignment, uptr size, StackTrace *stack);
if (hwasan_inited)                                     \
     stack.Unwind(pc, bp, nullptr, common_flags()->fast_unwind_on_fatal)
 
 void HwasanTSDInit();
 void HwasanTSDThreadInit();
 void HwasanAtExit();
 
 void HwasanOnDeadlySignal(int signo, void *info, void *context);
 
+void HwasanInstallAtForkHandler();
+
 void UpdateMemoryUsage();
 
 void AppendToErrorMessageBuffer(const char *buffer);
 
 void AndroidTestTlsSlot();
 
 // This is a compiler-generated struct that can be shared between hwasan
 // implementations.

compiler-rt/lib/hwasan/hwasan.cpp

__attribute__((constructor(0))) void __hwasan_init() {
   InstallAtExitHandler(); // Needs __cxa_atexit interceptor.
 
   InitializeCoverage(common_flags()->coverage, common_flags()->coverage_dir);
 
   HwasanTSDInit();
   HwasanTSDThreadInit();
 
   HwasanAllocatorInit();
+  HwasanInstallAtForkHandler();
 
 #if HWASAN_CONTAINS_UBSAN
   __ubsan::InitAsPlugin();
 #endif
 
   VPrintf(1, "HWAddressSanitizer init done\n");
 
   hwasan_init_is_running = 0;

compiler-rt/lib/hwasan/hwasan_allocator.cpp

atomic_store_relaxed(&hwasan_allocator_tagging_enabled,
                        !flags()->disable_allocator_tagging);
   SetAllocatorMayReturnNull(common_flags()->allocator_may_return_null);
   allocator.Init(common_flags()->allocator_release_to_os_interval_ms,
                  GetAliasRegionStart());
   for (uptr i = 0; i < sizeof(tail_magic); i++)
     tail_magic[i] = GetCurrentThread()->GenerateRandomTag();
 }
 
+void HwasanAllocatorLock() { allocator.ForceLock(); }
+
+void HwasanAllocatorUnlock() { allocator.ForceUnlock(); }
+
 void AllocatorSwallowThreadLocalCache(AllocatorCache *cache) {
   allocator.SwallowCache(cache);
 }
 
 static uptr TaggedSize(uptr size) {
   if (!size) size = 1;
   uptr new_size = RoundUpTo(size, kShadowAlignment);
   CHECK_GE(new_size, size);

compiler-rt/lib/hwasan/hwasan_fuchsia.cpp

 void HwasanTSDInit() {}
 void HwasanTSDThreadInit() {}
 
 // On linux, this just would call `atexit(HwasanAtExit)`. The functions in
 // HwasanAtExit are unimplemented for Fuchsia and effectively no-ops, so this
 // 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() {}
 
 }  // namespace __hwasan
 
 extern "C" {
 

compiler-rt/lib/hwasan/hwasan_interceptors.cpp

 INTERCEPTOR(void, longjmp, __hw_jmp_buf env, int val) {
   InternalLongjmp(env[0].__jmpbuf, val);
 }
 #undef SIG_BLOCK
 #undef SIG_SETMASK
 
 #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() {
   // FIXME: ask frontend whether we need to return failure.
   return 0;
 }
 
 } // namespace __hwasan
 
 namespace __hwasan {
 
 void InitializeInterceptors() {
   static int inited = 0;
   CHECK_EQ(inited, 0);
 
-  INTERCEPT_FUNCTION(fork);
-
 #if HWASAN_WITH_INTERCEPTORS
 #if defined(__linux__)
   INTERCEPT_FUNCTION(vfork);
 #endif  // __linux__
   INTERCEPT_FUNCTION(pthread_create);
 #endif
 
   inited = 1;
 }
 } // namespace __hwasan
 
 #endif  // #if !SANITIZER_FUCHSIA

compiler-rt/lib/hwasan/hwasan_linux.cpp

 #  include "hwasan_dynamic_shadow.h"
 #  include "hwasan_interface_internal.h"
 #  include "hwasan_mapping.h"
 #  include "hwasan_report.h"
 #  include "hwasan_thread.h"
 #  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.
 //
 // HWASAN_WITH_INTERCEPTORS=OFF, SANITIZER_ANDROID=OFF
 //   Not currently tested.
 // HWASAN_WITH_INTERCEPTORS=OFF, SANITIZER_ANDROID=ON
 //   Integration tests downstream exist.
 // HWASAN_WITH_INTERCEPTORS=ON, SANITIZER_ANDROID=OFF
if (SANITIZER_LINUX &&
     // Linux.
     ReleaseMemoryPagesToOSAndZeroFill(page_start, page_end);
   } else {
     internal_memset((void *)shadow_start, tag, shadow_size);
   }
   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

+// 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;
+}