From 78a85cdb5b4cdde5da61d90513be46be941acd8d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?King=20K=C3=A9vin?= <kingkevin@cuvoodoo.info>
Date: Sun, 24 Jul 2022 15:51:30 +0200
Subject: [PATCH] dfu: use task to handle download

---
 examples/device/dfu_freertos/src/main.c | 106 +++++++++++++++++-------
 1 file changed, 76 insertions(+), 30 deletions(-)

diff --git a/examples/device/dfu_freertos/src/main.c b/examples/device/dfu_freertos/src/main.c
index 733c447f3..3bfb85b27 100644
--- a/examples/device/dfu_freertos/src/main.c
+++ b/examples/device/dfu_freertos/src/main.c
@@ -31,7 +31,6 @@
 #include "bsp/board.h"
 #include "tusb.h"
 
-
 #if TU_CHECK_MCU(OPT_MCU_ESP32S2, OPT_MCU_ESP32S3)
   // ESP-IDF need "freertos/" prefix in include path.
   // CFG_TUSB_OS_INC_PATH should be defined accordingly.
@@ -88,6 +87,54 @@ void usb_device_task(void* param);
 // used to download image onto OTA partition
 esp_ota_handle_t ota_handle = 0;
 
+// download blocks to program
+static uint8_t dl_block[CFG_TUD_DFU_XFER_BUFSIZE] = {0};
+static uint16_t dl_block_len = 0;
+static uint16_t dl_block_num = 0;
+
+// task taking care of flashing during download
+// note: esp_ota_begin erase flash, which takes too long to be in tud_dfu_download_cb, causing a USB timeout, thus it's in a task
+static void dl_task(void* arg)
+{
+  esp_err_t rc;
+  while (true) {
+    if (dl_block_len) { // we received a new block
+      // get handle for OTA update
+      if (0 == ota_handle) {
+        const esp_partition_t *ota_part = esp_ota_get_next_update_partition(NULL);
+        if (NULL == ota_part) {
+          ESP_LOGE(TAG, "OTA not found");
+          tud_dfu_finish_flashing(DFU_STATUS_ERR_PROG);
+          return;
+        }
+        rc = esp_ota_begin(ota_part, OTA_SIZE_UNKNOWN, &ota_handle);
+        if (ESP_OK != rc) {
+          ESP_LOGE(TAG, "init OTA failed");
+          esp_ota_abort(ota_handle);
+          ota_handle = 0;
+          tud_dfu_finish_flashing(DFU_STATUS_ERR_ERASE);
+          return;
+        }
+      }
+
+      // write data to partition
+      rc = esp_ota_write_with_offset(ota_handle, dl_block, dl_block_len, dl_block_num * CFG_TUD_DFU_XFER_BUFSIZE);
+      if (ESP_OK != rc) {
+        ESP_LOGE(TAG, "writing OTA failed");
+        esp_ota_abort(ota_handle);
+        ota_handle = 0;
+        tud_dfu_finish_flashing(DFU_STATUS_ERR_WRITE);
+        return;
+      }
+
+      dl_block_len = 0; // ready for next block
+      tud_dfu_finish_flashing(DFU_STATUS_OK); // flashing op for download complete without error
+    } else { // no new block
+      vTaskDelay(1); // allow other tasks to run (and watchdog reset)
+    }
+  }
+}
+
 //--------------------------------------------------------------------+
 // Main
 //--------------------------------------------------------------------+
@@ -118,6 +165,9 @@ int main(void)
   vTaskStartScheduler();
 #endif
 
+  // create task to handle download progress
+  xTaskCreate(dl_task, "dl_task", 2 * 1024, NULL, 8, NULL);
+
   ESP_LOGI(TAG, "DFU mode");
 
   return 0;
@@ -196,46 +246,42 @@ uint32_t tud_dfu_get_timeout_cb(uint8_t alt, uint8_t state)
 void tud_dfu_download_cb(uint8_t alt, uint16_t block_num, uint8_t const* data, uint16_t length)
 {
   esp_err_t rc;
+  static bool dl_started = false;
+
   ESP_LOGD(TAG, "download, alt=%u block=%u", alt, block_num);
 
-  if (alt > 0)
-  {
+  if (alt > 0) {
     ESP_LOGW(TAG, "download to invalid alt %u", alt);
     tud_dfu_finish_flashing(DFU_STATUS_ERR_ADDRESS);
     return;
   }
 
-  // get handle for OTA update
-  if (0 == ota_handle) {
-    const esp_partition_t *ota_part = esp_ota_get_next_update_partition(NULL);
-    if (NULL == ota_part) {
-      ESP_LOGE(TAG, "OTA not found");
-      tud_dfu_finish_flashing(DFU_STATUS_ERR_PROG);
-      return;
-    }
-    rc = esp_ota_begin(ota_part, OTA_SIZE_UNKNOWN, &ota_handle);
-    if (ESP_OK != rc) {
-      ESP_LOGE(TAG, "init OTA failed");
-      esp_ota_abort(ota_handle);
+  if (0 == length) { // there is nothing to programm
+    // finish flashing
+    if (ota_handle) {
+      rc = esp_ota_end(ota_handle);
       ota_handle = 0;
-      tud_dfu_finish_flashing(DFU_STATUS_ERR_PROG);
-      return;
+      if (ESP_OK != rc) {
+        ESP_LOGE(TAG, "close OTA failed");
+        tud_dfu_finish_flashing(DFU_STATUS_ERR_PROG);
+        return;
+      }
     }
-  }
-  xTimerChangePeriod(blinky_tm, pdMS_TO_TICKS(BLINK_DOWNLOAD), 0);
-
-  // write data to partition
-  rc = esp_ota_write_with_offset(ota_handle, data, length, block_num * CFG_TUD_DFU_XFER_BUFSIZE);
-  if (ESP_OK != rc) {
-    ESP_LOGE(TAG, "writing OTA failed");
-    esp_ota_abort(ota_handle);
-    ota_handle = 0;
+    dl_started = false;
+    tud_dfu_finish_flashing(DFU_STATUS_OK); // flashing op for download complete without error
+  } else if (length > CFG_TUD_DFU_XFER_BUFSIZE) { // more data than we can handle
     tud_dfu_finish_flashing(DFU_STATUS_ERR_PROG);
-    return;
+  } else if (dl_block_len) { // there is already a block to programm
+    tud_dfu_finish_flashing(DFU_STATUS_ERR_STALLEDPKT);
+  } else { // all is fine to programm
+    if (!dl_started) {
+      xTimerChangePeriod(blinky_tm, pdMS_TO_TICKS(BLINK_DOWNLOAD), 0);
+      dl_started = true;
+    }
+    memcpy(dl_block, data, length); // copy data for dl_task
+    dl_block_num = block_num; // remember offset
+    dl_block_len = length; // notitfy dl_task block is ready
   }
-
-  // flashing op for download complete without error
-  tud_dfu_finish_flashing(DFU_STATUS_OK);
 }
 
 // Invoked when download process is complete, received DFU_DNLOAD (wLength=0) following by DFU_GETSTATUS (state=Manifest)