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.

1 files changed, 33 insertions, 12 deletions

diff --git a/embassy-extras/src/peripheral_shared.rs b/embassy-extras/src/peripheral_shared.rs
index c62113396..a9fca8ca6 100644
--- a/embassy-extras/src/peripheral_shared.rs
+++ b/embassy-extras/src/peripheral_shared.rs
@@ -1,16 +1,27 @@
-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(&self);
 }
 
-pub struct Peripheral<S: PeripheralState> {
-    state: UnsafeCell<S>,
+// `Peripheral` 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(&self);
+}
+
+pub struct Peripheral<S: PeripheralStateUnchecked> {
+    state: S,
 
     irq_setup_done: bool,
     irq: S::Interrupt,
@@ -19,13 +30,13 @@ pub struct Peripheral<S: PeripheralState> {
     _pinned: PhantomPinned,
 }
 
-impl<S: PeripheralState> Peripheral<S> {
+impl<S: PeripheralStateUnchecked> Peripheral<S> {
     pub fn new(irq: S::Interrupt, state: S) -> Self {
         Self {
             irq,
             irq_setup_done: false,
 
-            state: UnsafeCell::new(state),
+            state,
             _not_send: PhantomData,
             _pinned: PhantomPinned,
         }
@@ -39,8 +50,16 @@ impl<S: PeripheralState> Peripheral<S> {
 
         this.irq.disable();
         this.irq.set_handler(|p| {
-            let state = unsafe { &*(p as *const S) };
-            state.on_interrupt();
+            // 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();
+            });
         });
         this.irq
             .set_handler_context((&this.state) as *const _ as *mut ());
@@ -49,15 +68,17 @@ impl<S: PeripheralState> Peripheral<S> {
         this.irq_setup_done = true;
     }
 
-    pub fn state(self: Pin<&mut Self>) -> &S {
-        let this = unsafe { self.get_unchecked_mut() };
-        unsafe { &*this.state.get() }
+    pub fn state<'a>(self: Pin<&'a mut Self>) -> &'a S {
+        &self.into_ref().get_ref().state
     }
 }
 
-impl<S: PeripheralState> Drop for Peripheral<S> {
+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());
     }
 }