[msan] Use pthread_atfork instead of interceptor (#75398)

This is done for consistency with other sanitizers.
Also lock the allocator.

compiler-rt/lib/msan/msan.cpp

@@ -449,6 +449,7 @@ void __msan_init() {
   __sanitizer_set_report_path(common_flags()->log_path);
 
   InitializeInterceptors();
+  InstallAtForkHandler();
   CheckASLR();
   InitTlsSize();
   InstallDeadlySignalHandlers(MsanOnDeadlySignal);

compiler-rt/lib/msan/msan.h

@@ -336,6 +336,8 @@ void *MsanTSDGet();
 void MsanTSDSet(void *tsd);
 void MsanTSDDtor(void *tsd);
 
+void InstallAtForkHandler();
+
 }  // namespace __msan
 
 #endif  // MSAN_H

compiler-rt/lib/msan/msan_allocator.cpp

@@ -159,6 +159,10 @@ void MsanAllocatorInit() {
     max_malloc_size = kMaxAllowedMallocSize;
 }
 
+void LockAllocator() { allocator.ForceLock(); }
+
+void UnlockAllocator() { allocator.ForceUnlock(); }
+
 AllocatorCache *GetAllocatorCache(MsanThreadLocalMallocStorage *ms) {
   CHECK(ms);
   CHECK_LE(sizeof(AllocatorCache), sizeof(ms->allocator_cache));

compiler-rt/lib/msan/msan_allocator.h

@@ -28,5 +28,8 @@ struct MsanThreadLocalMallocStorage {
   MsanThreadLocalMallocStorage() {}
 };
 
+void LockAllocator();
+void UnlockAllocator();
+
 } // namespace __msan
 #endif // MSAN_ALLOCATOR_H

compiler-rt/lib/msan/msan_interceptors.cpp

@@ -1326,24 +1326,6 @@ static int setup_at_exit_wrapper(void(*f)(), void *arg, void *dso) {
   return res;
 }
 
-static void BeforeFork() {
-  StackDepotLockAll();
-  ChainedOriginDepotLockAll();
-}
-
-static void AfterFork() {
-  ChainedOriginDepotUnlockAll();
-  StackDepotUnlockAll();
-}
-
-INTERCEPTOR(int, fork, void) {
-  ENSURE_MSAN_INITED();
-  BeforeFork();
-  int pid = REAL(fork)();
-  AfterFork();
-  return pid;
-}
-
 // NetBSD ships with openpty(3) in -lutil, that needs to be prebuilt explicitly
 // with MSan.
 #if SANITIZER_LINUX
@@ -1933,7 +1915,6 @@ void InitializeInterceptors() {
   INTERCEPT_FUNCTION(atexit);
   INTERCEPT_FUNCTION(__cxa_atexit);
   INTERCEPT_FUNCTION(shmat);
-  INTERCEPT_FUNCTION(fork);
   MSAN_MAYBE_INTERCEPT_OPENPTY;
   MSAN_MAYBE_INTERCEPT_FORKPTY;
 

compiler-rt/lib/msan/msan_linux.cpp

@@ -26,10 +26,13 @@
 #  include <unwind.h>
 
 #  include "msan.h"
+#  include "msan_allocator.h"
+#  include "msan_chained_origin_depot.h"
 #  include "msan_report.h"
 #  include "msan_thread.h"
 #  include "sanitizer_common/sanitizer_common.h"
 #  include "sanitizer_common/sanitizer_procmaps.h"
+#  include "sanitizer_common/sanitizer_stackdepot.h"
 
 namespace __msan {
 
@@ -255,6 +258,22 @@ void MsanTSDDtor(void *tsd) {
 }
 #endif
 
+void InstallAtForkHandler() {
+  auto before = []() {
+    // Usually we lock ThreadRegistry, but msan does not have one.
+    LockAllocator();
+    StackDepotLockAll();
+    ChainedOriginDepotLockAll();
+  };
+  auto after = []() {
+    ChainedOriginDepotUnlockAll();
+    StackDepotUnlockAll();
+    UnlockAllocator();
+    // Usually we unlock ThreadRegistry, but msan does not have one.
+  };
+  pthread_atfork(before, after, after);
+}
+
 } // namespace __msan
 
 #endif // SANITIZER_FREEBSD || SANITIZER_LINUX || SANITIZER_NETBSD