Remove critical sections from `PeripheralMutex` interrupt handler by checking the interrupt's priority on startup.

Since `PeripheralMutex` is the only way to safely maintain state across interrupts, and it no longer allows setting the interrupt's priority, the priority changing isn't a concern.

This also prevents other causes of UB due to the interrupt being exposed during `with`, and allowing enabling the interrupt and setting its context to a bogus pointer.

5 files changed, 187 insertions, 48 deletions

diff --git a/embassy-extras/src/peripheral.rs b/embassy-extras/src/peripheral.rs
index 2358ab96f..2402ba9eb 100644
--- a/embassy-extras/src/peripheral.rs
+++ b/embassy-extras/src/peripheral.rs
@@ -1,8 +1,9 @@
 use core::cell::UnsafeCell;
 use core::marker::{PhantomData, PhantomPinned};
 use core::pin::Pin;
-use core::ptr;
 
+use cortex_m::peripheral::scb::{Exception, SystemHandler, VectActive};
+use cortex_m::peripheral::{NVIC, SCB};
 use embassy::interrupt::{Interrupt, InterruptExt};
 
 /// A version of `PeripheralState` without the `'static` bound,
@@ -50,8 +51,79 @@ pub struct PeripheralMutex<S: PeripheralStateUnchecked> {
     _pinned: PhantomPinned,
 }
 
