From 70e1dc99672235ae34b6b3981bb0805f97c20179 Mon Sep 17 00:00:00 2001
From: Thom Chiovoloni <thom@shift.click>
Date: Sun, 4 Jun 2023 14:54:28 -0700
Subject: [PATCH] Avoid unwind across `extern "C"` in
 `thread_local::fast_local.rs`

---
 .../src/sys/common/thread_local/fast_local.rs | 22 +++++++++----------
 .../std/src/sys/common/thread_local/mod.rs    | 21 ++++++++++++++++++
 library/std/src/thread/mod.rs                 |  2 +-
 3 files changed, 33 insertions(+), 12 deletions(-)

diff --git a/library/std/src/sys/common/thread_local/fast_local.rs b/library/std/src/sys/common/thread_local/fast_local.rs
index 447044a798ba6..bc5da1a189677 100644
--- a/library/std/src/sys/common/thread_local/fast_local.rs
+++ b/library/std/src/sys/common/thread_local/fast_local.rs
@@ -33,20 +33,21 @@ pub macro thread_local_inner {
             // 1 == dtor registered, dtor not run
             // 2 == dtor registered and is running or has run
             #[thread_local]
-            static mut STATE: $crate::primitive::u8 = 0;
+            static STATE: $crate::cell::Cell<$crate::primitive::u8> = $crate::cell::Cell::new(0);
 
+            // Safety: Performs `drop_in_place(ptr as *mut $t)`, and requires
+            // all that comes with it.
             unsafe extern "C" fn destroy(ptr: *mut $crate::primitive::u8) {
-                let ptr = ptr as *mut $t;
-
-                unsafe {
-                    $crate::debug_assert_eq!(STATE, 1);
-                    STATE = 2;
-                    $crate::ptr::drop_in_place(ptr);
-                }
+                $crate::thread::local_impl::abort_on_dtor_unwind(|| {
+                    let old_state = STATE.replace(2);
+                    $crate::debug_assert_eq!(old_state, 1);
+                    // Safety: safety requirement is passed on to caller.
+                    unsafe { $crate::ptr::drop_in_place(ptr.cast::<$t>()); }
+                });
             }
 
             unsafe {
-                match STATE {
+                match STATE.get() {
                     // 0 == we haven't registered a destructor, so do
                     //   so now.
                     0 => {
@@ -54,7 +55,7 @@ pub macro thread_local_inner {
                             $crate::ptr::addr_of_mut!(VAL) as *mut $crate::primitive::u8,
                             destroy,
                         );
-                        STATE = 1;
+                        STATE.set(1);
                         $crate::option::Option::Some(&VAL)
                     }
                     // 1 == the destructor is registered and the value
@@ -148,7 +149,6 @@ impl<T> fmt::Debug for Key<T> {
         f.debug_struct("Key").finish_non_exhaustive()
     }
 }
-
 impl<T> Key<T> {
     pub const fn new() -> Key<T> {
         Key { inner: LazyKeyInner::new(), dtor_state: Cell::new(DtorState::Unregistered) }
diff --git a/library/std/src/sys/common/thread_local/mod.rs b/library/std/src/sys/common/thread_local/mod.rs
index 77f645883102c..975509bd412b0 100644
--- a/library/std/src/sys/common/thread_local/mod.rs
+++ b/library/std/src/sys/common/thread_local/mod.rs
@@ -101,3 +101,24 @@ mod lazy {
         }
     }
 }
+
+/// Run a callback in a scenario which must not unwind (such as a `extern "C"
+/// fn` declared in a user crate). If the callback unwinds anyway, then
+/// `rtabort` with a message about thread local panicking on drop.
+#[inline]
+pub fn abort_on_dtor_unwind(f: impl FnOnce()) {
+    // Using a guard like this is lower cost.
+    let guard = DtorUnwindGuard;
+    f();
+    core::mem::forget(guard);
+
+    struct DtorUnwindGuard;
+    impl Drop for DtorUnwindGuard {
+        #[inline]
+        fn drop(&mut self) {
+            // This is not terribly descriptive, but it doesn't need to be as we'll
+            // already have printed a panic message at this point.
+            rtabort!("thread local panicked on drop");
+        }
+    }
+}
diff --git a/library/std/src/thread/mod.rs b/library/std/src/thread/mod.rs
index f712c872708ac..d9973185bc45e 100644
--- a/library/std/src/thread/mod.rs
+++ b/library/std/src/thread/mod.rs
@@ -206,7 +206,7 @@ cfg_if::cfg_if! {
         #[doc(hidden)]
         #[unstable(feature = "thread_local_internals", issue = "none")]
         pub mod local_impl {
-            pub use crate::sys::common::thread_local::{thread_local_inner, Key};
+            pub use crate::sys::common::thread_local::{thread_local_inner, Key, abort_on_dtor_unwind};
         }
     }
 }