[msan] Use pthread_atfork instead of interceptor

[vitalybuka]

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

[llvmbot]

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

Author: Vitaly Buka (vitalybuka)

Changes

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

---

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

6 Files Affected:

• (modified) compiler-rt/lib/msan/msan.cpp (+1)
• (modified) compiler-rt/lib/msan/msan.h (+2)
• (modified) compiler-rt/lib/msan/msan_allocator.cpp (+4)
• (modified) compiler-rt/lib/msan/msan_allocator.h (+3)
• (modified) compiler-rt/lib/msan/msan_interceptors.cpp (-19)
• (modified) compiler-rt/lib/msan/msan_linux.cpp (+19)

diff --git a/compiler-rt/lib/msan/msan.cpp b/compiler-rt/lib/msan/msan.cpp
index c4f47dea110432..3cdf10c149902c 100644
--- a/compiler-rt/lib/msan/msan.cpp
+++ b/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);
diff --git a/compiler-rt/lib/msan/msan.h b/compiler-rt/lib/msan/msan.h
index b3a9c641b4fb29..25fa2212bdadd3 100644
--- a/compiler-rt/lib/msan/msan.h
+++ b/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
diff --git a/compiler-rt/lib/msan/msan_allocator.cpp b/compiler-rt/lib/msan/msan_allocator.cpp
index c3b0f8512e82d8..72a7f980d39fb0 100644
--- a/compiler-rt/lib/msan/msan_allocator.cpp
+++ b/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));
diff --git a/compiler-rt/lib/msan/msan_allocator.h b/compiler-rt/lib/msan/msan_allocator.h
index 364331d964068e..c2a38a401f3b6b 100644
--- a/compiler-rt/lib/msan/msan_allocator.h
+++ b/compiler-rt/lib/msan/msan_allocator.h
@@ -28,5 +28,8 @@ struct MsanThreadLocalMallocStorage {
   MsanThreadLocalMallocStorage() {}
 };
 
+void LockAllocator();
+void UnlockAllocator();
+
 } // namespace __msan
 #endif // MSAN_ALLOCATOR_H
diff --git a/compiler-rt/lib/msan/msan_interceptors.cpp b/compiler-rt/lib/msan/msan_interceptors.cpp
index c2d740e7762b4b..2c9f2c01e14b06 100644
--- a/compiler-rt/lib/msan/msan_interceptors.cpp
+++ b/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;
 
diff --git a/compiler-rt/lib/msan/msan_linux.cpp b/compiler-rt/lib/msan/msan_linux.cpp
index 87a42affd237f4..04af6f4b27ac89 100644
--- a/compiler-rt/lib/msan/msan_linux.cpp
+++ b/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