std: detect stack overflows in TLS destructors on UNIX

library/std/src/rt.rs

@@ -117,6 +117,7 @@ pub(crate) fn thread_cleanup() {
     // print a nice message.
     panic::catch_unwind(|| {
         crate::thread::drop_current();
+        crate::sys::thread_cleanup();
     })
     .unwrap_or_else(handle_rt_panic);
 }

library/std/src/sys/pal/hermit/mod.rs

@@ -69,6 +69,8 @@ pub unsafe fn init(argc: isize, argv: *const *const u8, _sigpipe: u8) {
     }
 }
 
+pub fn thread_cleanup() {}
+
 // SAFETY: must be called only once during runtime cleanup.
 // NOTE: this is not guaranteed to run, for example when the program aborts.
 pub unsafe fn cleanup() {}

library/std/src/sys/pal/sgx/mod.rs

@@ -37,6 +37,8 @@ pub unsafe fn init(argc: isize, argv: *const *const u8, _sigpipe: u8) {
     }
 }
 
+pub fn thread_cleanup() {}
+
 // SAFETY: must be called only once during runtime cleanup.
 // NOTE: this is not guaranteed to run, for example when the program aborts.
 pub unsafe fn cleanup() {}

library/std/src/sys/pal/solid/mod.rs

@@ -38,6 +38,8 @@ pub mod time;
 // NOTE: this is not guaranteed to run, for example when Rust code is called externally.
 pub unsafe fn init(_argc: isize, _argv: *const *const u8, _sigpipe: u8) {}
 
+pub fn thread_cleanup() {}
+
 // SAFETY: must be called only once during runtime cleanup.
 pub unsafe fn cleanup() {}
 

library/std/src/sys/pal/teeos/mod.rs