+fn exception_to_system_handler(exception: Exception) -> Option<SystemHandler> {
+    match exception {
+        Exception::NonMaskableInt | Exception::HardFault => None,
+        #[cfg(not(armv6m))]
+        Exception::MemoryManagement => Some(SystemHandler::MemoryManagement),
+        #[cfg(not(armv6m))]
+        Exception::BusFault => Some(SystemHandler::BusFault),
+        #[cfg(not(armv6m))]
+        Exception::UsageFault => Some(SystemHandler::UsageFault),
+        #[cfg(any(armv8m, target_arch = "x86_64"))]
+        Exception::SecureFault => Some(SystemHandler::SecureFault),
+        Exception::SVCall => Some(SystemHandler::SVCall),
+        #[cfg(not(armv6m))]
+        Exception::DebugMonitor => Some(SystemHandler::DebugMonitor),
+        Exception::PendSV => Some(SystemHandler::PendSV),
+        Exception::SysTick => Some(SystemHandler::SysTick),
+    }
+}
+
+/// Whether `irq` can be preempted by the current interrupt.
+pub(crate) fn can_be_preempted(irq: &impl Interrupt) -> bool {
+    match SCB::vect_active() {
+        // Thread mode can't preempt each other
+        VectActive::ThreadMode => false,
+        VectActive::Exception(exception) => {
+            // `SystemHandler` is a subset of `Exception` for those with configurable priority.
+            // There's no built in way to convert between them, so `exception_to_system_handler` was necessary.
+            if let Some(system_handler) = exception_to_system_handler(exception) {
+                let current_prio = SCB::get_priority(system_handler);
+                let irq_prio = irq.get_priority().into();
+                if current_prio < irq_prio {
+                    true
+                } else if current_prio == irq_prio {
+                    // When multiple interrupts have the same priority number,
+                    // the pending interrupt with the lowest interrupt number takes precedence.
+                    (exception.irqn() as i16) < irq.number() as i16
+                } else {
+                    false
+                }
+            } else {
+                // There's no safe way I know of to maintain `!Send` state across invocations of HardFault or NMI, so that should be fine.
+                false
+            }
+        }
+        VectActive::Interrupt { irqn } => {
+            #[derive(Clone, Copy)]
+            struct NrWrap(u16);
+            unsafe impl cortex_m::interrupt::InterruptNumber for NrWrap {
+                fn number(self) -> u16 {
+                    self.0
+                }
+            }
+            let current_prio = NVIC::get_priority(NrWrap(irqn.into()));
+            let irq_prio = irq.get_priority().into();
+            if current_prio < irq_prio {
+                true
+            } else if current_prio == irq_prio {
+                // When multiple interrupts have the same priority number,
+                // the pending interrupt with the lowest interrupt number takes precedence.
+                (irqn as u16) < irq.number()
+            } else {
+                false
+            }
+        }
+    }
+}
+
 impl<S: PeripheralStateUnchecked> PeripheralMutex<S> {
     pub fn new(state: S, irq: S::Interrupt) -> Self {
+        if can_be_preempted(&irq) {
+            panic!("`PeripheralMutex` cannot be created in an interrupt with higher priority than the interrupt it wraps");
+        }
+
         Self {
             irq,
             irq_setup_done: false,
@@ -70,17 +142,13 @@ impl<S: PeripheralStateUnchecked> PeripheralMutex<S> {
 
         this.irq.disable();
         this.irq.set_handler(|p| {
-            critical_section::with(|_| {
-                if p.is_null() {
-                    // The state was dropped, so we can't operate on it.
-                    return;
-                }
-                // Safety: it's OK to get a &mut to the state, since
-                // - We're in a critical section, no one can preempt us (and call with())
-                // - We can't have preempted a with() call because the irq is disabled during it.
-                let state = unsafe { &mut *(p as *mut S) };
-                state.on_interrupt();
-            })
+            // Safety: it's OK to get a &mut to the state, since
+            // - We checked that the thread owning the `PeripheralMutex` can't preempt us in `new`.
+            //   Interrupts' priorities can only be changed with raw embassy `Interrupts`,
+            //   which can't safely store a `PeripheralMutex` across invocations.
+            // - We can't have preempted a with() call because the irq is disabled during it.
+            let state = unsafe { &mut *(p as *mut S) };
+            state.on_interrupt();
         });
         this.irq
             .set_handler_context((&mut this.state) as *mut _ as *mut ());
@@ -89,27 +157,63 @@ impl<S: PeripheralStateUnchecked> PeripheralMutex<S> {
         this.irq_setup_done = true;
     }
 
-    pub fn with<R>(self: Pin<&mut Self>, f: impl FnOnce(&mut S, &mut S::Interrupt) -> R) -> R {
+    pub fn with<R>(self: Pin<&mut Self>, f: impl FnOnce(&mut S) -> R) -> R {
         let this = unsafe { self.get_unchecked_mut() };
 
+        let was_enabled = this.irq.is_enabled();
         this.irq.disable();
 
         // Safety: it's OK to get a &mut to the state, since the irq is disabled.
         let state = unsafe { &mut *this.state.get() };
-        let r = f(state, &mut this.irq);
+        let r = f(state);
 
-        this.irq.enable();
+        if was_enabled {
+            this.irq.enable();
+        }
 
         r
     }
+
+    /// Enables the wrapped interrupt.
+    pub fn enable(&self) {
+        // This is fine to do before initialization, because we haven't set the handler yet.
+        self.irq.enable()
+    }
+
+    /// Disables the wrapped interrupt.
+    pub fn disable(&self) {
+        self.irq.disable()
+    }
+
+    /// Returns whether the wrapped interrupt is enabled.
+    pub fn is_enabled(&self) -> bool {
+        self.irq.is_enabled()
+    }
+
+    /// Returns whether the wrapped interrupt is currently in a pending state.
+    pub fn is_pending(&self) -> bool {
+        self.irq.is_pending()
+    }
+
+    /// Forces the wrapped interrupt into a pending state.
+    pub fn pend(&self) {
+        self.irq.pend()
+    }
+
+    /// Forces the wrapped interrupt out of a pending state.
+    pub fn unpend(&self) {
+        self.irq.unpend()
+    }
+
+    /// Gets the priority of the wrapped interrupt.
+    pub fn priority(&self) -> <S::Interrupt as Interrupt>::Priority {
+        self.irq.get_priority()
+    }
 }
 
 impl<S: PeripheralStateUnchecked> Drop for PeripheralMutex<S> {
     fn drop(&mut self) {
         self.irq.disable();
         self.irq.remove_handler();
-        // Set the context to null so that the interrupt will know we're dropped
-        // if we pre-empted it before it entered a critical section.
-        self.irq.set_handler_context(ptr::null_mut());
     }
 }
diff --git a/embassy-extras/src/peripheral_shared.rs b/embassy-extras/src/peripheral_shared.rs
index 5961441e0..ae9ae6935 100644
--- a/embassy-extras/src/peripheral_shared.rs
+++ b/embassy-extras/src/peripheral_shared.rs
@@ -1,9 +1,10 @@
 use core::marker::{PhantomData, PhantomPinned};
 use core::pin::Pin;
-use core::ptr;
 
 use embassy::interrupt::{Interrupt, InterruptExt};
 
+use crate::peripheral::can_be_preempted;
+
 /// A version of `PeripheralState` without the `'static` bound,
 /// for cases where the compiler can't statically make sure
 /// that `on_interrupt` doesn't reference anything which might be invalidated.
@@ -39,6 +40,10 @@ pub struct Peripheral<S: PeripheralStateUnchecked> {
 
 impl<S: PeripheralStateUnchecked> Peripheral<S> {
     pub fn new(irq: S::Interrupt, state: S) -> Self {
+        if can_be_preempted(&irq) {
+            panic!("`Peripheral` cannot be created in an interrupt with higher priority than the interrupt it wraps");
+        }
+
         Self {
             irq,
             irq_setup_done: false,
@@ -57,16 +62,11 @@ impl<S: PeripheralStateUnchecked> Peripheral<S> {
 
         this.irq.disable();
         this.irq.set_handler(|p| {
-            // We need to be in a critical section so that no one can preempt us
-            // and drop the state after we check whether `p.is_null()`.
-            critical_section::with(|_| {
-                if p.is_null() {
-                    // The state was dropped, so we can't operate on it.
-                    return;
-                }
-                let state = unsafe { &*(p as *const S) };
-                state.on_interrupt();
-            });
+            // The state can't have been dropped, otherwise the interrupt would have been disabled.
+            // We checked in `new` that the thread owning the `Peripheral` can't preempt the interrupt,
+            // so someone can't have preempted us before this point and dropped the `Peripheral`.
+            let state = unsafe { &*(p as *const S) };
+            state.on_interrupt();
         });
         this.irq
             .set_handler_context((&this.state) as *const _ as *mut ());
@@ -78,14 +78,47 @@ impl<S: PeripheralStateUnchecked> Peripheral<S> {
     pub fn state(self: Pin<&mut Self>) -> &S {
         &self.into_ref().get_ref().state
     }
+
+    /// Enables the wrapped interrupt.
+    pub fn enable(&self) {
+        // This is fine to do before initialization, because we haven't set the handler yet.
+        self.irq.enable()
+    }
+
+    /// Disables the wrapped interrupt.
+    pub fn disable(&self) {
+        self.irq.disable()
+    }
+
+    /// Returns whether the wrapped interrupt is enabled.
+    pub fn is_enabled(&self) -> bool {
+        self.irq.is_enabled()
+    }
+
+    /// Returns whether the wrapped interrupt is currently in a pending state.
+    pub fn is_pending(&self) -> bool {
+        self.irq.is_pending()
+    }
+
+    /// Forces the wrapped interrupt into a pending state.
+    pub fn pend(&self) {
+        self.irq.pend()
+    }
+
+    /// Forces the wrapped interrupt out of a pending state.
+    pub fn unpend(&self) {
+        self.irq.unpend()
+    }
+
+    /// Gets the priority of the wrapped interrupt.
+    pub fn priority(&self) -> <S::Interrupt as Interrupt>::Priority {
+        self.irq.get_priority()
+    }
 }
 
 impl<S: PeripheralStateUnchecked> Drop for Peripheral<S> {
     fn drop(&mut self) {
         self.irq.disable();
         self.irq.remove_handler();
-        // Set the context to null so that the interrupt will know we're dropped
-        // if we pre-empted it before it entered a critical section.
-        self.irq.set_handler_context(ptr::null_mut());
     }
 }
diff --git a/embassy-extras/src/usb/usb_serial.rs b/embassy-extras/src/usb/usb_serial.rs
index 9cbfb2da4..a229b2000 100644
--- a/embassy-extras/src/usb/usb_serial.rs
+++ b/embassy-extras/src/usb/usb_serial.rs
@@ -55,7 +55,7 @@ where
         let this = self.get_mut();
         let mut mutex = this.inner.borrow_mut();
         let mutex = unsafe { Pin::new_unchecked(&mut *mutex) };
-        mutex.with(|state, _irq| {
+        mutex.with(|state| {
             let serial = state.classes.get_serial();
             let serial = Pin::new(serial);
 
@@ -77,7 +77,7 @@ where
         let this = self.get_mut();
         let mut mutex = this.inner.borrow_mut();
         let mutex = unsafe { Pin::new_unchecked(&mut *mutex) };
-        mutex.with(|state, _irq| {
+        mutex.with(|state| {
             let serial = state.classes.get_serial();
             let serial = Pin::new(serial);
 
@@ -101,7 +101,7 @@ where
         let this = self.get_mut();
         let mut mutex = this.inner.borrow_mut();
         let mutex = unsafe { Pin::new_unchecked(&mut *mutex) };
-        mutex.with(|state, _irq| {
+        mutex.with(|state| {
             let serial = state.classes.get_serial();
             let serial = Pin::new(serial);
 
diff --git a/embassy-nrf/src/buffered_uarte.rs b/embassy-nrf/src/buffered_uarte.rs
index 1fa98a6b2..2cce122bc 100644
--- a/embassy-nrf/src/buffered_uarte.rs
+++ b/embassy-nrf/src/buffered_uarte.rs
@@ -176,7 +176,7 @@ impl<'d, U: UarteInstance, T: TimerInstance> BufferedUarte<'d, U, T> {
     pub fn set_baudrate(self: Pin<&mut Self>, baudrate: Baudrate) {
         let mut inner = self.inner();
         inner.as_mut().register_interrupt();
-        inner.with(|state, _irq| {
+        inner.with(|state| {
             let r = U::regs();
 
             let timeout = 0x8000_0000 / (baudrate as u32 / 40);
@@ -196,7 +196,7 @@ impl<'d, U: UarteInstance, T: TimerInstance> AsyncBufRead for BufferedUarte<'d,
     fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<&[u8]>> {
         let mut inner = self.inner();
         inner.as_mut().register_interrupt();
-        inner.with(|state, _irq| {
+        inner.with(|state| {
             // Conservative compiler fence to prevent optimizations that do not
             // take in to account actions by DMA. The fence has been placed here,
             // before any DMA action has started
@@ -221,11 +221,11 @@ impl<'d, U: UarteInstance, T: TimerInstance> AsyncBufRead for BufferedUarte<'d,
     fn consume(self: Pin<&mut Self>, amt: usize) {
         let mut inner = self.inner();
         inner.as_mut().register_interrupt();
-        inner.with(|state, irq| {
+        inner.as_mut().with(|state| {
             trace!("consume {:?}", amt);
             state.rx.pop(amt);
-            irq.pend();
-        })
+        });
+        inner.pend();
     }
 }
 
@@ -233,7 +233,7 @@ impl<'d, U: UarteInstance, T: TimerInstance> AsyncWrite for BufferedUarte<'d, U,
     fn poll_write(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) -> Poll<Result<usize>> {
         let mut inner = self.inner();
         inner.as_mut().register_interrupt();
-        inner.with(|state, irq| {
+        let poll = inner.as_mut().with(|state| {
             trace!("poll_write: {:?}", buf.len());
 
             let tx_buf = state.tx.push_buf();
@@ -254,10 +254,12 @@ impl<'d, U: UarteInstance, T: TimerInstance> AsyncWrite for BufferedUarte<'d, U,
             // before any DMA action has started
             compiler_fence(Ordering::SeqCst);
 
-            irq.pend();
-
             Poll::Ready(Ok(n))
-        })
+        });
+
+        inner.pend();
+
+        poll
     }
 }
 
diff --git a/embassy-stm32/src/eth/v2/mod.rs b/embassy-stm32/src/eth/v2/mod.rs
index 567088e8d..3ac09f944 100644
--- a/embassy-stm32/src/eth/v2/mod.rs
+++ b/embassy-stm32/src/eth/v2/mod.rs
@@ -161,7 +161,7 @@ impl<'d, P: PHY, const TX: usize, const RX: usize> Ethernet<'d, P, TX, RX> {
         let mut mutex = unsafe { Pin::new_unchecked(&mut this.state) };
         mutex.as_mut().register_interrupt();
 
-        mutex.with(|s, _| {
+        mutex.with(|s| {
             s.desc_ring.init();
 
             fence(Ordering::SeqCst);
@@ -237,7 +237,7 @@ impl<'d, P: PHY, const TX: usize, const RX: usize> Device for Pin<&mut Ethernet<
         let this = unsafe { self.as_mut().get_unchecked_mut() };
         let mutex = unsafe { Pin::new_unchecked(&mut this.state) };
 
-        mutex.with(|s, _| s.desc_ring.tx.available())
+        mutex.with(|s| s.desc_ring.tx.available())
     }
 
     fn transmit(&mut self, pkt: PacketBuf) {
@@ -245,7 +245,7 @@ impl<'d, P: PHY, const TX: usize, const RX: usize> Device for Pin<&mut Ethernet<
         let this = unsafe { self.as_mut().get_unchecked_mut() };
         let mutex = unsafe { Pin::new_unchecked(&mut this.state) };
 
-        mutex.with(|s, _| unwrap!(s.desc_ring.tx.transmit(pkt)));
+        mutex.with(|s| unwrap!(s.desc_ring.tx.transmit(pkt)));
     }
 
     fn receive(&mut self) -> Option<PacketBuf> {
@@ -253,7 +253,7 @@ impl<'d, P: PHY, const TX: usize, const RX: usize> Device for Pin<&mut Ethernet<
         let this = unsafe { self.as_mut().get_unchecked_mut() };
         let mutex = unsafe { Pin::new_unchecked(&mut this.state) };
 
-        mutex.with(|s, _| s.desc_ring.rx.pop_packet())
+        mutex.with(|s| s.desc_ring.rx.pop_packet())
     }
 
     fn register_waker(&mut self, waker: &Waker) {