3 files changed, 161 insertions, 130 deletions

diff --git a/embassy-nrf/src/usb.rs b/embassy-nrf/src/usb.rs
index 1057d880c..570d1c95b 100644
--- a/embassy-nrf/src/usb.rs
+++ b/embassy-nrf/src/usb.rs
@@ -186,7 +186,6 @@ impl<'d, T: Instance> driver::Driver<'d> for Driver<'d, T> {
         Ok(ControlPipe {
             _phantom: PhantomData,
             max_packet_size,
-            request: None,
         })
     }
 
@@ -502,72 +501,19 @@ impl<'d, T: Instance> driver::EndpointIn for Endpoint<'d, T, In> {
 pub struct ControlPipe<'d, T: Instance> {
     _phantom: PhantomData<&'d mut T>,
     max_packet_size: u16,
-    request: Option<Request>,
-}
-
-impl<'d, T: Instance> ControlPipe<'d, T> {
-    async fn read(&mut self, buf: &mut [u8]) -> Result<usize, ReadError> {
-        let regs = T::regs();
-
-        // Wait until ready
-        regs.intenset.write(|w| w.ep0datadone().set());
-        poll_fn(|cx| {
-            EP_OUT_WAKERS[0].register(cx.waker());
-            let regs = T::regs();
-            if regs
-                .events_ep0datadone
-                .read()
-                .events_ep0datadone()
-                .bit_is_set()
-            {
-                Poll::Ready(())
-            } else {
-                Poll::Pending
-            }
-        })
-        .await;
-
-        unsafe { read_dma::<T>(0, buf) }
-    }
-
-    async fn write(&mut self, buf: &[u8], last_chunk: bool) {
-        let regs = T::regs();
-        regs.events_ep0datadone.reset();
-        unsafe { write_dma::<T>(0, buf) }
-
-        regs.shorts
-            .modify(|_, w| w.ep0datadone_ep0status().bit(last_chunk));
-
-        regs.intenset.write(|w| w.ep0datadone().set());
-        let res = with_timeout(
-            Duration::from_millis(10),
-            poll_fn(|cx| {
-                EP_IN_WAKERS[0].register(cx.waker());
-                let regs = T::regs();
-                if regs.events_ep0datadone.read().bits() != 0 {
-                    Poll::Ready(())
-                } else {
-                    Poll::Pending
-                }
-            }),
-        )
-        .await;
-
-        if res.is_err() {
-            error!("ControlPipe::write timed out.");
-        }
-    }
 }
 
 impl<'d, T: Instance> driver::ControlPipe for ControlPipe<'d, T> {
     type SetupFuture<'a> = impl Future<Output = Request> + 'a where Self: 'a;
     type DataOutFuture<'a> = impl Future<Output = Result<usize, ReadError>> + 'a where Self: 'a;
-    type AcceptInFuture<'a> = impl Future<Output = ()> + 'a where Self: 'a;
+    type DataInFuture<'a> = impl Future<Output = ()> + 'a where Self: 'a;
+
+    fn max_packet_size(&self) -> usize {
+        usize::from(self.max_packet_size)
+    }
 
     fn setup<'a>(&'a mut self) -> Self::SetupFuture<'a> {
         async move {
-            assert!(self.request.is_none());
-
             let regs = T::regs();
 
             // Wait for SETUP packet
@@ -605,29 +551,65 @@ impl<'d, T: Instance> driver::ControlPipe for ControlPipe<'d, T> {
                     .write(|w| w.tasks_ep0rcvout().set_bit());
             }
 
-            self.request = Some(req);
             req
         }
     }
 
     fn data_out<'a>(&'a mut self, buf: &'a mut [u8]) -> Self::DataOutFuture<'a> {
         async move {
-            let req = unwrap!(self.request);
-            assert!(req.direction == UsbDirection::Out);
-            assert!(req.length > 0);
-
-            let req_length = usize::from(req.length);
-            let max_packet_size = usize::from(self.max_packet_size);
-            let mut total = 0;
-            for chunk in buf.chunks_mut(max_packet_size) {
-                let size = self.read(chunk).await?;
-                total += size;
-                if size < max_packet_size || total == req_length {
-                    break;
+            let regs = T::regs();
+
+            // Wait until ready
+            regs.intenset.write(|w| w.ep0datadone().set());
+            poll_fn(|cx| {
+                EP_OUT_WAKERS[0].register(cx.waker());
+                let regs = T::regs();
+                if regs
+                    .events_ep0datadone
+                    .read()
+                    .events_ep0datadone()
+                    .bit_is_set()
+                {
+                    Poll::Ready(())
+                } else {
+                    Poll::Pending
                 }
+            })
+            .await;
+
+            unsafe { read_dma::<T>(0, buf) }
+        }
+    }
+
+    fn data_in<'a>(&'a mut self, buf: &'a [u8], last_packet: bool) -> Self::DataInFuture<'a> {
+        async move {
+            let regs = T::regs();
+            regs.events_ep0datadone.reset();
+            unsafe {
+                write_dma::<T>(0, buf);
             }
 
-            Ok(total)
+            regs.shorts
+                .modify(|_, w| w.ep0datadone_ep0status().bit(last_packet));
+
+            regs.intenset.write(|w| w.ep0datadone().set());
+            let res = with_timeout(
+                Duration::from_millis(10),
+                poll_fn(|cx| {
+                    EP_IN_WAKERS[0].register(cx.waker());
+                    let regs = T::regs();
+                    if regs.events_ep0datadone.read().bits() != 0 {
+                        Poll::Ready(())
+                    } else {
+                        Poll::Pending
+                    }
+                }),
+            )
+            .await;
+
+            if res.is_err() {
+                error!("ControlPipe::data_in timed out.");
+            }
         }
     }
 
@@ -636,37 +618,12 @@ impl<'d, T: Instance> driver::ControlPipe for ControlPipe<'d, T> {
         let regs = T::regs();
         regs.tasks_ep0status
             .write(|w| w.tasks_ep0status().bit(true));
-        self.request = None;
-    }
-
-    fn accept_in<'a>(&'a mut self, buf: &'a [u8]) -> Self::AcceptInFuture<'a> {
-        async move {
-            #[cfg(feature = "defmt")]
-            debug!("control in accept {:x}", buf);
-            #[cfg(not(feature = "defmt"))]
-            debug!("control in accept {:x?}", buf);
-            let req = unwrap!(self.request);
-            assert!(req.direction == UsbDirection::In);
-
-            let req_len = usize::from(req.length);
-            let len = buf.len().min(req_len);
-            let need_zlp = len != req_len && (len % usize::from(self.max_packet_size)) == 0;
-            let mut chunks = buf[0..len]
-                .chunks(usize::from(self.max_packet_size))
-                .chain(need_zlp.then(|| -> &[u8] { &[] }));
-            while let Some(chunk) = chunks.next() {
-                self.write(chunk, chunks.size_hint().0 == 0).await;
-            }
-
-            self.request = None;
-        }
     }
 
     fn reject(&mut self) {
         debug!("control reject");
         let regs = T::regs();
         regs.tasks_ep0stall.write(|w| w.tasks_ep0stall().bit(true));
-        self.request = None;
     }
 }
 
diff --git a/embassy-usb/src/driver.rs b/embassy-usb/src/driver.rs
index 1c6ba1f52..82b59bd1e 100644
--- a/embassy-usb/src/driver.rs
+++ b/embassy-usb/src/driver.rs
@@ -137,30 +137,35 @@ pub trait ControlPipe {
     type DataOutFuture<'a>: Future<Output = Result<usize, ReadError>> + 'a
     where
         Self: 'a;
-    type AcceptInFuture<'a>: Future<Output = ()> + 'a
+    type DataInFuture<'a>: Future<Output = ()> + 'a
     where
         Self: 'a;
 
+    /// Maximum packet size for the control pipe
+    fn max_packet_size(&self) -> usize;
+
     /// Reads a single setup packet from the endpoint.
     fn setup<'a>(&'a mut self) -> Self::SetupFuture<'a>;
 
-    /// Reads the data packet of a control write sequence.
+    /// Reads a DATA OUT packet into `buf` in response to a control write request.
     ///
     /// Must be called after `setup()` for requests with `direction` of `Out`
     /// and `length` greater than zero.
-    ///
-    /// `buf.len()` must be greater than or equal to the request's `length`.
     fn data_out<'a>(&'a mut self, buf: &'a mut [u8]) -> Self::DataOutFuture<'a>;
 
-    /// Accepts a control request.
-    fn accept(&mut self);
+    /// Sends a DATA IN packet with `data` in response to a control read request.
+    ///
+    /// If `last_packet` is true, the STATUS packet will be ACKed following the transfer of `data`.
+    fn data_in<'a>(&'a mut self, data: &'a [u8], last_packet: bool) -> Self::DataInFuture<'a>;
 
-    /// Accepts a control read request with `data`.
+    /// Accepts a control request.
     ///
-    /// `data.len()` must be less than or equal to the request's `length`.
-    fn accept_in<'a>(&'a mut self, data: &'a [u8]) -> Self::AcceptInFuture<'a>;
+    /// Causes the STATUS packet for the current request to be ACKed.
+    fn accept(&mut self);
 
     /// Rejects a control request.
+    ///
+    /// Sets a STALL condition on the pipe to indicate an error.
     fn reject(&mut self);
 }
 
