From 2721ad46ea6ac5e4adc8171e71adf65c1f2cacba Mon Sep 17 00:00:00 2001
From: Gavin Hayes <gavin@computoid.com>
Date: Wed, 10 Apr 2024 22:26:23 -0400
Subject: [PATCH] Fix fork locking on win32

- __enable_threads / set __threaded in __proc_setup as threads are required for
  win32 subprocess management
- move mmi/fds locking out of pthread_atfork.c into fork.c so it's done anytime
  __threaded is set instead of being dependent of pthreads
- explicitly yoink _pthread_onfork_prepare, _pthread_onfork_parent, and
  _pthread_onfork_child in pthread_create.c so they are linked in in-case they
  are separated from _pthread_atfork

Big Thanks to @dfyz for help with locating the issue, testing, and devising a fix!
---
 libc/proc/fork.c             | 26 ++++++++++++++++++++++----
 libc/proc/proc.c             |  1 +
 libc/thread/pthread_atfork.c |  5 -----
 libc/thread/pthread_create.c |  3 +++
 4 files changed, 26 insertions(+), 9 deletions(-)

diff --git a/libc/proc/fork.c b/libc/proc/fork.c
index 6c4706963..aa437dc6d 100644
--- a/libc/proc/fork.c
+++ b/libc/proc/fork.c
@@ -18,6 +18,7 @@
 ╚─────────────────────────────────────────────────────────────────────────────*/
 #include "libc/assert.h"
 #include "libc/calls/calls.h"
+#include "libc/calls/state.internal.h"
 #include "libc/calls/struct/sigset.h"
 #include "libc/calls/struct/sigset.internal.h"
 #include "libc/calls/syscall-nt.internal.h"
@@ -34,12 +35,29 @@
 #include "libc/nt/thread.h"
 #include "libc/proc/proc.internal.h"
 #include "libc/runtime/internal.h"
+#include "libc/runtime/memtrack.internal.h"
 #include "libc/runtime/runtime.h"
 #include "libc/runtime/syslib.internal.h"
 #include "libc/sysv/consts/sig.h"
 #include "libc/thread/posixthread.internal.h"
 #include "libc/thread/tls.h"
 
+static void _onfork_prepare(void) {
+  if (_weaken(_pthread_onfork_prepare)) {
+    _weaken(_pthread_onfork_prepare)();
+  }
+  __fds_lock();
+  __mmi_lock();
+}
+
+static void _onfork_parent(void) {
+  __mmi_unlock();
+  __fds_unlock();
+  if (_weaken(_pthread_onfork_parent)) {
+    _weaken(_pthread_onfork_parent)();
+  }
+}
+
 int _fork(uint32_t dwCreationFlags) {
   struct Dll *e;
   int ax, dx, tid, parent;
@@ -47,8 +65,8 @@ int _fork(uint32_t dwCreationFlags) {
   BLOCK_SIGNALS;
   if (IsWindows())
     __proc_lock();
-  if (__threaded && _weaken(_pthread_onfork_prepare)) {
-    _weaken(_pthread_onfork_prepare)();
+  if (__threaded) {
+    _onfork_prepare();
   }
   if (!IsWindows()) {
     ax = sys_fork();
@@ -105,8 +123,8 @@ int _fork(uint32_t dwCreationFlags) {
     STRACE("fork() → 0 (child of %d)", parent);
   } else {
     // this is the parent process
-    if (__threaded && _weaken(_pthread_onfork_parent)) {
-      _weaken(_pthread_onfork_parent)();
+    if (__threaded) {
+      _onfork_parent();
     }
     if (IsWindows())
       __proc_unlock();
diff --git a/libc/proc/proc.c b/libc/proc/proc.c
index 33aa5f15c..837957482 100644
--- a/libc/proc/proc.c
+++ b/libc/proc/proc.c
@@ -233,6 +233,7 @@ static textwindows dontinstrument uint32_t __proc_worker(void *arg) {
  * Lazy initializes process tracker data structures and worker.
  */
 static textwindows void __proc_setup(void) {
+  __enable_threads();
   __proc.onbirth = CreateEvent(0, 0, 0, 0);     // auto reset
   __proc.haszombies = CreateEvent(0, 1, 0, 0);  // manual reset
   __proc.thread = CreateThread(0, 65536, __proc_worker, 0,
diff --git a/libc/thread/pthread_atfork.c b/libc/thread/pthread_atfork.c
index 71adf7339..c20ec1859 100644
--- a/libc/thread/pthread_atfork.c
+++ b/libc/thread/pthread_atfork.c
@@ -28,7 +28,6 @@
 #include "libc/macros.internal.h"
 #include "libc/mem/mem.h"
 #include "libc/proc/proc.internal.h"
-#include "libc/runtime/memtrack.internal.h"
 #include "libc/runtime/runtime.h"
 #include "libc/str/str.h"
 #include "libc/thread/posixthread.internal.h"
@@ -66,13 +65,9 @@ static void _pthread_onfork(int i) {
 void _pthread_onfork_prepare(void) {
   _pthread_onfork(0);
   _pthread_lock();
-  __fds_lock();
-  __mmi_lock();
 }
 
 void _pthread_onfork_parent(void) {
-  __mmi_unlock();
-  __fds_unlock();
   _pthread_unlock();
   _pthread_onfork(1);
 }
diff --git a/libc/thread/pthread_create.c b/libc/thread/pthread_create.c
index b59d74845..9a41a9a1f 100644
--- a/libc/thread/pthread_create.c
+++ b/libc/thread/pthread_create.c
@@ -60,6 +60,9 @@ __static_yoink("nsync_mu_trylock");
 __static_yoink("nsync_mu_rlock");
 __static_yoink("nsync_mu_runlock");
 __static_yoink("_pthread_atfork");
+__static_yoink("_pthread_onfork_prepare");
+__static_yoink("_pthread_onfork_parent");
+__static_yoink("_pthread_onfork_child");
 
 #define MAP_ANON_OPENBSD  0x1000
 #define MAP_STACK_OPENBSD 0x4000