From 50f6e24df67df8088e3aedcf3c0a82138a39293b Mon Sep 17 00:00:00 2001
From: Haobo Gu <haobogu@outlook.com>
Date: Fri, 23 Aug 2024 11:44:40 +0800
Subject: [PATCH] fix(usb): check ep status before transfer

Signed-off-by: Haobo Gu <haobogu@outlook.com>
---
 src/usb/bus.rs      | 16 +++++++++++++++-
 src/usb/endpoint.rs | 23 +++++++++++++++++++++++
 2 files changed, 38 insertions(+), 1 deletion(-)

diff --git a/src/usb/bus.rs b/src/usb/bus.rs
index 228df47..4ada95a 100644
--- a/src/usb/bus.rs
+++ b/src/usb/bus.rs
@@ -8,7 +8,7 @@ use embedded_hal::delay::DelayNs;
 use hpm_metapac::usb::regs::*;
 use riscv::delay::McycleDelay;
 
-use super::{init_qhd, Instance, ENDPOINT_COUNT};
+use super::{init_qhd, Instance, ENDPOINT_COUNT, EP_IN_WAKERS, EP_OUT_WAKERS};
 use crate::usb::{reset_dcd_data, EpConfig, BUS_WAKER, DCD_DATA, IRQ_RESET, IRQ_SUSPEND};
 
 /// USB bus
@@ -75,6 +75,12 @@ impl<T: Instance> embassy_usb_driver::Bus for Bus<T> {
             if IRQ_RESET.load(Ordering::Acquire) {
                 IRQ_RESET.store(false, Ordering::Relaxed);
 
+                // Disable all endpoints except ep0
+                for i in 1..ENDPOINT_COUNT {
+                    self.endpoint_close(EndpointAddress::from_parts(i, Direction::In));
+                    self.endpoint_close(EndpointAddress::from_parts(i, Direction::Out));
+                }
+
                 // Set device addr to 0
                 self.device_set_address(0);
 
@@ -90,6 +96,13 @@ impl<T: Instance> embassy_usb_driver::Bus for Bus<T> {
                     max_packet_size: 64,
                 });
 
+                for w in &EP_IN_WAKERS {
+                    w.wake()
+                }
+                for w in &EP_OUT_WAKERS {
+                    w.wake()
+                }
+
                 // Reset bus
                 self.device_bus_reset(64);
 
@@ -355,6 +368,7 @@ impl<T: Instance> Bus<T> {
 
     /// Open the endpoint
     fn endpoint_open(&mut self, ep_config: EpConfig) {
+        // defmt::info!("Enabling endpoint: {:?}", ep_config.ep_addr);
         if ep_config.ep_addr.index() >= ENDPOINT_COUNT {
             return;
         }
diff --git a/src/usb/endpoint.rs b/src/usb/endpoint.rs
index 33e9d15..ebe363d 100644
--- a/src/usb/endpoint.rs
+++ b/src/usb/endpoint.rs
@@ -146,6 +146,16 @@ impl<'d, T: Instance> Endpoint<'d, T> {
             r.endptctrl(self.info.addr.index() as usize).modify(|w| w.set_rxs(true));
         }
     }
+
+    pub(crate) fn enabled(&self) -> bool {
+        let r = T::info().regs;
+        let ep_num = self.info.addr.index();
+        if self.info.addr.is_in() {
+            r.endptctrl(ep_num).read().txe()
+        } else {
+            r.endptctrl(ep_num).read().rxe()
+        }
+    }
 }
 
 impl<'d, T: Instance> embassy_usb_driver::Endpoint for Endpoint<'d, T> {
@@ -156,6 +166,9 @@ impl<'d, T: Instance> embassy_usb_driver::Endpoint for Endpoint<'d, T> {
 
     /// Wait for the endpoint to be enabled.
     async fn wait_enabled(&mut self) {
+        if self.enabled() {
+            return;
+        }
         let i = self.info.addr.index();
         poll_fn(|cx| {
             let r = T::info().regs;
@@ -183,6 +196,9 @@ impl<'d, T: Instance> EndpointOut for Endpoint<'d, T> {
     ///
     /// This should also clear any NAK flags and prepare the endpoint to receive the next packet.
     async fn read(&mut self, buf: &mut [u8]) -> Result<usize, embassy_usb_driver::EndpointError> {
+        if !self.enabled() {
+            return Err(embassy_usb_driver::EndpointError::Disabled);
+        }
         let r = T::info().regs;
         let ep_num = self.info.addr.index();
 
@@ -195,6 +211,8 @@ impl<'d, T: Instance> EndpointOut for Endpoint<'d, T> {
                 // Clear the flag
                 r.endptcomplete().modify(|w| w.set_erce(1 << ep_num));
                 Poll::Ready(())
+            } else if !r.endptctrl(ep_num).read().rxe() {
+                Poll::Ready(())
             } else {
                 Poll::Pending
             }
@@ -212,6 +230,9 @@ impl<'d, T: Instance> EndpointOut for Endpoint<'d, T> {
 impl<'d, T: Instance> EndpointIn for Endpoint<'d, T> {
     /// Write a single packet of data to the endpoint.
     async fn write(&mut self, buf: &[u8]) -> Result<(), embassy_usb_driver::EndpointError> {
+        if !self.enabled() {
+            return Err(embassy_usb_driver::EndpointError::Disabled);
+        }
         let r = T::info().regs;
         let ep_num = self.info.addr.index();
 
@@ -223,6 +244,8 @@ impl<'d, T: Instance> EndpointIn for Endpoint<'d, T> {
             if r.endptcomplete().read().etce() & (1 << ep_num) != 0 {
                 r.endptcomplete().modify(|w| w.set_etce(1 << ep_num));
                 Poll::Ready(())
+            } else if !r.endptctrl(ep_num).read().txe() {
+                Poll::Ready(())
             } else {
                 Poll::Pending
             }