diff --git a/embassy-usb/src/lib.rs b/embassy-usb/src/lib.rs
index 32b67a766..77a9c33be 100644
--- a/embassy-usb/src/lib.rs
+++ b/embassy-usb/src/lib.rs
@@ -55,7 +55,7 @@ pub const MAX_INTERFACE_COUNT: usize = 4;
 
 pub struct UsbDevice<'d, D: Driver<'d>> {
     bus: D::Bus,
-    control: D::ControlPipe,
+    control: ControlPipe<D::ControlPipe>,
 
     config: Config<'d>,
     device_descriptor: &'d [u8],
@@ -92,7 +92,10 @@ impl<'d, D: Driver<'d>> UsbDevice<'d, D> {
         Self {
             bus: driver,
             config,
-            control,
+            control: ControlPipe {
+                control,
+                request: None,
+            },
             device_descriptor,
             config_descriptor,
             bos_descriptor,
@@ -140,32 +143,19 @@ impl<'d, D: Driver<'d>> UsbDevice<'d, D> {
         }
     }
 
-    async fn control_in_accept_writer(
-        &mut self,
-        req: Request,
-        f: impl FnOnce(&mut DescriptorWriter),
-    ) {
-        let mut buf = [0; 256];
-        let mut w = DescriptorWriter::new(&mut buf);
-        f(&mut w);
-        let pos = w.position().min(usize::from(req.length));
-        self.control.accept_in(&buf[..pos]).await;
-    }
-
     async fn handle_control_out(&mut self, req: Request) {
         const CONFIGURATION_NONE_U16: u16 = CONFIGURATION_NONE as u16;
         const CONFIGURATION_VALUE_U16: u16 = CONFIGURATION_VALUE as u16;
 
         // If the request has a data state, we must read it.
         let data = if req.length > 0 {
-            let size = match self.control.data_out(self.control_buf).await {
-                Ok(size) => size,
+            match self.control.data_out(self.control_buf).await {
+                Ok(data) => data,
                 Err(_) => {
                     warn!("usb: failed to read CONTROL OUT data stage.");
                     return;
                 }
-            };
-            &self.control_buf[0..size]
+            }
         } else {
             &[]
         };
@@ -315,10 +305,11 @@ impl<'d, D: Driver<'d>> UsbDevice<'d, D> {
             descriptor_type::CONFIGURATION => self.control.accept_in(self.config_descriptor).await,
             descriptor_type::STRING => {
                 if index == 0 {
-                    self.control_in_accept_writer(req, |w| {
-                        w.write(descriptor_type::STRING, &lang_id::ENGLISH_US.to_le_bytes())
-                    })
-                    .await
+                    self.control
+                        .accept_in_writer(req, |w| {
+                            w.write(descriptor_type::STRING, &lang_id::ENGLISH_US.to_le_bytes());
+                        })
+                        .await
                 } else {
                     let s = match index {
                         1 => self.config.manufacturer,
@@ -333,7 +324,7 @@ impl<'d, D: Driver<'d>> UsbDevice<'d, D> {
                     };
 
                     if let Some(s) = s {
-                        self.control_in_accept_writer(req, |w| w.string(s)).await;
+                        self.control.accept_in_writer(req, |w| w.string(s)).await;
                     } else {
                         self.control.reject()
                     }
@@ -343,3 +334,81 @@ impl<'d, D: Driver<'d>> UsbDevice<'d, D> {
         }
     }
 }
