Remove channel and make run future cancelable

6 files changed, 129 insertions, 196 deletions

diff --git a/embassy-usb-hid/src/lib.rs b/embassy-usb-hid/src/lib.rs
index 4a1df0eab..e870becf5 100644
--- a/embassy-usb-hid/src/lib.rs
+++ b/embassy-usb-hid/src/lib.rs
@@ -11,7 +11,6 @@ use core::mem::MaybeUninit;
 use core::ops::Range;
 use core::sync::atomic::{AtomicUsize, Ordering};
 
-use embassy::blocking_mutex::raw::RawMutex;
 use embassy::time::Duration;
 use embassy_usb::driver::EndpointOut;
 use embassy_usb::{
@@ -89,8 +88,8 @@ impl<'d, D: Driver<'d>, const IN_N: usize> HidClass<'d, D, (), IN_N> {
     /// high performance uses, and a value of 255 is good for best-effort usecases.
     ///
     /// This allocates an IN endpoint only.
-    pub fn new<M: RawMutex, const OUT_N: usize>(
-        builder: &mut UsbDeviceBuilder<'d, D, M>,
+    pub fn new<const OUT_N: usize>(
+        builder: &mut UsbDeviceBuilder<'d, D>,
         state: &'d mut State<'d, IN_N, OUT_N>,
         report_descriptor: &'static [u8],
         request_handler: Option<&'d dyn RequestHandler>,
@@ -133,8 +132,8 @@ impl<'d, D: Driver<'d>, const IN_N: usize, const OUT_N: usize>
     /// high performance uses, and a value of 255 is good for best-effort usecases.
     ///
     /// This allocates two endpoints (IN and OUT).
-    pub fn with_output_ep<M: RawMutex>(
-        builder: &mut UsbDeviceBuilder<'d, D, M>,
+    pub fn with_output_ep(
+        builder: &mut UsbDeviceBuilder<'d, D>,
         state: &'d mut State<'d, IN_N, OUT_N>,
         report_descriptor: &'static [u8],
         request_handler: Option<&'d dyn RequestHandler>,
@@ -393,9 +392,9 @@ impl<'a> Control<'a> {
         }
     }
 
-    fn build<'d, D: Driver<'d>, M: RawMutex>(
+    fn build<'d, D: Driver<'d>>(
         &'d mut self,
-        builder: &mut UsbDeviceBuilder<'d, D, M>,
+        builder: &mut UsbDeviceBuilder<'d, D>,
         ep_out: Option<&D::EndpointOut>,
         ep_in: &D::EndpointIn,
     ) {
diff --git a/embassy-usb-serial/src/lib.rs b/embassy-usb-serial/src/lib.rs
index 7f006c0fd..7b25398d0 100644
--- a/embassy-usb-serial/src/lib.rs
+++ b/embassy-usb-serial/src/lib.rs
@@ -8,7 +8,6 @@ pub(crate) mod fmt;
 use core::cell::Cell;
 use core::mem::{self, MaybeUninit};
 use core::sync::atomic::{AtomicBool, Ordering};
-use embassy::blocking_mutex::raw::RawMutex;
 use embassy::blocking_mutex::CriticalSectionMutex;
 use embassy_usb::control::{self, ControlHandler, InResponse, OutResponse, Request};
 use embassy_usb::driver::{Endpoint, EndpointError, EndpointIn, EndpointOut};
@@ -163,8 +162,8 @@ impl<'d> ControlHandler for Control<'d> {
 impl<'d, D: Driver<'d>> CdcAcmClass<'d, D> {
     /// Creates a new CdcAcmClass with the provided UsbBus and max_packet_size in bytes. For
     /// full-speed devices, max_packet_size has to be one of 8, 16, 32 or 64.
-    pub fn new<M: RawMutex>(
-        builder: &mut UsbDeviceBuilder<'d, D, M>,
+    pub fn new(
+        builder: &mut UsbDeviceBuilder<'d, D>,
         state: &'d mut State<'d>,
         max_packet_size: u16,
     ) -> Self {
diff --git a/embassy-usb/src/builder.rs b/embassy-usb/src/builder.rs
index 30d31ac74..486672055 100644
--- a/embassy-usb/src/builder.rs
+++ b/embassy-usb/src/builder.rs
@@ -1,9 +1,5 @@
-use embassy::blocking_mutex::raw::{NoopRawMutex, RawMutex};
-use embassy::channel::Channel;
 use heapless::Vec;
 
-use crate::DeviceCommand;
-
 use super::control::ControlHandler;
 use super::descriptor::{BosWriter, DescriptorWriter};
 use super::driver::{Driver, EndpointAllocError};
@@ -98,11 +94,6 @@ pub struct Config<'a> {
     /// Default: 100mA
     /// Max: 500mA
     pub max_power: u16,
-
-    /// Whether the USB bus should be enabled when built.
-    ///
-    /// Default: true
-    pub start_enabled: bool,
 }
 
 impl<'a> Config<'a> {
@@ -122,18 +113,16 @@ impl<'a> Config<'a> {
             supports_remote_wakeup: false,
             composite_with_iads: false,
             max_power: 100,
-            start_enabled: true,
         }
     }
 }
 
 /// Used to build new [`UsbDevice`]s.
-pub struct UsbDeviceBuilder<'d, D: Driver<'d>, M: RawMutex> {
+pub struct UsbDeviceBuilder<'d, D: Driver<'d>> {
     config: Config<'d>,
     handler: Option<&'d dyn DeviceStateHandler>,
     interfaces: Vec<(u8, &'d mut dyn ControlHandler), MAX_INTERFACE_COUNT>,
     control_buf: &'d mut [u8],
-    commands: Option<&'d Channel<M, DeviceCommand, 1>>,
 
     driver: D,
     next_interface_number: u8,
@@ -145,7 +134,7 @@ pub struct UsbDeviceBuilder<'d, D: Driver<'d>, M: RawMutex> {
     pub bos_descriptor: BosWriter<'d>,
 }
 
-impl<'d, D: Driver<'d>> UsbDeviceBuilder<'d, D, NoopRawMutex> {
+impl<'d, D: Driver<'d>> UsbDeviceBuilder<'d, D> {
     /// Creates a builder for constructing a new [`UsbDevice`].
     ///
     /// `control_buf` is a buffer used for USB control request data. It should be sized
@@ -160,57 +149,6 @@ impl<'d, D: Driver<'d>> UsbDeviceBuilder<'d, D, NoopRawMutex> {
         control_buf: &'d mut [u8],
         handler: Option<&'d dyn DeviceStateHandler>,
     ) -> Self {
-        Self::new_inner(
-            driver,
-            config,
-            device_descriptor_buf,
-            config_descriptor_buf,
-            bos_descriptor_buf,
-            control_buf,
-            handler,
-            None,
-        )
-    }
-}
-
-impl<'d, D: Driver<'d>, M: RawMutex> UsbDeviceBuilder<'d, D, M> {
-    /// Creates a builder for constructing a new [`UsbDevice`].
-    ///
-    /// `control_buf` is a buffer used for USB control request data. It should be sized
-    /// large enough for the length of the largest control request (in or out)
-    /// anticipated by any class added to the device.
-    pub fn new_with_channel(
-        driver: D,
-        config: Config<'d>,
-        device_descriptor_buf: &'d mut [u8],
-        config_descriptor_buf: &'d mut [u8],
-        bos_descriptor_buf: &'d mut [u8],
-        control_buf: &'d mut [u8],
-        handler: Option<&'d dyn DeviceStateHandler>,
-        channel: &'d Channel<M, DeviceCommand, 1>,
-    ) -> Self {
-        Self::new_inner(
-            driver,
-            config,
-            device_descriptor_buf,
-            config_descriptor_buf,
-            bos_descriptor_buf,
-            control_buf,
-            handler,
-            Some(channel),
-        )
-    }
-
-    fn new_inner(
-        driver: D,
-        config: Config<'d>,
-        device_descriptor_buf: &'d mut [u8],
-        config_descriptor_buf: &'d mut [u8],
-        bos_descriptor_buf: &'d mut [u8],
-        control_buf: &'d mut [u8],
-        handler: Option<&'d dyn DeviceStateHandler>,
-        channel: Option<&'d Channel<M, DeviceCommand, 1>>,
-    ) -> Self {
         // Magic values specified in USB-IF ECN on IADs.
         if config.composite_with_iads
             && (config.device_class != 0xEF
@@ -243,8 +181,6 @@ impl<'d, D: Driver<'d>, M: RawMutex> UsbDeviceBuilder<'d, D, M> {
             config,
             interfaces: Vec::new(),
             control_buf,
-            commands: channel,
-
             next_interface_number: 0,
             next_string_index: 4,
 
@@ -255,7 +191,7 @@ impl<'d, D: Driver<'d>, M: RawMutex> UsbDeviceBuilder<'d, D, M> {
     }
 
     /// Creates the [`UsbDevice`] instance with the configuration in this builder.
-    pub fn build(mut self) -> UsbDevice<'d, D, M> {
+    pub fn build(mut self) -> UsbDevice<'d, D> {
         self.config_descriptor.end_configuration();
         self.bos_descriptor.end_bos();
 
@@ -263,7 +199,6 @@ impl<'d, D: Driver<'d>, M: RawMutex> UsbDeviceBuilder<'d, D, M> {
             self.driver,
             self.config,
             self.handler,
-            self.commands,
             self.device_descriptor.into_buf(),
             self.config_descriptor.into_buf(),
             self.bos_descriptor.writer.into_buf(),
diff --git a/embassy-usb/src/lib.rs b/embassy-usb/src/lib.rs
index 102ccbb3a..ccea8bc7a 100644
--- a/embassy-usb/src/lib.rs
+++ b/embassy-usb/src/lib.rs
@@ -10,19 +10,14 @@ pub mod control;
 pub mod descriptor;
 pub mod driver;
 pub mod types;
-mod util;
 
-use driver::Unsupported;
-use embassy::blocking_mutex::raw::{NoopRawMutex, RawMutex};
-use embassy::channel::Channel;
-use embassy::util::{select3, Either3};
+use embassy::util::{select, Either};
 use heapless::Vec;
 
 use self::control::*;
 use self::descriptor::*;
 use self::driver::{Bus, Driver, Event};
 use self::types::*;
-use self::util::*;
 
 pub use self::builder::Config;
 pub use self::builder::UsbDeviceBuilder;
@@ -47,6 +42,19 @@ pub enum UsbDeviceState {
     Configured,
 }
 
+#[derive(PartialEq, Eq, Copy, Clone, Debug)]
+#[cfg_attr(feature = "defmt", derive(defmt::Format))]
+pub enum RemoteWakeupError {
+    InvalidState,
+    Unsupported,
+}
+
+impl From<driver::Unsupported> for RemoteWakeupError {
+    fn from(_: driver::Unsupported) -> Self {
+        RemoteWakeupError::Unsupported
+    }
+}
+
 /// The bConfiguration value for the not configured state.
 pub const CONFIGURATION_NONE: u8 = 0;
 
@@ -58,16 +66,11 @@ pub const DEFAULT_ALTERNATE_SETTING: u8 = 0;
 
 pub const MAX_INTERFACE_COUNT: usize = 4;
 
-#[derive(PartialEq, Eq, Copy, Clone, Debug)]
-#[cfg_attr(feature = "defmt", derive(defmt::Format))]
-pub enum DeviceCommand {
-    Enable,
-    Disable,
-    RemoteWakeup,
-}
-
 /// A handler trait for changes in the device state of the [UsbDevice].
 pub trait DeviceStateHandler {
+    /// Called when the USB device has been enabled or disabled.
+    fn enabled(&self, _enabled: bool) {}
+
     /// Called when the host resets the device.
     fn reset(&self) {}
 
@@ -82,15 +85,11 @@ pub trait DeviceStateHandler {
 
     /// Called when remote wakeup feature is enabled or disabled.
     fn remote_wakeup_enabled(&self, _enabled: bool) {}
-
-    /// Called when the USB device has been disabled.
-    fn disabled(&self) {}
 }
 
-pub struct UsbDevice<'d, D: Driver<'d>, M: RawMutex = NoopRawMutex> {
+pub struct UsbDevice<'d, D: Driver<'d>> {
     bus: D::Bus,
     handler: Option<&'d dyn DeviceStateHandler>,
-    commands: Option<&'d Channel<M, DeviceCommand, 1>>,
     control: ControlPipe<D::ControlPipe>,
 
     config: Config<'d>,
@@ -101,6 +100,7 @@ pub struct UsbDevice<'d, D: Driver<'d>, M: RawMutex = NoopRawMutex> {
 
     device_state: UsbDeviceState,
     suspended: bool,
+    in_control_handler: bool,
     remote_wakeup_enabled: bool,
     self_powered: bool,
     pending_address: u8,
@@ -108,18 +108,17 @@ pub struct UsbDevice<'d, D: Driver<'d>, M: RawMutex = NoopRawMutex> {
     interfaces: Vec<(u8, &'d mut dyn ControlHandler), MAX_INTERFACE_COUNT>,
 }
 
-impl<'d, D: Driver<'d>, M: RawMutex> UsbDevice<'d, D, M> {
+impl<'d, D: Driver<'d>> UsbDevice<'d, D> {
     pub(crate) fn build(
         mut driver: D,
         config: Config<'d>,
         handler: Option<&'d dyn DeviceStateHandler>,
-        commands: Option<&'d Channel<M, DeviceCommand, 1>>,
         device_descriptor: &'d [u8],
         config_descriptor: &'d [u8],
         bos_descriptor: &'d [u8],
         interfaces: Vec<(u8, &'d mut dyn ControlHandler), MAX_INTERFACE_COUNT>,
         control_buf: &'d mut [u8],
-    ) -> UsbDevice<'d, D, M> {
+    ) -> UsbDevice<'d, D> {
         let control = driver
             .alloc_control_pipe(config.max_packet_size_0 as u16)
             .expect("failed to alloc control endpoint");
@@ -132,14 +131,14 @@ impl<'d, D: Driver<'d>, M: RawMutex> UsbDevice<'d, D, M> {
             bus,
             config,
             handler,
-            commands,
             control: ControlPipe::new(control),
             device_descriptor,
             config_descriptor,
             bos_descriptor,
             control_buf,
-            device_state: UsbDeviceState::Default,
+            device_state: UsbDeviceState::Disabled,
             suspended: false,
+            in_control_handler: false,
             remote_wakeup_enabled: false,
             self_powered: false,
             pending_address: 0,
@@ -148,19 +147,24 @@ impl<'d, D: Driver<'d>, M: RawMutex> UsbDevice<'d, D, M> {
     }
 
     pub async fn run(&mut self) -> ! {
-        if self.config.start_enabled {
+        if self.device_state == UsbDeviceState::Disabled {
             self.bus.enable().await;
-        } else {
-            self.wait_for_enable().await
+            self.device_state = UsbDeviceState::Default;
+
+            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();
-            let commands_fut = recv_or_wait(self.commands);
-
-            match select3(bus_fut, control_fut, commands_fut).await {
-                Either3::First(evt) => match evt {
+            match select(bus_fut, control_fut).await {
+                Either::First(evt) => match evt {
                     Event::Reset => {
                         trace!("usb: reset");
                         self.device_state = UsbDeviceState::Default;
@@ -191,54 +195,48 @@ impl<'d, D: Driver<'d>, M: RawMutex> UsbDevice<'d, D, M> {
                         }
                     }
                 },
-                Either3::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,
-                },
-                Either3::Third(cmd) => match cmd {
-                    DeviceCommand::Enable => warn!("usb: Enable command received while enabled."),
-                    DeviceCommand::Disable => {
-                        trace!("usb: disable");
-                        self.bus.disable().await;
-                        self.device_state = UsbDeviceState::Disabled;
-                        if let Some(h) = &self.handler {
-                            h.disabled();
-                        }
-                        self.wait_for_enable().await;
+                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,
                     }
-                    DeviceCommand::RemoteWakeup => {
-                        trace!("usb: remote wakeup");
-                        if self.suspended && self.remote_wakeup_enabled {
-                            match self.bus.remote_wakeup().await {
-                                Ok(()) => {
-                                    self.suspended = false;
-                                    if let Some(h) = &self.handler {
-                                        h.suspended(false);
-                                    }
-                                }
-                                Err(Unsupported) => warn!("Remote wakeup is unsupported!"),
-                            }
-                        } else {
-                            warn!("Remote wakeup requested when not enabled or not suspended.");
-                        }
-                    }
-                },
+                    self.in_control_handler = false;
+                }
             }
         }
     }
 
-    async fn wait_for_enable(&mut self) {
-        loop {
-            // When disabled just wait until we're told to re-enable
-            match recv_or_wait(self.commands).await {
-                DeviceCommand::Enable => break,
-                cmd => warn!("usb: {:?} received while disabled", cmd),
+    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);
             }
         }
+    }
+
+    pub async fn remote_wakeup(&mut self) -> Result<(), RemoteWakeupError> {
+        if self.device_state == UsbDeviceState::Configured
+            && self.suspended
+            && self.remote_wakeup_enabled
+        {
+            self.bus.remote_wakeup().await?;
+            self.suspended = false;
 
-        trace!("usb: enable");
-        self.bus.enable().await;
-        self.device_state = UsbDeviceState::Default;
+            if let Some(h) = &self.handler {
+                h.suspended(false);
+            }
+
+            Ok(())
+        } else {
+            Err(RemoteWakeupError::InvalidState)
+        }
     }
 
     async fn handle_control_out(&mut self, req: Request, stage: DataOutStage) {
diff --git a/embassy-usb/src/util.rs b/embassy-usb/src/util.rs
deleted file mode 100644
index 3f20262c6..000000000
--- a/embassy-usb/src/util.rs
+++ /dev/null
@@ -1,33 +0,0 @@
-use core::future::Future;
-use core::marker::PhantomData;
-use core::pin::Pin;
-use core::task::{Context, Poll};
-
-use embassy::blocking_mutex::raw::RawMutex;
-use embassy::channel::Channel;
-
-pub struct Pending<T> {
-    _phantom: PhantomData<T>,
-}
-
-impl<T> Pending<T> {
-    fn new() -> Self {
-        Pending {
-            _phantom: PhantomData,
-        }
-    }
-}
-
-impl<T> Future for Pending<T> {
-    type Output = T;
-    fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Self::Output> {
-        Poll::Pending
-    }
-}
-
-pub async fn recv_or_wait<M: RawMutex, T, const N: usize>(ch: Option<&Channel<M, T, N>>) -> T {
-    match ch {
-        Some(ch) => ch.recv().await,
-        None => Pending::new().await,
-    }
-}
diff --git a/examples/nrf/src/bin/usb_hid_keyboard.rs b/examples/nrf/src/bin/usb_hid_keyboard.rs
index fb3a198a7..32659dfbb 100644
--- a/examples/nrf/src/bin/usb_hid_keyboard.rs
+++ b/examples/nrf/src/bin/usb_hid_keyboard.rs
@@ -11,13 +11,14 @@ use embassy::channel::Channel;
 use embassy::executor::Spawner;
 use embassy::interrupt::InterruptExt;
 use embassy::time::Duration;
+use embassy::util::select;
 use embassy_nrf::gpio::{Input, Pin, Pull};
 use embassy_nrf::interrupt;
 use embassy_nrf::pac;
 use embassy_nrf::usb::Driver;
 use embassy_nrf::Peripherals;
 use embassy_usb::control::OutResponse;
-use embassy_usb::{Config, DeviceCommand, DeviceStateHandler, UsbDeviceBuilder};
+use embassy_usb::{Config, DeviceStateHandler, UsbDeviceBuilder};
 use embassy_usb_hid::{HidClass, ReportId, RequestHandler, State};
 use futures::future::join;
 use usbd_hid::descriptor::{KeyboardReport, SerializedDescriptor};
@@ -25,7 +26,14 @@ use usbd_hid::descriptor::{KeyboardReport, SerializedDescriptor};
 use defmt_rtt as _; // global logger
 use panic_probe as _;
 
-static USB_COMMANDS: Channel<CriticalSectionRawMutex, DeviceCommand, 1> = Channel::new();
+#[derive(defmt::Format)]
+enum Command {
+    Enable,
+    Disable,
+    RemoteWakeup,
+}
+
+static USB_COMMANDS: Channel<CriticalSectionRawMutex, Command, 1> = Channel::new();
 static SUSPENDED: AtomicBool = AtomicBool::new(false);
 
 fn on_power_interrupt(_: *mut ()) {
@@ -34,7 +42,7 @@ 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(DeviceCommand::Enable).is_err() {
+        if USB_COMMANDS.try_send(Command::Enable).is_err() {
             warn!("Failed to send enable command to USB channel");
         }
     }
@@ -42,7 +50,7 @@ fn on_power_interrupt(_: *mut ()) {
     if regs.events_usbremoved.read().bits() != 0 {
         regs.events_usbremoved.reset();
         info!("Vbus removed, disabling USB...");
-        if USB_COMMANDS.try_send(DeviceCommand::Disable).is_err() {
+        if USB_COMMANDS.try_send(Command::Disable).is_err() {
             warn!("Failed to send disable command to USB channel");
         };
     }
@@ -69,7 +77,6 @@ async fn main(_spawner: Spawner, p: Peripherals) {
     config.max_power = 100;
     config.max_packet_size_0 = 64;
     config.supports_remote_wakeup = true;
-    config.start_enabled = false;
 
     // Create embassy-usb DeviceBuilder using the driver and config.
     // It needs some buffers for building the descriptors.
@@ -82,7 +89,7 @@ async fn main(_spawner: Spawner, p: Peripherals) {
 
     let mut state = State::<8, 1>::new();
 
-    let mut builder = UsbDeviceBuilder::new_with_channel(
+    let mut builder = UsbDeviceBuilder::new(
         driver,
         config,
         &mut device_descriptor,
@@ -90,7 +97,6 @@ async fn main(_spawner: Spawner, p: Peripherals) {
         &mut bos_descriptor,
         &mut control_buf,
         Some(&device_state_handler),
-        &USB_COMMANDS,
     );
 
     // Create classes on the builder.
@@ -107,7 +113,22 @@ async fn main(_spawner: Spawner, p: Peripherals) {
     let mut usb = builder.build();
 
     // Run the USB device.
-    let usb_fut = usb.run();
+    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 => {
+                        usb.disable().await;
+                        enable_command().await;
+                    }
+                    Command::RemoteWakeup => unwrap!(usb.remote_wakeup().await),
+                },
+            }
+        }
+    };
 
     let mut button = Input::new(p.P0_11.degrade(), Pull::Up);
 
@@ -121,7 +142,7 @@ async fn main(_spawner: Spawner, p: Peripherals) {
 
             if SUSPENDED.load(Ordering::Acquire) {
                 info!("Triggering remote wakeup");
-                USB_COMMANDS.send(DeviceCommand::RemoteWakeup).await;
+                USB_COMMANDS.send(Command::RemoteWakeup).await;
             } else {
                 let report = KeyboardReport {
                     keycodes: [4, 0, 0, 0, 0, 0],
@@ -168,6 +189,15 @@ async fn main(_spawner: Spawner, p: Peripherals) {
     join(usb_fut, join(in_fut, out_fut)).await;
 }
 
+async fn enable_command() {
+    loop {
+        match USB_COMMANDS.recv().await {
+            Command::Enable => break,
+            cmd => warn!("Received command {:?} when disabled!", cmd),
+        }
+    }
+}
+
 struct MyRequestHandler {}
 
 impl RequestHandler for MyRequestHandler {
@@ -204,6 +234,16 @@ impl MyDeviceStateHandler {
 }
 
 impl DeviceStateHandler for MyDeviceStateHandler {
+    fn enabled(&self, enabled: bool) {
+        self.configured.store(false, Ordering::Relaxed);
+        SUSPENDED.store(false, Ordering::Release);
+        if enabled {
+            info!("Device enabled");
+        } else {
+            info!("Device disabled");
+        }
+    }
+
     fn reset(&self) {
         self.configured.store(false, Ordering::Relaxed);
         info!("Bus reset, the Vbus current limit is 100mA");
@@ -240,9 +280,4 @@ impl DeviceStateHandler for MyDeviceStateHandler {
             }
         }
     }
-
-    fn disabled(&self) {
-        self.configured.store(false, Ordering::Relaxed);
-        info!("Device disabled");
-    }
 }