Auto merge of #160976 - maxdexh:fix-thread-local-global-alloc-reenter-2, r=<try>

Ensure TLS accesses don't call the global allocator through panic (part 2)


try-job: test-various

Author: bors

library/std/src/sys/thread_local/destructors/list.rs

@@ -11,6 +11,11 @@ pub unsafe fn register(t: *mut u8, dtor: unsafe extern "C" fn(*mut u8)) {
         rtabort!("the System allocator may not use TLS with destructors")
     };
     guard::enable();
+
+    // Avoid calling the alloc error hook
+    if dtors.capacity() == dtors.len() {
+        dtors.try_reserve(1).unwrap_or_else(|_| rtabort!("Failed to grow TLS destructor list"))
+    }
     dtors.push((t, dtor));
 }
 

library/std/src/sys/thread_local/key/tests.rs

@@ -1,3 +1,13 @@
+#![allow(
+    clippy::arithmetic_side_effects,
+    clippy::expect_used,
+    clippy::unwrap_used,
+    clippy::indexing_slicing,
+    clippy::panic,
+    clippy::unreachable,
+    clippy::unimplemented
+)]
+
 use super::{LazyKey, get, set};
 use crate::ptr;
 

library/std/src/sys/thread_local/key/unix.rs

@@ -31,13 +31,13 @@ pub fn create(dtor: Option<unsafe extern "C" fn(*mut u8)>) -> Key {
     key
 }
 