+
+struct ControlPipe<C: driver::ControlPipe> {
+    control: C,
+    request: Option<Request>,
+}
+
+impl<C: driver::ControlPipe> ControlPipe<C> {
+    async fn setup(&mut self) -> Request {
+        assert!(self.request.is_none());
+        let req = self.control.setup().await;
+        self.request = Some(req);
+        req
+    }
+
+    async fn data_out<'a>(&mut self, buf: &'a mut [u8]) -> Result<&'a [u8], ReadError> {
+        let req = self.request.unwrap();
+        assert_eq!(req.direction, UsbDirection::Out);
+        assert!(req.length > 0);
+        let req_length = usize::from(req.length);
+
+        let max_packet_size = self.control.max_packet_size();
+        let mut total = 0;
+
+        for chunk in buf.chunks_mut(max_packet_size) {
+            let size = self.control.data_out(chunk).await?;
+            total += size;
+            if size < max_packet_size || total == req_length {
+                break;
+            }
+        }
+
+        Ok(&buf[0..total])
+    }
+
+    async fn accept_in(&mut self, buf: &[u8]) -> () {
+        #[cfg(feature = "defmt")]
+        debug!("control in accept {:x}", buf);
+        #[cfg(not(feature = "defmt"))]
+        debug!("control in accept {:x?}", buf);
+        let req = unwrap!(self.request);
+        assert!(req.direction == UsbDirection::In);
+
+        let req_len = usize::from(req.length);
+        let len = buf.len().min(req_len);
+        let max_packet_size = self.control.max_packet_size();
+        let need_zlp = len != req_len && (len % usize::from(max_packet_size)) == 0;
+
+        let mut chunks = buf[0..len]
+            .chunks(max_packet_size)
+            .chain(need_zlp.then(|| -> &[u8] { &[] }));
+
+        while let Some(chunk) = chunks.next() {
+            self.control.data_in(chunk, chunks.size_hint().0 == 0).await;
+        }
+
+        self.request = None;
+    }
+
+    async fn accept_in_writer(&mut self, req: Request, f: impl FnOnce(&mut DescriptorWriter)) {
+        let mut buf = [0; 256];
+        let mut w = DescriptorWriter::new(&mut buf);
+        f(&mut w);
+        let pos = w.position().min(usize::from(req.length));
+        self.accept_in(&buf[..pos]).await;
+    }
+
+    fn accept(&mut self) {
+        assert!(self.request.is_some());
+        self.control.accept();
+        self.request = None;
+    }
+
+    fn reject(&mut self) {
+        assert!(self.request.is_some());
+        self.control.reject();
+        self.request = None;
+    }
+}