@@ -45,6 +45,8 @@ pub fn abort_internal() -> ! {
 // so this should never be called.
 pub fn init(argc: isize, argv: *const *const u8, sigpipe: u8) {}
 
+pub fn thread_cleanup() {}
+
 // SAFETY: must be called only once during runtime cleanup.
 // this is not guaranteed to run, for example when the program aborts.
 pub unsafe fn cleanup() {

library/std/src/sys/pal/uefi/mod.rs

@@ -76,6 +76,8 @@ pub(crate) unsafe fn init(argc: isize, argv: *const *const u8, _sigpipe: u8) {
     }
 }
 
+pub fn thread_cleanup() {}
+
 /// # SAFETY
 /// this is not guaranteed to run, for example when the program aborts.
 /// - must be called only once during runtime cleanup.

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

@@ -55,7 +55,7 @@ pub unsafe fn init(argc: isize, argv: *const *const u8, sigpipe: u8) {
     // behavior.
     reset_sigpipe(sigpipe);
 
-    stack_overflow::init();
+    stack_overflow::protect(true);
     args::init(argc, argv);
 
     // Normally, `thread::spawn` will call `Thread::set_name` but since this thread
@@ -229,12 +229,14 @@ pub(crate) fn on_broken_pipe_flag_used() -> bool {
     ON_BROKEN_PIPE_FLAG_USED.load(crate::sync::atomic::Ordering::Relaxed)
 }
 
-// SAFETY: must be called only once during runtime cleanup.
-// NOTE: this is not guaranteed to run, for example when the program aborts.
-pub unsafe fn cleanup() {
+pub fn thread_cleanup() {
     stack_overflow::cleanup();
 }
 
+// SAFETY: must be called only once during runtime cleanup.
+// NOTE: this is not guaranteed to run, for example when the program aborts.
+pub unsafe fn cleanup() {}
+
 #[allow(unused_imports)]
 pub use libc::signal;
 

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

@@ -1,29 +1,6 @@
 #![cfg_attr(test, allow(dead_code))]
 
-pub use self::imp::{cleanup, init};
-use self::imp::{drop_handler, make_handler};
-
-pub struct Handler {
-    data: *mut libc::c_void,
-}
-
-impl Handler {
-    pub unsafe fn new() -> Handler {
-        make_handler(false)
-    }
-
-    fn null() -> Handler {
-        Handler { data: crate::ptr::null_mut() }
-    }
-}
-
-impl Drop for Handler {
-    fn drop(&mut self) {
-        unsafe {
-            drop_handler(self.data);
-        }
-    }
-}
+pub use self::imp::{cleanup, protect};
 
 #[cfg(any(
     target_os = "linux",
@@ -45,22 +22,23 @@ mod imp {
     #[cfg(all(target_os = "linux", target_env = "gnu"))]
     use libc::{mmap64, mprotect, munmap};
 
-    use super::Handler;
-    use crate::cell::Cell;
     use crate::ops::Range;
     use crate::sync::OnceLock;
-    use crate::sync::atomic::{AtomicBool, AtomicPtr, AtomicUsize, Ordering};
+    use crate::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
     use crate::sys::pal::unix::os;
+    use crate::sys::thread_local::{guard, local_pointer};
     use crate::{io, mem, ptr, thread};
 
     // We use a TLS variable to store the address of the guard page. While TLS
     // variables are not guaranteed to be signal-safe, this works out in practice
     // since we make sure to write to the variable before the signal stack is
     // installed, thereby ensuring that the variable is always allocated when
     // the signal handler is called.
-    thread_local! {
-        // FIXME: use `Range` once that implements `Copy`.
-        static GUARD: Cell<(usize, usize)> = const { Cell::new((0, 0)) };
+    local_pointer! {
+        static GUARD_START;
+        static GUARD_END;
+
+        static SIGALTSTACK;
     }
 
     // Signal handler for the SIGSEGV and SIGBUS handlers. We've got guard pages
@@ -93,7 +71,9 @@ mod imp {
         info: *mut libc::siginfo_t,
         _data: *mut libc::c_void,
     ) {
-        let (start, end) = GUARD.get();
+        let start = GUARD_START.get().addr();
+        let end = GUARD_END.get().addr();
+
         // SAFETY: this pointer is provided by the system and will always point to a valid `siginfo_t`.
         let addr = unsafe { (*info).si_addr().addr() };
 
@@ -119,51 +99,72 @@ mod imp {
     }
 
     static PAGE_SIZE: AtomicUsize = AtomicUsize::new(0);
-    static MAIN_ALTSTACK: AtomicPtr<libc::c_void> = AtomicPtr::new(ptr::null_mut());
     static NEED_ALTSTACK: AtomicBool = AtomicBool::new(false);
 
+    /// Set up stack overflow protection for the current thread
+    ///
     /// # Safety
-    /// Must be called only once
+    /// May only be called once per thread.
     #[forbid(unsafe_op_in_unsafe_fn)]
-    pub unsafe fn init() {
-        PAGE_SIZE.store(os::page_size(), Ordering::Relaxed);
-
-        // Always write to GUARD to ensure the TLS variable is allocated.
-        let guard = unsafe { install_main_guard().unwrap_or(0..0) };
-        GUARD.set((guard.start, guard.end));
-
-        // SAFETY: assuming all platforms define struct sigaction as "zero-initializable"
-        let mut action: sigaction = unsafe { mem::zeroed() };
-        for &signal in &[SIGSEGV, SIGBUS] {
-            // SAFETY: just fetches the current signal handler into action
-            unsafe { sigaction(signal, ptr::null_mut(), &mut action) };
-            // Configure our signal handler if one is not already set.
-            if action.sa_sigaction == SIG_DFL {
-                if !NEED_ALTSTACK.load(Ordering::Relaxed) {
-                    // haven't set up our sigaltstack yet
-                    NEED_ALTSTACK.store(true, Ordering::Release);
-                    let handler = unsafe { make_handler(true) };
-                    MAIN_ALTSTACK.store(handler.data, Ordering::Relaxed);
-                    mem::forget(handler);
+    pub unsafe fn protect(main_thread: bool) {
+        if main_thread {
+            PAGE_SIZE.store(os::page_size(), Ordering::Relaxed);
+        // Use acquire ordering to observe the page size store above,
+        // which is propagated by a release store to NEED_ALTSTACK.
+        } else if !NEED_ALTSTACK.load(Ordering::Acquire) {
+            return;
+        }
+
+        let guard = if main_thread {
+            unsafe { install_main_guard().unwrap_or(0..0) }
+        } else {
+            unsafe { current_guard().unwrap_or(0..0) }
+        };
+
+        // Always store the guard range to ensure the TLS variables are allocated.
+        GUARD_START.set(ptr::without_provenance_mut(guard.start));
+        GUARD_END.set(ptr::without_provenance_mut(guard.end));
+
+        if main_thread {
+            // SAFETY: assuming all platforms define struct sigaction as "zero-initializable"
+            let mut action: sigaction = unsafe { mem::zeroed() };
+            for &signal in &[SIGSEGV, SIGBUS] {
+                // SAFETY: just fetches the current signal handler into action
+                unsafe { sigaction(signal, ptr::null_mut(), &mut action) };
+                // Configure our signal handler if one is not already set.
+                if action.sa_sigaction == SIG_DFL {
+                    if !NEED_ALTSTACK.load(Ordering::Relaxed) {
+                        // Set up the signal stack and tell other threads to set
+                        // up their own. This uses a release store to propagate
+                        // the store to PAGE_SIZE above.
+                        NEED_ALTSTACK.store(true, Ordering::Release);
+                        unsafe { setup_sigaltstack() };
+                    }
+
+                    action.sa_flags = SA_SIGINFO | SA_ONSTACK;
+                    action.sa_sigaction = signal_handler as sighandler_t;
+                    // SAFETY: only overriding signals if the default is set
+                    unsafe { sigaction(signal, &action, ptr::null_mut()) };
                 }
-                action.sa_flags = SA_SIGINFO | SA_ONSTACK;
-                action.sa_sigaction = signal_handler as sighandler_t;
-                // SAFETY: only overriding signals if the default is set
-                unsafe { sigaction(signal, &action, ptr::null_mut()) };
             }
+        } else {
+            unsafe { setup_sigaltstack() };
         }
     }
 
     /// # Safety
-    /// Must be called only once
+    /// Mutates the alternate signal stack
     #[forbid(unsafe_op_in_unsafe_fn)]
-    pub unsafe fn cleanup() {
-        // FIXME: I probably cause more bugs than I'm worth!
-        // see https://github.com/rust-lang/rust/issues/111272
-        unsafe { drop_handler(MAIN_ALTSTACK.load(Ordering::Relaxed)) };
-    }
+    unsafe fn setup_sigaltstack() {
+        // SAFETY: assuming stack_t is zero-initializable
+        let mut stack = unsafe { mem::zeroed() };
+        // SAFETY: reads current stack_t into stack
+        unsafe { sigaltstack(ptr::null(), &mut stack) };
+        // Do not overwrite the stack if one is already set.
+        if stack.ss_flags & SS_DISABLE == 0 {
+            return;
+        }
 
-    unsafe fn get_stack() -> libc::stack_t {
         // OpenBSD requires this flag for stack mapping
         // otherwise the said mapping will fail as a no-op on most systems
         // and has a different meaning on FreeBSD
@@ -185,82 +186,60 @@ mod imp {
         let sigstack_size = sigstack_size();
         let page_size = PAGE_SIZE.load(Ordering::Relaxed);
 
-        let stackp = mmap64(
-            ptr::null_mut(),
-            sigstack_size + page_size,
-            PROT_READ | PROT_WRITE,
-            flags,
-            -1,
-            0,
-        );
-        if stackp == MAP_FAILED {
+        let allocation = unsafe {
+            mmap64(ptr::null_mut(), sigstack_size + page_size, PROT_READ | PROT_WRITE, flags, -1, 0)
+        };
+        if allocation == MAP_FAILED {
             panic!("failed to allocate an alternative stack: {}", io::Error::last_os_error());
         }
-        let guard_result = libc::mprotect(stackp, page_size, PROT_NONE);
+
+        let guard_result = unsafe { libc::mprotect(allocation, page_size, PROT_NONE) };
         if guard_result != 0 {
             panic!("failed to set up alternative stack guard page: {}", io::Error::last_os_error());
         }
-        let stackp = stackp.add(page_size);
 
-        libc::stack_t { ss_sp: stackp, ss_flags: 0, ss_size: sigstack_size }
+        let stack = libc::stack_t {
+            // Reserve a guard page at the bottom of the allocation.
+            ss_sp: unsafe { allocation.add(page_size) },
+            ss_flags: 0,
+            ss_size: sigstack_size,
+        };
+        // SAFETY: We warned our caller this would happen!
+        unsafe {
+            sigaltstack(&stack, ptr::null_mut());
+        }
+
+        // Ensure that `rt::thread_cleanup` gets called, which will in turn call
+        // cleanup, where this signal stack will be freed.
+        guard::enable();
+        SIGALTSTACK.set(allocation.cast());
     }
 
-    /// # Safety
-    /// Mutates the alternate signal stack
-    #[forbid(unsafe_op_in_unsafe_fn)]
-    pub unsafe fn make_handler(main_thread: bool) -> Handler {
-        if !NEED_ALTSTACK.load(Ordering::Acquire) {
-            return Handler::null();
+    pub fn cleanup() {
+        let allocation = SIGALTSTACK.get();
+        if allocation.is_null() {
+            return;
         }
 
-        if !main_thread {
-            // Always write to GUARD to ensure the TLS variable is allocated.
-            let guard = unsafe { current_guard() }.unwrap_or(0..0);
-            GUARD.set((guard.start, guard.end));
-        }
+        SIGALTSTACK.set(ptr::null_mut());
 
-        // SAFETY: assuming stack_t is zero-initializable
-        let mut stack = unsafe { mem::zeroed() };
-        // SAFETY: reads current stack_t into stack
-        unsafe { sigaltstack(ptr::null(), &mut stack) };
-        // Configure alternate signal stack, if one is not already set.
-        if stack.ss_flags & SS_DISABLE != 0 {
-            // SAFETY: We warned our caller this would happen!
-            unsafe {
-                stack = get_stack();
-                sigaltstack(&stack, ptr::null_mut());
-            }
-            Handler { data: stack.ss_sp as *mut libc::c_void }
-        } else {
-            Handler::null()
-        }
-    }
+        let sigstack_size = sigstack_size();
+        let page_size = PAGE_SIZE.load(Ordering::Relaxed);
 
-    /// # Safety
-    /// Must be called
-    /// - only with our handler or nullptr
-    /// - only when done with our altstack
-    /// This disables the alternate signal stack!
-    #[forbid(unsafe_op_in_unsafe_fn)]
-    pub unsafe fn drop_handler(data: *mut libc::c_void) {
-        if !data.is_null() {
-            let sigstack_size = sigstack_size();
-            let page_size = PAGE_SIZE.load(Ordering::Relaxed);
-            let disabling_stack = libc::stack_t {
-                ss_sp: ptr::null_mut(),
-                ss_flags: SS_DISABLE,
-                // Workaround for bug in macOS implementation of sigaltstack
-                // UNIX2003 which returns ENOMEM when disabling a stack while
-                // passing ss_size smaller than MINSIGSTKSZ. According to POSIX
-                // both ss_sp and ss_size should be ignored in this case.
-                ss_size: sigstack_size,
-            };
-            // SAFETY: we warned the caller this disables the alternate signal stack!
-            unsafe { sigaltstack(&disabling_stack, ptr::null_mut()) };
-            // SAFETY: We know from `get_stackp` that the alternate stack we installed is part of
-            // a mapping that started one page earlier, so walk back a page and unmap from there.
-            unsafe { munmap(data.sub(page_size), sigstack_size + page_size) };
-        }
+        let disabling_stack = libc::stack_t {
+            ss_sp: ptr::null_mut(),
+            ss_flags: SS_DISABLE,
+            // Workaround for bug in macOS implementation of sigaltstack
+            // UNIX2003 which returns ENOMEM when disabling a stack while
+            // passing ss_size smaller than MINSIGSTKSZ. According to POSIX
+            // both ss_sp and ss_size should be ignored in this case.
+            ss_size: sigstack_size,
+        };
+        unsafe { sigaltstack(&disabling_stack, ptr::null_mut()) };
+
+        // SAFETY: we created this mapping in `setup_sigaltstack` above with
+        // this exact size.
+        unsafe { munmap(allocation.cast(), sigstack_size + page_size) };
     }
 
     /// Modern kernels on modern hardware can have dynamic signal stack sizes.
@@ -577,13 +556,6 @@ mod imp {
     target_os = "illumos",
 )))]
 mod imp {
-    pub unsafe fn init() {}
-
-    pub unsafe fn cleanup() {}
-
-    pub unsafe fn make_handler(_main_thread: bool) -> super::Handler {
-        super::Handler::null()
-    }
-
-    pub unsafe fn drop_handler(_data: *mut libc::c_void) {}
+    pub unsafe fn protect(_main_thread: bool) {}
+    pub fn cleanup() {}
 }