2 files changed, 113 insertions, 76 deletions

diff --git a/embassy-usb/src/lib.rs b/embassy-usb/src/lib.rs
index ccea8bc7a..baeca099f 100644
--- a/embassy-usb/src/lib.rs
+++ b/embassy-usb/src/lib.rs
@@ -100,7 +100,6 @@ pub struct UsbDevice<'d, D: Driver<'d>> {
 
     device_state: UsbDeviceState,
     suspended: bool,
-    in_control_handler: bool,
     remote_wakeup_enabled: bool,
     self_powered: bool,
     pending_address: u8,
@@ -138,7 +137,6 @@ impl<'d, D: Driver<'d>> UsbDevice<'d, D> {
             control_buf,
             device_state: UsbDeviceState::Disabled,
             suspended: false,
-            in_control_handler: false,
             remote_wakeup_enabled: false,
             self_powered: false,
             pending_address: 0,
@@ -146,7 +144,26 @@ impl<'d, D: Driver<'d>> UsbDevice<'d, D> {
         }
     }
 
+    /// Runs the `UsbDevice` forever.
+    ///
+    /// This future may leave the bus in an invalid state if it is dropped.
+    /// After dropping the future, [`UsbDevice::disable()`] should be called
+    /// before calling any other `UsbDevice` methods to fully reset the
+    /// peripheral.
     pub async fn run(&mut self) -> ! {
+        loop {
+            self.run_until_suspend().await;
+            self.wait_resume().await;
+        }
+    }
+
+    /// Runs the `UsbDevice` until the bus is suspended.
+    ///
+    /// This future may leave the bus in an invalid state if it is dropped.
+    /// After dropping the future, [`UsbDevice::disable()`] should be called
+    /// before calling any other `UsbDevice` methods to fully reset the
+    /// peripheral.
+    pub async fn run_until_suspend(&mut self) -> () {
         if self.device_state == UsbDeviceState::Disabled {
             self.bus.enable().await;
             self.device_state = UsbDeviceState::Default;
@@ -154,66 +171,33 @@ impl<'d, D: Driver<'d>> UsbDevice<'d, D> {
             if let Some(h) = &self.handler {
                 h.enabled(true);
             }
-        } else if self.in_control_handler {
-            warn!("usb: control request interrupted");
-            self.control.reject();
-            self.in_control_handler = false;
         }
 
         loop {
             let control_fut = self.control.setup();
             let bus_fut = self.bus.poll();
             match select(bus_fut, control_fut).await {
-                Either::First(evt) => match evt {
-                    Event::Reset => {
-                        trace!("usb: reset");
-                        self.device_state = UsbDeviceState::Default;
-                        self.suspended = false;
-                        self.remote_wakeup_enabled = false;
-                        self.pending_address = 0;
-
-                        for (_, h) in self.interfaces.iter_mut() {
-                            h.reset();
-                        }
-
-                        if let Some(h) = &self.handler {
-                            h.reset();
-                        }
-                    }
-                    Event::Resume => {
-                        trace!("usb: resume");
-                        self.suspended = false;
-                        if let Some(h) = &self.handler {
-                            h.suspended(false);
-                        }
-                    }
-                    Event::Suspend => {
-                        trace!("usb: suspend");
-                        self.suspended = true;
-                        if let Some(h) = &self.handler {
-                            h.suspended(true);
-                        }
-                    }
-                },
-                Either::Second(req) => {
-                    self.in_control_handler = true;
-                    match req {
-                        Setup::DataIn(req, stage) => self.handle_control_in(req, stage).await,
-                        Setup::DataOut(req, stage) => self.handle_control_out(req, stage).await,
+                Either::First(evt) => {
+                    self.handle_bus_event(evt);
+                    if self.suspended {
+                        return;
                     }
-                    self.in_control_handler = false;
                 }
+                Either::Second(req) => match req {
+                    Setup::DataIn(req, stage) => self.handle_control_in(req, stage).await,
+                    Setup::DataOut(req, stage) => self.handle_control_out(req, stage).await,
+                },
             }
         }
     }
 
+    /// Disables the USB peripheral.
     pub async fn disable(&mut self) {
         if self.device_state != UsbDeviceState::Disabled {
             self.bus.disable().await;
             self.device_state = UsbDeviceState::Disabled;
             self.suspended = false;
             self.remote_wakeup_enabled = false;
-            self.in_control_handler = false;
 
             if let Some(h) = &self.handler {
                 h.enabled(false);
@@ -221,11 +205,26 @@ impl<'d, D: Driver<'d>> UsbDevice<'d, D> {
         }
     }
 
+    /// Waits for a resume condition on the USB bus.
+    ///
+    /// This future is cancel-safe.
+    pub async fn wait_resume(&mut self) {
+        while self.suspended {
+            let evt = self.bus.poll().await;
+            self.handle_bus_event(evt);
+        }
+    }
+
+    /// Initiates a device remote wakeup on the USB bus.
+    ///
+    /// If the bus is not suspended or remote wakeup is not enabled, an error
+    /// will be returned.
+    ///
+    /// This future may leave the bus in an inconsistent state if dropped.
+    /// After dropping the future, [`UsbDevice::disable()`] should be called
+    /// before calling any other `UsbDevice` methods to fully reset the peripheral.
     pub async fn remote_wakeup(&mut self) -> Result<(), RemoteWakeupError> {
-        if self.device_state == UsbDeviceState::Configured
-            && self.suspended
-            && self.remote_wakeup_enabled
-        {
+        if self.suspended && self.remote_wakeup_enabled {
             self.bus.remote_wakeup().await?;
             self.suspended = false;
 
@@ -239,6 +238,40 @@ impl<'d, D: Driver<'d>> UsbDevice<'d, D> {
         }
     }
 
+    fn handle_bus_event(&mut self, evt: Event) {
+        match evt {
+            Event::Reset => {
+                trace!("usb: reset");
+                self.device_state = UsbDeviceState::Default;
+                self.suspended = false;
+                self.remote_wakeup_enabled = false;
+                self.pending_address = 0;
+
+                for (_, h) in self.interfaces.iter_mut() {
+                    h.reset();
+                }
+
+                if let Some(h) = &self.handler {
+                    h.reset();
+                }
+            }
+            Event::Resume => {
+                trace!("usb: resume");
+                self.suspended = false;
+                if let Some(h) = &self.handler {
+                    h.suspended(false);
+                }
+            }
+            Event::Suspend => {
+                trace!("usb: suspend");
+                self.suspended = true;
+                if let Some(h) = &self.handler {
+                    h.suspended(true);
+                }
+            }
+        }
+    }
+
     async fn handle_control_out(&mut self, req: Request, stage: DataOutStage) {
         const CONFIGURATION_NONE_U16: u16 = CONFIGURATION_NONE as u16;
         const CONFIGURATION_VALUE_U16: u16 = CONFIGURATION_VALUE as u16;
diff --git a/examples/nrf/src/bin/usb_hid_keyboard.rs b/examples/nrf/src/bin/usb_hid_keyboard.rs
index 32659dfbb..5f03f5126 100644
--- a/examples/nrf/src/bin/usb_hid_keyboard.rs
+++ b/examples/nrf/src/bin/usb_hid_keyboard.rs
@@ -6,12 +6,11 @@
 use core::mem;
 use core::sync::atomic::{AtomicBool, Ordering};
 use defmt::*;
-use embassy::blocking_mutex::raw::CriticalSectionRawMutex;
-use embassy::channel::Channel;
+use embassy::channel::Signal;
 use embassy::executor::Spawner;
 use embassy::interrupt::InterruptExt;
 use embassy::time::Duration;
-use embassy::util::select;
+use embassy::util::{select, select3, Either, Either3};
 use embassy_nrf::gpio::{Input, Pin, Pull};
 use embassy_nrf::interrupt;
 use embassy_nrf::pac;
@@ -26,14 +25,7 @@ use usbd_hid::descriptor::{KeyboardReport, SerializedDescriptor};
 use defmt_rtt as _; // global logger
 use panic_probe as _;
 
-#[derive(defmt::Format)]
-enum Command {
-    Enable,
-    Disable,
-    RemoteWakeup,
-}
-
-static USB_COMMANDS: Channel<CriticalSectionRawMutex, Command, 1> = Channel::new();
+static ENABLE_USB: Signal<bool> = Signal::new();
 static SUSPENDED: AtomicBool = AtomicBool::new(false);
 
 fn on_power_interrupt(_: *mut ()) {
@@ -42,17 +34,13 @@ fn on_power_interrupt(_: *mut ()) {
     if regs.events_usbdetected.read().bits() != 0 {
         regs.events_usbdetected.reset();
         info!("Vbus detected, enabling USB...");
-        if USB_COMMANDS.try_send(Command::Enable).is_err() {
-            warn!("Failed to send enable command to USB channel");
-        }
+        ENABLE_USB.signal(true);
     }
 
     if regs.events_usbremoved.read().bits() != 0 {
         regs.events_usbremoved.reset();
         info!("Vbus removed, disabling USB...");
-        if USB_COMMANDS.try_send(Command::Disable).is_err() {
-            warn!("Failed to send disable command to USB channel");
-        };
+        ENABLE_USB.signal(false);
     }
 }
 
@@ -112,20 +100,35 @@ async fn main(_spawner: Spawner, p: Peripherals) {
     // Build the builder.
     let mut usb = builder.build();
 
+    let remote_wakeup = Signal::new();
+
     // Run the USB device.
     let usb_fut = async {
         enable_command().await;
         loop {
-            match select(usb.run(), USB_COMMANDS.recv()).await {
-                embassy::util::Either::First(_) => defmt::unreachable!(),
-                embassy::util::Either::Second(cmd) => match cmd {
-                    Command::Enable => warn!("Enable when already enabled!"),
-                    Command::Disable => {
+            match select(usb.run_until_suspend(), ENABLE_USB.wait()).await {
+                Either::First(_) => {}
+                Either::Second(enable) => {
+                    if enable {
+                        warn!("Enable when already enabled!");
+                    } else {
+                        usb.disable().await;
+                        enable_command().await;
+                    }
+                }
+            }
+
+            match select3(usb.wait_resume(), ENABLE_USB.wait(), remote_wakeup.wait()).await {
+                Either3::First(_) => (),
+                Either3::Second(enable) => {
+                    if enable {
+                        warn!("Enable when already enabled!");
+                    } else {
                         usb.disable().await;
                         enable_command().await;
                     }
-                    Command::RemoteWakeup => unwrap!(usb.remote_wakeup().await),
-                },
+                }
+                Either3::Third(_) => unwrap!(usb.remote_wakeup().await),
             }
         }
     };
@@ -142,7 +145,7 @@ async fn main(_spawner: Spawner, p: Peripherals) {
 
             if SUSPENDED.load(Ordering::Acquire) {
                 info!("Triggering remote wakeup");
-                USB_COMMANDS.send(Command::RemoteWakeup).await;
+                remote_wakeup.signal(());
             } else {
                 let report = KeyboardReport {
                     keycodes: [4, 0, 0, 0, 0, 0],
@@ -191,9 +194,10 @@ async fn main(_spawner: Spawner, p: Peripherals) {
 
 async fn enable_command() {
     loop {
-        match USB_COMMANDS.recv().await {
-            Command::Enable => break,
-            cmd => warn!("Received command {:?} when disabled!", cmd),
+        if ENABLE_USB.wait().await {
+            break;
+        } else {
+            warn!("Received disable signal when already disabled!");
         }
     }
 }