[PATCH] D72364: [scudo][standalone] Modify malloc_{enable,disable} wrt fork

[Kostya Kortchinsky via Phabricator via llvm-commits]

cryptoad updated this revision to Diff 236834.
cryptoad marked 9 inline comments as done.
cryptoad added a comment.

Updating test to exercise both disable and no-disable pre-fork.
Added an `_exit` in fork child.


Repository:
  rG LLVM Github Monorepo

CHANGES SINCE LAST ACTION
  https://reviews.llvm.org/D72364/new/

https://reviews.llvm.org/D72364

Files:
  compiler-rt/lib/scudo/standalone/tests/wrappers_c_test.cpp
  compiler-rt/lib/scudo/standalone/wrappers_c.inc


Index: compiler-rt/lib/scudo/standalone/wrappers_c.inc
===================================================================
--- compiler-rt/lib/scudo/standalone/wrappers_c.inc
+++ compiler-rt/lib/scudo/standalone/wrappers_c.inc
@@ -141,12 +141,33 @@
   return 0;
 }
 
+static pthread_once_t SCUDO_PREFIX(OnceControl) = PTHREAD_ONCE_INIT;
+static bool SCUDO_PREFIX(Disabled) = false;
+static scudo::HybridMutex SCUDO_PREFIX(Mutex);
+
+INTERFACE WEAK void SCUDO_PREFIX(malloc_enable)() {
+  scudo::ScopedLock L(SCUDO_PREFIX(Mutex));
+  if (!SCUDO_PREFIX(Disabled))
+    return;
+  SCUDO_ALLOCATOR.enable();
+  SCUDO_PREFIX(Disabled) = false;
+}
+
+// Execute malloc_enable for a fork child.
+static void SCUDO_PREFIX(setupAtFork)() {
+  pthread_atfork(nullptr, nullptr, SCUDO_PREFIX(malloc_enable));
+}
+
 INTERFACE WEAK void SCUDO_PREFIX(malloc_disable)() {
+  scudo::ScopedLock L(SCUDO_PREFIX(Mutex));
+  if (SCUDO_PREFIX(Disabled))
+    return;
+  CHECK_EQ(pthread_once(&SCUDO_PREFIX(OnceControl), SCUDO_PREFIX(setupAtFork)),
+           0);
   SCUDO_ALLOCATOR.disable();
+  SCUDO_PREFIX(Disabled) = true;
 }
 
-INTERFACE WEAK void SCUDO_PREFIX(malloc_enable)() { SCUDO_ALLOCATOR.enable(); }
-
 INTERFACE WEAK int SCUDO_PREFIX(mallopt)(int param, UNUSED int value) {
   if (param == M_DECAY_TIME) {
     // TODO(kostyak): set release_to_os_interval_ms accordingly.
Index: compiler-rt/lib/scudo/standalone/tests/wrappers_c_test.cpp
===================================================================
--- compiler-rt/lib/scudo/standalone/tests/wrappers_c_test.cpp
+++ compiler-rt/lib/scudo/standalone/tests/wrappers_c_test.cpp
@@ -310,4 +310,31 @@
   fclose(F);
   EXPECT_EQ(strncmp(Buffer, "<malloc version=\"scudo-", 23), 0);
 }
+
+// If a fork occurs while the allocator was disabled, we want to ensure that the
+// child can use allocation primitives without enabling the allocator first.
+TEST(ScudoWrappersCTest, MallocPostFork) {
+  pid_t Pid;
+  void *P;
+  for (bool DisablePreFork : {false, true}) {
+    if (DisablePreFork)
+      malloc_disable();
+    Pid = fork();
+    EXPECT_GE(Pid, 0);
+    if (Pid == 0) {
+      P = malloc(Size);
+      EXPECT_NE(P, nullptr);
+      memset(P, 0x42, Size);
+      free(P);
+      _exit(0);
+    }
+    waitpid(Pid, nullptr, 0);
+    if (DisablePreFork)
+      malloc_enable();
+    P = malloc(Size);
+    EXPECT_NE(P, nullptr);
+    memset(P, 0x42, Size);
+    free(P);
+  }
+}
 #endif


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D72364.236834.patch
Type: text/x-patch
Size: 2436 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20200108/b5cf4447/attachment.bin>