1 files changed, 33 insertions, 16 deletions

diff --git a/src/class/dfu/dfu_device.c b/src/class/dfu/dfu_device.c
index 0d2b63b57..ee57621b8 100644
--- a/src/class/dfu/dfu_device.c
+++ b/src/class/dfu/dfu_device.c
@@ -50,21 +50,19 @@
 typedef struct {
   uint8_t attrs;
   uint8_t alt;
+  uint8_t state;
+  uint8_t status;
 
-  dfu_state_t state;
-  dfu_status_t status;
-
-  bool flashing_in_progress;
+  bool     flashing_in_progress;
   uint16_t block;
   uint16_t length;
 } dfu_state_ctx_t;
 
-// Only a single dfu state is allowed
 static dfu_state_ctx_t _dfu_ctx;
 
-CFG_TUD_MEM_SECTION static struct {
-  TUD_EPBUF_DEF(transfer_buf, CFG_TUD_DFU_XFER_BUFSIZE);
-} _dfu_epbuf;
+#if CFG_TUD_DFU_XFER_BUFSIZE > CFG_TUD_ENDPOINT0_BUFSIZE
+TU_ATTR_ALIGNED(4) uint8_t _transfer_buf[CFG_TUD_DFU_XFER_BUFSIZE];
+#endif
 
 static void reset_state(void) {
   _dfu_ctx.state = DFU_IDLE;
@@ -72,6 +70,15 @@ static void reset_state(void) {
   _dfu_ctx.flashing_in_progress = false;
 }
 
+static inline uint8_t* get_xfer_buffer(void) {
+  // Use EP0 buffer if it is large enough, otherwise use dedicated buffer
+  #if CFG_TUD_DFU_XFER_BUFSIZE > CFG_TUD_ENDPOINT0_BUFSIZE
+  return _transfer_buf;
+  #else
+  return usbd_get_ctrl_buf();
+  #endif
+}
+
 static bool reply_getstatus(uint8_t rhport, const tusb_control_request_t* request, dfu_state_t state, dfu_status_t status, uint32_t timeout);
 static bool process_download_get_status(uint8_t rhport, uint8_t stage, const tusb_control_request_t* request);
 static bool process_manifest_get_status(uint8_t rhport, uint8_t stage, const tusb_control_request_t* request);
@@ -253,6 +260,8 @@ bool dfu_moded_control_xfer_cb(uint8_t rhport, uint8_t stage, const tusb_control
           tud_control_status(rhport, request);
         } else if (stage == CONTROL_STAGE_ACK) {
           tud_dfu_detach_cb();
+        } else {
+          // nothing to do
         }
         break;
 
@@ -275,6 +284,8 @@ bool dfu_moded_control_xfer_cb(uint8_t rhport, uint8_t stage, const tusb_control
           tud_control_status(rhport, request);
         } else if (stage == CONTROL_STAGE_ACK) {
           tud_dfu_abort_cb(_dfu_ctx.alt);
+        } else {
+          // nothing to do
         }
         break;
 
@@ -283,10 +294,10 @@ bool dfu_moded_control_xfer_cb(uint8_t rhport, uint8_t stage, const tusb_control
           TU_VERIFY(_dfu_ctx.attrs & DFU_ATTR_CAN_UPLOAD);
           TU_VERIFY(request->wLength <= CFG_TUD_DFU_XFER_BUFSIZE);
 
-          const uint16_t xfer_len = tud_dfu_upload_cb(_dfu_ctx.alt, request->wValue, _dfu_epbuf.transfer_buf,
+          const uint16_t xfer_len = tud_dfu_upload_cb(_dfu_ctx.alt, request->wValue, get_xfer_buffer(),
                                                       request->wLength);
 
-          return tud_control_xfer(rhport, request, _dfu_epbuf.transfer_buf, xfer_len);
+          return tud_control_xfer(rhport, request, get_xfer_buffer(), xfer_len);
         }
         break;
 
@@ -303,10 +314,10 @@ bool dfu_moded_control_xfer_cb(uint8_t rhport, uint8_t stage, const tusb_control
           _dfu_ctx.block = request->wValue;
           _dfu_ctx.length = request->wLength;
 
-          if (request->wLength) {
+          if (request->wLength > 0) {
             // Download with payload -> transition to DOWNLOAD SYNC
             _dfu_ctx.state = DFU_DNLOAD_SYNC;
-            return tud_control_xfer(rhport, request, _dfu_epbuf.transfer_buf, request->wLength);
+            return tud_control_xfer(rhport, request, get_xfer_buffer(), request->wLength);
           } else {
             // Download is complete -> transition to MANIFEST SYNC
             _dfu_ctx.state = DFU_MANIFEST_SYNC;
@@ -327,7 +338,7 @@ bool dfu_moded_control_xfer_cb(uint8_t rhport, uint8_t stage, const tusb_control
 
           default:
             if (stage == CONTROL_STAGE_SETUP) {
-              return reply_getstatus(rhport, request, _dfu_ctx.state, _dfu_ctx.status, 0);
+              return reply_getstatus(rhport, request, (dfu_state_t) _dfu_ctx.state, (dfu_status_t) _dfu_ctx.status, 0);
             }
             break;
         }
@@ -352,6 +363,8 @@ void tud_dfu_finish_flashing(uint8_t status) {
       _dfu_ctx.state = (_dfu_ctx.attrs & DFU_ATTR_MANIFESTATION_TOLERANT)
                          ? DFU_MANIFEST_SYNC
                          : DFU_MANIFEST_WAIT_RESET;
+    } else {
+      // nothing to do
     }
   } else {
     // failed while flashing, move to dfuError
@@ -374,14 +387,16 @@ static bool process_download_get_status(uint8_t rhport, uint8_t stage, const tus
       timeout = 0;
     }
 
-    return reply_getstatus(rhport, request, next_state, _dfu_ctx.status, timeout);
+    return reply_getstatus(rhport, request, next_state, (dfu_status_t) _dfu_ctx.status, timeout);
   } else if (stage == CONTROL_STAGE_ACK) {
     if (_dfu_ctx.flashing_in_progress) {
       _dfu_ctx.state = DFU_DNBUSY;
-      tud_dfu_download_cb(_dfu_ctx.alt, _dfu_ctx.block, _dfu_epbuf.transfer_buf, _dfu_ctx.length);
+      tud_dfu_download_cb(_dfu_ctx.alt, _dfu_ctx.block, get_xfer_buffer(), _dfu_ctx.length);
     } else {
       _dfu_ctx.state = DFU_DNLOAD_IDLE;
     }
+  } else {
+    // nothing to do
   }
 
   return true;
@@ -401,7 +416,7 @@ static bool process_manifest_get_status(uint8_t rhport, uint8_t stage, const tus
       timeout = 0;
     }
 
-    return reply_getstatus(rhport, request, next_state, _dfu_ctx.status, timeout);
+    return reply_getstatus(rhport, request, next_state, (dfu_status_t) _dfu_ctx.status, timeout);
   } else if (stage == CONTROL_STAGE_ACK) {
     if (_dfu_ctx.flashing_in_progress) {
       _dfu_ctx.state = DFU_MANIFEST;
@@ -409,6 +424,8 @@ static bool process_manifest_get_status(uint8_t rhport, uint8_t stage, const tus
     } else {
       _dfu_ctx.state = DFU_IDLE;
     }
+  } else {
+    // nothing to do
   }
 
   return true;