From 5c676be217c53553e8e3a8efa7e6d38b355a01ad Mon Sep 17 00:00:00 2001
From: Zachary S <zasample18+github@gmail.com>
Date: Tue, 18 Jun 2024 12:43:06 -0500
Subject: [PATCH] On `unix`, ensure that only one Rust thread calls
 `libc::exit` or returns from `main`.

---
 library/std/src/rt.rs                        |  3 +
 library/std/src/sys/pal/common/exit_guard.rs | 85 ++++++++++++++++++++
 library/std/src/sys/pal/common/mod.rs        |  1 +
 library/std/src/sys/pal/unix/os.rs           |  1 +
 4 files changed, 90 insertions(+)
 create mode 100644 library/std/src/sys/pal/common/exit_guard.rs

diff --git a/library/std/src/rt.rs b/library/std/src/rt.rs
index b03fa1c01f26..376bf3440693 100644
--- a/library/std/src/rt.rs
+++ b/library/std/src/rt.rs
@@ -161,5 +161,8 @@ fn lang_start<T: crate::process::Termination + 'static>(
         argv,
         sigpipe,
     );
+    // Guard against multple threads calling `libc::exit` concurrently.
+    // See the documentation for `unique_thread_exit` for more information.
+    crate::sys::common::exit_guard::unique_thread_exit();
     v
 }
diff --git a/library/std/src/sys/pal/common/exit_guard.rs b/library/std/src/sys/pal/common/exit_guard.rs
new file mode 100644
index 000000000000..fdaa9e186240
--- /dev/null
+++ b/library/std/src/sys/pal/common/exit_guard.rs
@@ -0,0 +1,85 @@
+/// Mitigation for https://github.com/rust-lang/rust/issues/126600
+///
+/// On `unix` (where `libc::exit` is not thread-safe), ensure that only one Rust thread
+/// calls `libc::exit` (or returns from `main`) by calling this function before calling
+/// `libc::exit` (or returning from `main`).
+///
+/// Technically not enough to ensure soundness, since other code directly calling
+/// libc::exit will still race with this.
+///
+/// *This function does not itself call `libc::exit`.* This is so it can also be used
+/// to guard returning from `main`.
+///
+/// This function will return only the first time it is called in a process.
+///
+/// * If it is called again on the same thread as the first call, it will abort.
+/// * If it is called again on a different thread, it will `thread::park()` in a loop
+///   (waiting for the process to exit).
+#[cfg(unix)]
+pub(crate) fn unique_thread_exit() {
+    let this_thread_id = unsafe { libc::gettid() };
+    debug_assert_ne!(this_thread_id, 0, "thread ID cannot be zero");
+    #[cfg(target_has_atomic = "32")]
+    {
+        use crate::sync::atomic::{AtomicI32, Ordering};
+        static EXITING_THREAD_ID: AtomicI32 = AtomicI32::new(0);
+        match EXITING_THREAD_ID.compare_exchange(
+            0,
+            this_thread_id,
+            Ordering::Relaxed,
+            Ordering::Relaxed,
+        ) {
+            Ok(_zero) => {
+                // This is the first thread to call `unique_thread_exit`,
+                // and this is the first time it is called.
+                // Set EXITING_THREAD_ID to this thread's ID (done by the
+                // compare_exchange) and return.
+            }
+            Err(id) if id == this_thread_id => {
+                // This is the first thread to call `unique_thread_exit`,
+                // but this is the second time it is called.
+                // Abort the process.
+                core::panicking::panic_nounwind("std::process::exit called re-entrantly")
+            }
+            Err(_) => {
+                // This is not the first thread to call `unique_thread_exit`.
+                // Park until the process exits.
+                loop {
+                    crate::thread::park();
+                }
+            }
+        }
+    }
+    #[cfg(not(target_has_atomic = "32"))]
+    {
+        use crate::sync::{Mutex, PoisonError};
+        static EXITING_THREAD_ID: Mutex<i32> = Mutex::new(0);
+        let mut exiting_thread_id =
+            EXITING_THREAD_ID.lock().unwrap_or_else(PoisonError::into_inner);
+        if *exiting_thread_id == 0 {
+            // This is the first thread to call `unique_thread_exit`,
+            // and this is the first time it is called.
+            // Set EXITING_THREAD_ID to this thread's ID and return.
+            *exiting_thread_id = this_thread_id;
+        } else if *exiting_thread_id == this_thread_id {
+            // This is the first thread to call `unique_thread_exit`,
+            // but this is the second time it is called.
+            // Abort the process.
+            core::panicking::panic_nounwind("std::process::exit called re-entrantly")
+        } else {
+            // This is not the first thread to call `unique_thread_exit`.
+            // Park until the process exits.
+            drop(exiting_thread_id);
+            loop {
+                crate::thread::park();
+            }
+        }
+    }
+}
+
+/// Mitigation for https://github.com/rust-lang/rust/issues/126600
+/// Not required on platforms where `libc::exit` is thread-safe.
+#[cfg(not(unix))]
+pub(crate) fn unique_thread_exit() {
+    // Mitigation not required on platforms where `exit` is thread-safe.
+}
diff --git a/library/std/src/sys/pal/common/mod.rs b/library/std/src/sys/pal/common/mod.rs
index 29fc0835d766..cc1dceb63e2f 100644
--- a/library/std/src/sys/pal/common/mod.rs
+++ b/library/std/src/sys/pal/common/mod.rs
@@ -11,6 +11,7 @@
 #![allow(dead_code)]
 
 pub mod alloc;
+pub mod exit_guard;
 pub mod small_c_string;
 
 #[cfg(test)]
diff --git a/library/std/src/sys/pal/unix/os.rs b/library/std/src/sys/pal/unix/os.rs
index 2e71ceceb58b..3f598a095c15 100644
--- a/library/std/src/sys/pal/unix/os.rs
+++ b/library/std/src/sys/pal/unix/os.rs
@@ -758,6 +758,7 @@ pub fn home_dir() -> Option<PathBuf> {
 }
 
 pub fn exit(code: i32) -> ! {
+    crate::sys::common::exit_guard::unique_thread_exit();
     unsafe { libc::exit(code as c_int) }
 }