Add bwPollTimeout set callback, postpone download callback after GETSTATUS

3 files changed, 18 insertions, 4 deletions

diff --git a/src/class/dfu/dfu_device.c b/src/class/dfu/dfu_device.c
index 834648c5a..34472bcdd 100644
--- a/src/class/dfu/dfu_device.c
+++ b/src/class/dfu/dfu_device.c
@@ -270,7 +270,6 @@ bool dfu_moded_control_xfer_cb(uint8_t rhport, uint8_t stage, tusb_control_reque
            && ((_dfu_state_ctx.attrs & DFU_FUNC_ATTR_CAN_DOWNLOAD_BITMASK) != 0)
            && (_dfu_state_ctx.state == DFU_DNLOAD_SYNC))
       {
-        dfu_req_dnload_reply(rhport, request);
         return true;
       }
     } // fallthrough
@@ -314,8 +313,16 @@ static void dfu_req_getstatus_reply(uint8_t rhport, tusb_control_request_t const
 {
   dfu_status_req_payload_t resp;
 
+  uint16_t timeout = 0;
   resp.bStatus = _dfu_state_ctx.status;
-  memset((uint8_t *)&resp.bwPollTimeout, 0x00, 3);
+  if(_dfu_state_ctx.state == DFU_DNBUSY && tud_dfu_set_timeout_cb)
+  {
+    timeout = tud_dfu_set_timeout_cb(_dfu_state_ctx.alt);
+    
+  }
+  resp.bwPollTimeout[0] = TU_U16_LOW(timeout);
+  resp.bwPollTimeout[1] = TU_U16_HIGH(timeout);
+  resp.bwPollTimeout[2] = 0;
   resp.bState = _dfu_state_ctx.state;
   resp.iString = 0;
 
@@ -433,6 +440,7 @@ static bool dfu_state_machine(uint8_t rhport, tusb_control_request_t const * req
           {
             _dfu_state_ctx.state = DFU_DNBUSY;
             dfu_req_getstatus_reply(rhport, request);
+            dfu_req_dnload_reply(rhport, request);
           } else {
             _dfu_state_ctx.state = DFU_DNLOAD_IDLE;
             dfu_req_getstatus_reply(rhport, request);
@@ -537,7 +545,8 @@ static bool dfu_state_machine(uint8_t rhport, tusb_control_request_t const * req
           {
             _dfu_state_ctx.state = DFU_MANIFEST;
             dfu_req_getstatus_reply(rhport, request);
-          } else {
+          } else 
+          {
             if ( tud_dfu_firmware_valid_check_cb(_dfu_state_ctx.alt) )
             {
               _dfu_state_ctx.state = DFU_IDLE;
diff --git a/src/class/dfu/dfu_device.h b/src/class/dfu/dfu_device.h
index f41b4b651..c09b9a823 100644
--- a/src/class/dfu/dfu_device.h
+++ b/src/class/dfu/dfu_device.h
@@ -48,6 +48,11 @@
 // alt is used as the partition number, in order to support multiple partitions like FLASH, EEPROM, etc.
 bool tud_dfu_firmware_valid_check_cb(uint8_t alt);
 
+// Invoked when a DFU_GETSTATUS request is received in DFU_DNBUSY state
+// Used to set the bwPollTimeout value, useful for slow Flash in order to make host wait longer 
+// alt is used as the partition number, in order to support multiple partitions like FLASH, EEPROM, etc.
+TU_ATTR_WEAK uint16_t tud_dfu_set_timeout_cb(uint8_t alt);
+
 // Invoked when a DFU_DNLOAD request is received
 // alt is used as the partition number, in order to support multiple partitions like FLASH, EEPROM, etc.
 // This callback takes the wBlockNum chunk of length length and provides it
diff --git a/src/device/usbd.h b/src/device/usbd.h
index fe1863706..fd3aee6d7 100644
--- a/src/device/usbd.h
+++ b/src/device/usbd.h
@@ -640,7 +640,7 @@ TU_ATTR_WEAK bool tud_vendor_control_xfer_cb(uint8_t rhport, uint8_t stage, tusb
 #define TUD_DFU_MODE_ALTS(_itfnum, ...) \
   TU_XSTRCAT(TUD_DFU_MODE_ALT_, CFG_TUD_DFU_ALT_COUNT)(_itfnum, __VA_ARGS__)
 
-// Interface number, detach timeout, transfer size, string index 1, [string index 2, string index n]
+// Interface number, attributes, detach timeout, transfer size, string index 0, [string index 1, string index n]
 #define TUD_DFU_MODE_DESCRIPTOR(_itfnum, _attr, _timeout, _xfer_size, _stridx, ...) \
   TUD_DFU_MODE_ALTS(_itfnum, _TUD_DFU_COMBINE(_stridx, __VA_ARGS__)) \
   TUD_DFU_MODE_FUNC(_attr, _timeout, _xfer_size)