extras: Fix UB in Peripheral

`Peripheral` assumed that interrupts can't be preempted,
when they can be preempted by higher priority interrupts.
So I put the interrupt handler inside a critical section,
and also added checks for whether the state had been dropped
before the critical section was entered.

I also added a `'static` bound to `PeripheralState`,
since `Pin` only guarantees that the memory it directly references
will not be invalidated.
It doesn't guarantee that memory its pointee references also won't be invalidated.

There were already some implementations of `PeripheralState`
that weren't `'static`, though,
so I added an unsafe `PeripheralStateUnchecked` trait
and forwarded the `unsafe` to the constructors of the implementors.

embassy-extras/src/peripheral.rs

@@ -1,15 +1,38 @@
 use core::cell::UnsafeCell;
 use core::marker::{PhantomData, PhantomPinned};
 use core::pin::Pin;
+use core::ptr;
 
 use embassy::interrupt::{Interrupt, InterruptExt};
 
-pub trait PeripheralState {
+/// # Safety
+/// When types implementing this trait are used with `Peripheral` or `PeripheralMutex`,
+/// their lifetime must not end without first calling `Drop` on the `Peripheral` or `PeripheralMutex`.
+pub unsafe trait PeripheralStateUnchecked {
     type Interrupt: Interrupt;
     fn on_interrupt(&mut self);
 }
 
-pub struct PeripheralMutex<S: PeripheralState> {
+// `PeripheralMutex` is safe because `Pin` guarantees that the memory it references will not be invalidated or reused
+// without calling `Drop`. However, it provides no guarantees about references contained within the state still being valid,
+// so this `'static` bound is necessary.
+pub trait PeripheralState: 'static {
+    type Interrupt: Interrupt;
+    fn on_interrupt(&mut self);
+}
+
+// SAFETY: `T` has to live for `'static` to implement `PeripheralState`, thus its lifetime cannot end.
+unsafe impl<T> PeripheralStateUnchecked for T
+where
+    T: PeripheralState,
+{
+    type Interrupt = T::Interrupt;
+    fn on_interrupt(&mut self) {
+        self.on_interrupt()
+    }
+}
+
+pub struct PeripheralMutex<S: PeripheralStateUnchecked> {
     state: UnsafeCell<S>,
 
     irq_setup_done: bool,
@@ -19,7 +42,7 @@ pub struct PeripheralMutex<S: PeripheralState> {
     _pinned: PhantomPinned,
 }
 
-impl<S: PeripheralState> PeripheralMutex<S> {
+impl<S: PeripheralStateUnchecked> PeripheralMutex<S> {
     pub fn new(state: S, irq: S::Interrupt) -> Self {
         Self {
             irq,
@@ -39,11 +62,17 @@ impl<S: PeripheralState> PeripheralMutex<S> {
 
         this.irq.disable();
         this.irq.set_handler(|p| {
-            // Safety: it's OK to get a &mut to the state, since
-            // - We're in the IRQ, no one else can't preempt us
-            // - 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();
+            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();
+            })
         });
         this.irq
             .set_handler_context((&mut this.state) as *mut _ as *mut ());
@@ -67,9 +96,12 @@ impl<S: PeripheralState> PeripheralMutex<S> {
     }
 }
 
-impl<S: PeripheralState> Drop for PeripheralMutex<S> {
+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());
     }
 }