dfu: use ep0 buffer if it is large enough

Signed-off-by: HiFiPhile <[email protected]>

1 files changed, 15 insertions, 4 deletions

diff --git a/src/class/dfu/dfu_device.c b/src/class/dfu/dfu_device.c
index d3cc53918..6c42ca4e6 100644
--- a/src/class/dfu/dfu_device.c
+++ b/src/class/dfu/dfu_device.c
@@ -60,7 +60,9 @@ typedef struct {
 
 static dfu_state_ctx_t _dfu_ctx;
 
+#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;
@@ -68,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_control_get_buffer();
+  #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);
@@ -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, _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, _transfer_buf, xfer_len);
+          return tud_control_xfer(rhport, request, get_xfer_buffer(), xfer_len);
         }
         break;
 
@@ -306,7 +317,7 @@ bool dfu_moded_control_xfer_cb(uint8_t rhport, uint8_t stage, const tusb_control
           if (request->wLength > 0) {
             // Download with payload -> transition to DOWNLOAD SYNC
             _dfu_ctx.state = DFU_DNLOAD_SYNC;
-            return tud_control_xfer(rhport, request, _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;
@@ -380,7 +391,7 @@ static bool process_download_get_status(uint8_t rhport, uint8_t stage, const tus
   } 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, _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;
     }