+#[cold]
+fn fail() -> ! {
+    rtabort!("Unexpected TLS failure")
+}
+
 #[inline]
 pub unsafe fn set(key: Key, value: *mut u8) {
-    #[cold]
-    fn fail() -> ! {
-        rtabort!("Failed to set value of thread local")
-    }
-
     let r = unsafe { libc::pthread_setspecific(key, value as *mut _) };
     // May happen on memory exhaustion
     if r != 0 {
@@ -54,5 +54,8 @@ pub unsafe fn get(key: Key) -> *mut u8 {
 #[inline]
 pub unsafe fn destroy(key: Key) {
     let r = unsafe { libc::pthread_key_delete(key) };
-    debug_assert_eq!(r, 0);
+    // only documented error is for invalid keys
+    if r != 0 {
+        fail()
+    }
 }

library/std/src/sys/thread_local/key/windows.rs

@@ -47,6 +47,11 @@ pub struct LazyKey {
     once: UnsafeCell<c::INIT_ONCE>,
 }
 
+#[cold]
+fn fail() -> ! {
+    rtabort!("Unexpected TLS failure")
+}
+
 impl LazyKey {
     #[inline]
     pub const fn new(dtor: Option<Dtor>) -> LazyKey {
@@ -66,9 +71,9 @@ impl LazyKey {
             guard::enable();
         }
 
-        match self.key.load(Acquire) {
-            0 => unsafe { self.init() },
-            key => key - 1,
+        match self.key.load(Acquire).checked_sub(1) {
+            None => unsafe { self.init() },
+            Some(dec) => dec,
         }
     }
 
@@ -79,11 +84,13 @@ impl LazyKey {
             let r = unsafe {
                 c::InitOnceBeginInitialize(self.once.get(), 0, &mut pending, ptr::null_mut())
             };
-            assert_eq!(r, c::TRUE);
+            if r != c::TRUE {
+                fail()
+            }
 
             if pending == c::FALSE {
                 // Some other thread initialized the key, load it.
-                self.key.load(Relaxed) - 1
+                self.key.load(Relaxed).wrapping_sub(1)
             } else {
                 let key = unsafe { c::TlsAlloc() };
                 if key == c::TLS_OUT_OF_INDEXES {
@@ -104,10 +111,12 @@ impl LazyKey {
                 // and if that sees this write then it will entirely bypass the `InitOnce`. We thus
                 // need to establish synchronization through `key`. In particular that acquire load
                 // must happen-after the register_dtor above, to ensure the dtor actually runs!
-                self.key.store(key + 1, Release);
+                self.key.store(key.wrapping_add(1), Release);
 
                 let r = unsafe { c::InitOnceComplete(self.once.get(), 0, ptr::null_mut()) };
-                debug_assert_eq!(r, c::TRUE);
+                if r != c::TRUE {
+                    fail()
+                }
 
                 key
             }
@@ -119,14 +128,16 @@ impl LazyKey {
                 rtabort!("out of TLS indexes");
             }
 
-            match self.key.compare_exchange(0, key + 1, AcqRel, Acquire) {
+            match self.key.compare_exchange(0, key.wrapping_add(1), AcqRel, Acquire) {
                 Ok(_) => key,
                 Err(new) => unsafe {
                     // Some other thread completed initialization first, so destroy
                     // our key and use theirs.
                     let r = c::TlsFree(key);
-                    debug_assert_eq!(r, c::TRUE);
-                    new - 1
+                    if r != c::TRUE {
+                        fail()
+                    }
+                    new.wrapping_sub(1)
                 },
             }
         }
@@ -138,10 +149,6 @@ unsafe impl Sync for LazyKey {}
 
 #[inline]
 pub unsafe fn set(key: Key, val: *mut u8) {
-    #[cold]
-    fn fail() -> ! {
-        rtabort!("Failed to set value of thread local")
-    }
     let r = unsafe { c::TlsSetValue(key, val.cast()) };
     // According to MS documentation, `TlsSetValue` returns zero "if it fails"
     if r != c::TRUE {
@@ -181,22 +188,21 @@ pub unsafe fn run_dtors() {
         let mut cur = DTORS.load(Acquire);
         while !cur.is_null() {
             let pre_key = unsafe { (*cur).key.load(Acquire) };
-            let dtor = unsafe { (*cur).dtor.unwrap() };
+            let dtor = unsafe { rtunwrap!(Some, (*cur).dtor) };
             cur = unsafe { (*cur).next.load(Relaxed) };
 
             // In LazyKey::init, we register the dtor before setting `key`.
             // So if one thread's `run_dtors` races with another thread executing `init` on the same
             // `LazyKey`, we can encounter a key of 0 here. That means this key was never
             // initialized in this thread so we can safely skip it.
-            if pre_key == 0 {
+            let Some(key) = pre_key.checked_sub(1) else {
                 continue;
-            }
+            };
+
             // If this is non-zero, then via the `Acquire` load above we synchronized with
             // everything relevant for this key. (It's not clear that this is needed, since the
             // release-acquire pair on DTORS also establishes synchronization, but better safe than
             // sorry.)
-            let key = pre_key - 1;
-
             let ptr = unsafe { c::TlsGetValue(key) };
             if !ptr.is_null() {
                 unsafe {

library/std/src/sys/thread_local/key/xous.rs

@@ -104,11 +104,11 @@ fn tls_table_slow() -> &'static mut [*mut u8] {
             TLS_MEMORY_SIZE / size_of::<*mut u8>(),
             MemoryFlags::R | MemoryFlags::W,
         )
-        .expect("Unable to allocate memory for thread local storage")
+        .unwrap_or_else(|_| rtabort!("Unable to allocate memory for thread local storage"))
     };
 
     for val in tp.iter() {
-        assert!((*val).is_null());
+        rtassert!((*val).is_null());
     }
 
     unsafe {
@@ -136,13 +136,14 @@ pub fn create(dtor: Option<Dtor>) -> Key {
 pub unsafe fn set(key: Key, value: *mut u8) {
     rtassert!((key < 1022) && (key >= 1));
     let table = tls_table();
-    table[key] = value;
+    *rtunwrap!(Some, table.get_mut(key)) = value;
 }
 
 #[inline]
 pub unsafe fn get(key: Key) -> *mut u8 {
     rtassert!((key < 1022) && (key >= 1));
-    tls_table()[key]
+    let table = tls_table();
+    *rtunwrap!(Some, table.get(key))
 }
 
 #[inline]
@@ -186,10 +187,10 @@ pub unsafe fn destroy_tls() {
     unsafe { run_dtors() };
 
     // Finally, free the TLS array
-    unsafe {
+    let result = unsafe {
         unmap_memory(core::slice::from_raw_parts_mut(tp, TLS_MEMORY_SIZE / size_of::<usize>()))
-            .unwrap()
     };
+    rtunwrap!(Ok, result);
 }
 
 // This is marked inline(never) to prevent dealloc calls from being reordered

library/std/src/sys/thread_local/mod.rs

@@ -22,6 +22,16 @@
     reason = "internal details of the thread_local macro",
     issue = "none"
 )]
+#![deny(
+    clippy::arithmetic_side_effects,
+    clippy::expect_used,
+    clippy::unwrap_used,
+    clippy::indexing_slicing,
+    clippy::panic,
+    clippy::unreachable,
+    clippy::unimplemented,
+    reason = "TLS accesses must not call the global allocator, including via panic (#160930)"
+)]
 
 cfg_select! {
     any(

library/std/src/sys/thread_local/no_threads.rs

@@ -95,7 +95,7 @@ impl<T> LazyStorage<T> {
         let value = i.and_then(Option::take).unwrap_or_else(f);
 
         // Destroy the old value if it is initialized
-        // FIXME(#110897): maybe panic on recursive initialization.
+        // FIXME(#110897): maybe abort on recursive initialization.
         if self.state.get() == State::Alive {
             self.state.set(State::Destroying);
             // Safety: we check for no initialization during drop below
@@ -107,7 +107,7 @@ impl<T> LazyStorage<T> {
 
         // Guard against initialization during drop
         if self.state.get() == State::Destroying {
-            panic!("Attempted to initialize thread-local while it is being dropped");
+            rtabort!("Attempted to initialize thread-local while it is being dropped");
         }
 
         unsafe {

library/std/src/sys/thread_local/os.rs

@@ -1,6 +1,6 @@
 use super::key::{Key, LazyKey, get, set};
 use super::{abort_on_dtor_unwind, guard};
-use crate::alloc::{self, GlobalAlloc, Layout, System};
+use crate::alloc::{GlobalAlloc, Layout, System};
 use crate::cell::Cell;
 use crate::marker::PhantomData;
 use crate::mem::ManuallyDrop;
@@ -103,13 +103,14 @@ struct AlignedSystemBox<T: 'static, const ALIGN: usize> {
 impl<T: 'static, const ALIGN: usize> AlignedSystemBox<T, ALIGN> {
     #[inline]
     fn new(v: Value<T>) -> Self {
-        let layout = Layout::new::<Value<T>>().align_to(ALIGN).unwrap();
+        let layout = rtunwrap!(Ok, Layout::new::<Value<T>>().align_to(ALIGN));
 
         // We use the System allocator here to avoid interfering with a potential
         // Global allocator using thread-local storage.
         let ptr: *mut Value<T> = (unsafe { System.alloc(layout) }).cast();
         let Some(ptr) = NonNull::new(ptr) else {
-            alloc::handle_alloc_error(layout);
+            // Do not call the alloc error hook here. It may allocate!
+            rtabort!("Allocation failure");
         };
         unsafe { ptr.write(v) };
         Self { ptr }
@@ -139,7 +140,7 @@ impl<T: 'static, const ALIGN: usize> Deref for AlignedSystemBox<T, ALIGN> {
 impl<T: 'static, const ALIGN: usize> Drop for AlignedSystemBox<T, ALIGN> {
     #[inline]
     fn drop(&mut self) {
-        let layout = Layout::new::<Value<T>>().align_to(ALIGN).unwrap();
+        let layout = rtunwrap!(Ok, Layout::new::<Value<T>>().align_to(ALIGN));
 
         unsafe {
             let unwind_result = catch_unwind(AssertUnwindSafe(|| self.ptr.drop_in_place()));