Ensure that only one Rust thread calls libc::exit or returns from main.

Author: zachs18

library/std/src/rt.rs

@@ -161,5 +161,9 @@ 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.
+    #[cfg(unix)]
+    crate::sys::os::unique_thread_exit();
     v
 }

library/std/src/sys/pal/unix/os.rs

@@ -757,7 +757,87 @@ pub fn home_dir() -> Option<PathBuf> {
     }
 }
 
+/// Mitigation for https://github.com/rust-lang/rust/issues/126600
+///
+/// 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).
+/// * If it is called in a situation where `std::thread::current()` fails, it will abort.
+pub(crate) fn unique_thread_exit() {
+    let this_thread_id =
+        crate::thread::try_current().unwrap_or_else(|| crate::process::abort()).id().as_u64().get();
+    // Defense against refactors of `.as_u64()`
+    debug_assert_ne!(this_thread_id, 0, "thread ID cannot be zero");
+    #[cfg(target_has_atomic = "64")]
+    {
+        use crate::sync::atomic::{AtomicU64, Ordering};
+        static EXITING_THREAD_ID: AtomicU64 = AtomicU64::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.
+                crate::process::abort();
+            }
+            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 = "64"))]
+    {
+        use crate::sync::{Mutex, PoisonError};
+        static EXITING_THREAD_ID: Mutex<u64> = 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.
+            crate::process::abort();
+        } 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();
+            }
+        }
+    }
+}
+
 pub fn exit(code: i32) -> ! {
+    unique_thread_exit();
     unsafe { libc::exit(code as c_int) }
 }