From 4b202f4faf40fd7af8b84491360186aed8ce5733 Mon Sep 17 00:00:00 2001
From: Kyle Evans <kevans@FreeBSD.org>
Date: Wed, 13 Nov 2024 19:33:44 -0600
Subject: [PATCH] libthr: allow very early atfork registration

LSan wants to be able to register atfork handlers at __lsan_init time,
which can happen either at the first intercepted function call or in a
.preinit_array function.  Both of these end up being very early in rtld
and executed with the bind lock held, which ends up causing problems
when we go to _libpthread_init().

Instead of requiring libpthread to be initialized, just insert the new
atfork handler straight into the list if it's not ready yet.  The
critical section and locking should not be necessary if we're really
executing this early, as there won't be any threads to contend with.

Reviewed by:	kib (earlier version), markj
Differential Revision:	https://reviews.freebsd.org/D47349
---
 lib/libthr/tests/atfork_test.c | 53 ++++++++++++++++++++++++++++++++++
 lib/libthr/thread/thr_fork.c   | 20 ++++++++-----
 lib/libthr/thread/thr_init.c   |  1 -
 3 files changed, 65 insertions(+), 9 deletions(-)

diff --git a/lib/libthr/tests/atfork_test.c b/lib/libthr/tests/atfork_test.c
index 5133330b124..cb0fcb7e62d 100644
--- a/lib/libthr/tests/atfork_test.c
+++ b/lib/libthr/tests/atfork_test.c
@@ -25,6 +25,58 @@ static int child;
 static int forked;
 static int parent;
 
+/*
+ * We'll disable prefork unless we're specifically running the preinit test to
+ * be sure that we don't mess up any other tests' results.
+ */
+static bool prefork_enabled;
+
+static void
+prefork(void)
+{
+	if (prefork_enabled)
+		forked++;
+}
+
+static void
+registrar(void)
+{
+	pthread_atfork(prefork, NULL, NULL);
+}
+
+static __attribute__((section(".preinit_array"), used))
+void (*preinitfn)(void) = &registrar;
+
+/*
+ * preinit_atfork() just enables the prepare handler that we registered in a
+ * .preinit_array entry and checks that forking actually invoked that callback.
+ * We don't bother testing all three callbacks here because the implementation
+ * doesn't really lend itself to the kind of error where we only have a partial
+ * set of callbacks registered.
+ */
+ATF_TC(preinit_atfork);
+ATF_TC_HEAD(preinit_atfork, tc)
+{
+	atf_tc_set_md_var(tc, "descr",
+	    "Checks that atfork callbacks may be registered in .preinit_array functions");
+}
+ATF_TC_BODY(preinit_atfork, tc)
+{
+	pid_t p;
+
+	(void)signal(SIGCHLD, SIG_IGN);
+	prefork_enabled = true;
+	p = fork();
+
+	ATF_REQUIRE(p >= 0);
+	if (p == 0)
+		_exit(0);
+
+	prefork_enabled = false;
+
+	ATF_REQUIRE(forked != 0);
+}
+
 static void
 basic_prepare(void)
 {
@@ -221,6 +273,7 @@ ATF_TC_BODY(multi_atfork, tc)
 
 ATF_TP_ADD_TCS(tp)
 {
+	ATF_TP_ADD_TC(tp, preinit_atfork);
 	ATF_TP_ADD_TC(tp, basic_atfork);
 	ATF_TP_ADD_TC(tp, multi_atfork);
 	return (atf_no_error());
diff --git a/lib/libthr/thread/thr_fork.c b/lib/libthr/thread/thr_fork.c
index 614841cc314..aa54b0444cf 100644
--- a/lib/libthr/thread/thr_fork.c
+++ b/lib/libthr/thread/thr_fork.c
@@ -84,20 +84,24 @@ _thr_atfork(void (*prepare)(void), void (*parent)(void),
 	struct pthread *curthread;
 	struct pthread_atfork *af;
 
-	_thr_check_init();
-
 	if ((af = malloc(sizeof(struct pthread_atfork))) == NULL)
 		return (ENOMEM);
 
-	curthread = _get_curthread();
 	af->prepare = prepare;
 	af->parent = parent;
 	af->child = child;
-	THR_CRITICAL_ENTER(curthread);
-	_thr_rwl_wrlock(&_thr_atfork_lock);
-	TAILQ_INSERT_TAIL(&_thr_atfork_list, af, qe);
-	_thr_rwl_unlock(&_thr_atfork_lock);
-	THR_CRITICAL_LEAVE(curthread);
+
+	if (_thr_initial != NULL) {
+		curthread = _get_curthread();
+		THR_CRITICAL_ENTER(curthread);
+		_thr_rwl_wrlock(&_thr_atfork_lock);
+		TAILQ_INSERT_TAIL(&_thr_atfork_list, af, qe);
+		_thr_rwl_unlock(&_thr_atfork_lock);
+		THR_CRITICAL_LEAVE(curthread);
+	} else {
+		TAILQ_INSERT_TAIL(&_thr_atfork_list, af, qe);
+	}
+
 	return (0);
 }
 
diff --git a/lib/libthr/thread/thr_init.c b/lib/libthr/thread/thr_init.c
index 007cf5be81c..ff59288d919 100644
--- a/lib/libthr/thread/thr_init.c
+++ b/lib/libthr/thread/thr_init.c
@@ -519,7 +519,6 @@ init_private(void)
 		env = getenv("LIBPTHREAD_QUEUE_FIFO");
 		if (env)
 			_thr_queuefifo = atoi(env);
-		TAILQ_INIT(&_thr_atfork_list);
 		env = getenv("LIBPTHREAD_UMTX_MIN_TIMEOUT");
 		if (env) {
 			char *endptr;