fix: Prevent double-enter with same &mut ExtCtxt (which would be UB)

Author: futile

compiler/rustc_expand/src/derive_macro_expansion.rs

@@ -12,7 +12,7 @@ use rustc_span::LocalExpnId;
 use crate::base::ExtCtxt;
 use crate::errors;
 
-pub(super) fn expand<'tcx>(
+pub(super) fn provide_derive_macro_expansion<'tcx>(
     tcx: TyCtxt<'tcx>,
     key: (LocalExpnId, &'tcx TokenStream),
 ) -> Result<&'tcx TokenStream, ()> {
@@ -51,7 +51,6 @@ pub(super) fn expand<'tcx>(
 type CLIENT = pm::bridge::client::Client<pm::TokenStream, pm::TokenStream>;
 
 // based on rust/compiler/rustc_middle/src/ty/context/tls.rs
-// #[cfg(not(parallel_compiler))]
 thread_local! {
     /// A thread local variable that stores a pointer to the current `CONTEXT`.
     static TLV: Cell<(*mut (), Option<CLIENT>)> = const { Cell::new((ptr::null_mut(), None)) };
@@ -67,21 +66,29 @@ unsafe fn downcast<'a>(context: *mut ()) -> &'a mut ExtCtxt<'a> {
     unsafe { &mut *(context as *mut ExtCtxt<'a>) }
 }
 
-/// Sets `context` as the new current `CONTEXT` for the duration of the function `f`.
 #[inline]
-pub fn enter_context<'a, F, R>(context: (&mut ExtCtxt<'a>, CLIENT), f: F) -> R
+fn enter_context_erased<F, R>(erased: (*mut (), Option<CLIENT>), f: F) -> R
 where
     F: FnOnce() -> R,
 {
-    let (ectx, client) = context;
-    let erased = (erase(ectx), Some(client));
     TLV.with(|tlv| {
         let old = tlv.replace(erased);
         let _reset = rustc_data_structures::defer(move || tlv.set(old));
         f()
     })
 }
 
+/// Sets `context` as the new current `CONTEXT` for the duration of the function `f`.
+#[inline]
+pub fn enter_context<'a, F, R>(context: (&mut ExtCtxt<'a>, CLIENT), f: F) -> R
+where
+    F: FnOnce() -> R,
+{
+    let (ectx, client) = context;
+    let erased = (erase(ectx), Some(client));
+    enter_context_erased(erased, f)
+}
+
 /// Allows access to the current `CONTEXT` in a closure if one is available.
 #[inline]
 #[track_caller]
@@ -98,7 +105,12 @@ where
         // TODO: we should not be able to?
         // sync::assert_dyn_sync::<CONTEXT<'_>>();
 
-        unsafe { f(Some(&mut (downcast(ectx), client_opt.unwrap()))) }
+        // prevent double entering, as that would allow creating two `&mut ExtCtxt`s
+        // TODO: probably use a RefCell instead (which checks this properly)?
+        enter_context_erased((ptr::null_mut(), None), || unsafe {
+            let ectx = downcast(ectx);
+            f(Some(&mut (ectx, client_opt.unwrap())))
+        })
     }
 }
 

compiler/rustc_expand/src/lib.rs

@@ -36,7 +36,7 @@ pub mod module;
 pub mod proc_macro;
 
 pub fn provide(providers: &mut rustc_middle::util::Providers) {
-    providers.derive_macro_expansion = derive_macro_expansion::expand;
+    providers.derive_macro_expansion = derive_macro_expansion::provide_derive_macro_expansion;
 }
 
 rustc_fluent_macro::fluent_messages! { "../messages.ftl" }