From f6d11dfba56b2b04868e87a14d10395e1916306d Mon Sep 17 00:00:00 2001
From: Dario Nieuwenhuis <dirbaio@dirbaio.net>
Date: Wed, 6 Apr 2022 03:02:13 +0200
Subject: [PATCH] usb: fix slow enumeration with EP0 max_packet_size of 8 or
 16.

---
 embassy-usb/src/control.rs |  2 +-
 embassy-usb/src/lib.rs     | 15 ++++++++++++++-
 2 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/embassy-usb/src/control.rs b/embassy-usb/src/control.rs
index da48dccaf..b15ba4463 100644
--- a/embassy-usb/src/control.rs
+++ b/embassy-usb/src/control.rs
@@ -198,7 +198,7 @@ pub trait ControlHandler {
 #[derive(Debug)]
 #[cfg_attr(feature = "defmt", derive(defmt::Format))]
 pub(crate) struct DataInStage {
-    length: usize,
+    pub(crate) length: usize,
 }
 
 /// Typestate representing a ControlPipe in the DATA OUT stage
diff --git a/embassy-usb/src/lib.rs b/embassy-usb/src/lib.rs
index cf8d12539..f833a86d8 100644
--- a/embassy-usb/src/lib.rs
+++ b/embassy-usb/src/lib.rs
@@ -222,7 +222,20 @@ impl<'d, D: Driver<'d>> UsbDevice<'d, D> {
         }
     }
 
-    async fn handle_control_in(&mut self, req: Request, stage: DataInStage) {
+    async fn handle_control_in(&mut self, req: Request, mut stage: DataInStage) {
+        // If we don't have an address yet, respond with max 1 packet.
+        // The host doesn't know our EP0 max packet size yet, and might assume
+        // a full-length packet is a short packet, thinking we're done sending data.
+        // See https://github.com/hathach/tinyusb/issues/184
+        const DEVICE_DESCRIPTOR_LEN: u8 = 18;
+        if self.pending_address == 0
+            && self.config.max_packet_size_0 < DEVICE_DESCRIPTOR_LEN
+            && (self.config.max_packet_size_0 as usize) < stage.length
+        {
+            trace!("received control req while not addressed: capping response to 1 packet.");
+            stage.length = self.config.max_packet_size_0 as _;
+        }
+
         match (req.request_type, req.recipient) {
             (RequestType::Standard, Recipient::Device) => match req.request {
                 Request::GET_STATUS => {