From fcce84322742882010d6b524bbfd9d4355b213c7 Mon Sep 17 00:00:00 2001
From: Vitaly Buka <vitalybuka@google.com>
Date: Wed, 13 Dec 2023 15:36:38 -0800
Subject: [PATCH] [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              |  1 +
 compiler-rt/lib/msan/msan.h                |  2 ++
 compiler-rt/lib/msan/msan_allocator.cpp    |  4 ++++
 compiler-rt/lib/msan/msan_allocator.h      |  3 +++
 compiler-rt/lib/msan/msan_interceptors.cpp | 19 -------------------
 compiler-rt/lib/msan/msan_linux.cpp        | 19 +++++++++++++++++++
 6 files changed, 29 insertions(+), 19 deletions(-)

diff --git a/compiler-rt/lib/msan/msan.cpp b/compiler-rt/lib/msan/msan.cpp
index c4f47dea1104..3cdf10c14990 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 b3a9c641b4fb..25fa2212bdad 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 c3b0f8512e82..72a7f980d39f 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 364331d96406..c2a38a401f3b 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 c2d740e7762b..2c9f2c01e14b 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 87a42affd237..04af6f4b27ac 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