Allow cancellation of download, dont block the device

src/activities/settings/OtaUpdateActivity.cpp

@@ -140,20 +140,15 @@ void OtaUpdateActivity::loop() {
   // TODO @ngxson : refactor this logic later
   if (updater.getRender()) {
     requestUpdate();
+    updater.clearRender();
   }
 
   if (state == WAITING_CONFIRMATION) {
     if (mappedInput.wasPressed(MappedInputManager::Button::Confirm)) {
       LOG_DBG("OTA", "New update available, starting download...");
-      {
+      const auto beginResult = updater.beginInstallUpdate();
-        RenderLock lock(*this);
+      if (beginResult != OtaUpdater::UPDATE_IN_PROGRESS) {
-        state = UPDATE_IN_PROGRESS;
+        LOG_DBG("OTA", "Update begin failed: %d", beginResult);
-      }
-      requestUpdateAndWait();
-      const auto res = updater.installUpdate();
-
-      if (res != OtaUpdater::OK) {
-        LOG_DBG("OTA", "Update failed: %d", res);
         {
           RenderLock lock(*this);
           state = FAILED;
@@ -164,9 +159,10 @@ void OtaUpdateActivity::loop() {
 
       {
         RenderLock lock(*this);
-        state = FINISHED;
+        state = UPDATE_IN_PROGRESS;
       }
       requestUpdate();
+      return;
     }
 
     if (mappedInput.wasPressed(MappedInputManager::Button::Back)) {
@@ -176,6 +172,45 @@ void OtaUpdateActivity::loop() {
     return;
   }
 
+  if (state == UPDATE_IN_PROGRESS) {
+    if (mappedInput.wasPressed(MappedInputManager::Button::Back)) {
+      updater.cancelUpdate();
+      finish();
+      return;
+    }
+
+    const auto res = updater.performInstallUpdateStep();
+    if (res == OtaUpdater::UPDATE_IN_PROGRESS) {
+      if (updater.getRender()) {
+        requestUpdate();
+        updater.clearRender();
+      }
+      return;
+    }
+
+    if (res == OtaUpdater::OK) {
+      {
+        RenderLock lock(*this);
+        state = FINISHED;
+      }
+      requestUpdate();
+      return;
+    }
+
+    if (res == OtaUpdater::UPDATE_CANCELLED) {
+      finish();
+      return;
+    }
+
+    LOG_DBG("OTA", "Update failed: %d", res);
+    {
+      RenderLock lock(*this);
+      state = FAILED;
+    }
+    requestUpdate();
+    return;
+  }
+
   if (state == FAILED) {
     if (mappedInput.wasPressed(MappedInputManager::Button::Back)) {
       finish();

src/network/OtaUpdater.cpp

@@ -28,34 +28,25 @@ esp_err_t http_client_set_header_cb(esp_http_client_handle_t http_client) {
 }
 
 esp_err_t event_handler(esp_http_client_event_t* event) {
-  /* We do interested in only HTTP_EVENT_ON_DATA event only */
+  /* We are only interested in HTTP_EVENT_ON_DATA event */
   if (event->event_id != HTTP_EVENT_ON_DATA) return ESP_OK;
 
-  if (!esp_http_client_is_chunked_response(event->client)) {
+  if (event->data == nullptr || event->data_len == 0) {
-    int content_len = esp_http_client_get_content_length(event->client);
+    return ESP_OK;
-    int copy_len = 0;
-
-    if (local_buf == NULL) {
-      /* local_buf life span is tracked by caller checkForUpdate */
-      local_buf = static_cast<char*>(calloc(content_len + 1, sizeof(char)));
-      output_len = 0;
-      if (local_buf == NULL) {
-        LOG_ERR("OTA", "HTTP Client Out of Memory Failed, Allocation %d", content_len);
-        return ESP_ERR_NO_MEM;
-      }
-    }
-    copy_len = min(event->data_len, (content_len - output_len));
-    if (copy_len) {
-      memcpy(local_buf + output_len, event->data, copy_len);
-    }
-    output_len += copy_len;
-  } else {
-    /* Code might be hits here, It happened once (for version checking) but I need more logs to handle that */
-    int chunked_len;
-    esp_http_client_get_chunk_length(event->client, &chunked_len);
-    LOG_DBG("OTA", "esp_http_client_is_chunked_response failed, chunked_len: %d", chunked_len);
   }
 
+  const int newSize = output_len + event->data_len + 1;
+  char* newBuf = static_cast<char*>(realloc(local_buf, static_cast<size_t>(newSize)));
+  if (newBuf == nullptr) {
+    LOG_ERR("OTA", "HTTP Client Out of Memory Failed, Allocation %d", newSize);
+    return ESP_ERR_NO_MEM;
+  }
+
+  local_buf = newBuf;
+  memcpy(local_buf + output_len, event->data, event->data_len);
+  output_len += event->data_len;
+  local_buf[output_len] = '\0';
+
   return ESP_OK;
 } /* event_handler */
 } /* namespace */
@@ -68,6 +59,7 @@ OtaUpdater::OtaUpdaterError OtaUpdater::checkForUpdate() {
   esp_http_client_config_t client_config = {
       .url = latestReleaseUrl,
       .event_handler = event_handler,
+      .timeout_ms = 10000,
       /* Default HTTP client buffer size 512 byte only */
       .buffer_size = 8192,
       .buffer_size_tx = 8192,
@@ -199,28 +191,38 @@ bool OtaUpdater::isUpdateNewer() const {
 
 const std::string& OtaUpdater::getLatestVersion() const { return latestVersion; }
 
-OtaUpdater::OtaUpdaterError OtaUpdater::installUpdate() {
+void OtaUpdater::cleanupUpdate() {
+  if (otaHandle) {
+    esp_https_ota_finish(otaHandle);
+    otaHandle = nullptr;
+  }
+  cancelRequested = false;
+  esp_wifi_set_ps(WIFI_PS_MIN_MODEM);
+}
+
+void OtaUpdater::cancelUpdate() {
+  if (otaHandle) {
+    cleanupUpdate();
+  } else {
+    cancelRequested = true;
+  }
+}
+
+OtaUpdater::OtaUpdaterError OtaUpdater::beginInstallUpdate() {
   if (!isUpdateNewer()) {
     return UPDATE_OLDER_ERROR;
   }
 
-  esp_https_ota_handle_t ota_handle = NULL;
+  cleanupUpdate();
-  esp_err_t esp_err;
-  /* Signal for OtaUpdateActivity */
   render = false;
+  cancelRequested = false;
 
   esp_http_client_config_t client_config = {
       .url = otaUrl.c_str(),
-      .timeout_ms = 30000,
+      .timeout_ms = 10000,
-      /* Default HTTP client buffer size 512 byte only
-       * not sufficient to handle URL redirection cases or
-       * parsing of large HTTP headers.
-       */
       .max_redirection_count = 5,
       .buffer_size = 8192,
       .buffer_size_tx = 8192,
-      /* GitHub release assets redirect to objects.githubusercontent.com CDN.
-       * Without max_redirection_count, esp_https_ota downloads 0 bytes and stalls. */
       .crt_bundle_attach = esp_crt_bundle_attach,
       .keep_alive_enable = true,
   };
@@ -233,41 +235,73 @@ OtaUpdater::OtaUpdaterError OtaUpdater::installUpdate() {
   /* For better timing and connectivity, we disable power saving for WiFi */
   esp_wifi_set_ps(WIFI_PS_NONE);
 
-  esp_err = esp_https_ota_begin(&ota_config, &ota_handle);
+  esp_err_t esp_err = esp_https_ota_begin(&ota_config, &otaHandle);
   if (esp_err != ESP_OK) {
     LOG_DBG("OTA", "HTTP OTA Begin Failed: %s", esp_err_to_name(esp_err));
+    cleanupUpdate();
     return INTERNAL_UPDATE_ERROR;
   }
 
-  do {
+  return UPDATE_IN_PROGRESS;
-    esp_err = esp_https_ota_perform(ota_handle);
+}
-    processedSize = esp_https_ota_get_image_len_read(ota_handle);
+
-    /* Sent signal to  OtaUpdateActivity */
+OtaUpdater::OtaUpdaterError OtaUpdater::performInstallUpdateStep() {
-    render = true;
+  if (cancelRequested) {
-    delay(100);  // TODO: should we replace this with something better?
+    cleanupUpdate();
-  } while (esp_err == ESP_ERR_HTTPS_OTA_IN_PROGRESS);
+    return UPDATE_CANCELLED;
+  }
+
+  if (!otaHandle) {
+    return INTERNAL_UPDATE_ERROR;
+  }
+
+  esp_err_t esp_err = esp_https_ota_perform(otaHandle);
+  processedSize = esp_https_ota_get_image_len_read(otaHandle);
+  render = true;
+
+  if (esp_err == ESP_ERR_HTTPS_OTA_IN_PROGRESS) {
+    return UPDATE_IN_PROGRESS;
+  }
 
-  /* Return back to default power saving for WiFi in case of failing */
   esp_wifi_set_ps(WIFI_PS_MIN_MODEM);
 
   if (esp_err != ESP_OK) {
     LOG_ERR("OTA", "esp_https_ota_perform Failed: %s", esp_err_to_name(esp_err));
-    esp_https_ota_finish(ota_handle);
+    cleanupUpdate();
     return HTTP_ERROR;
   }
 
-  if (!esp_https_ota_is_complete_data_received(ota_handle)) {
+  if (!esp_https_ota_is_complete_data_received(otaHandle)) {
-    LOG_ERR("OTA", "esp_https_ota_is_complete_data_received Failed: %s", esp_err_to_name(esp_err));
+    LOG_ERR("OTA", "esp_https_ota_is_complete_data_received Failed");
-    esp_https_ota_finish(ota_handle);
+    cleanupUpdate();
     return INTERNAL_UPDATE_ERROR;
   }
 
-  esp_err = esp_https_ota_finish(ota_handle);
+  esp_err_t finish_err = esp_https_ota_finish(otaHandle);
-  if (esp_err != ESP_OK) {
+  otaHandle = nullptr;
-    LOG_ERR("OTA", "esp_https_ota_finish Failed: %s", esp_err_to_name(esp_err));
+  if (finish_err != ESP_OK) {
+    LOG_ERR("OTA", "esp_https_ota_finish Failed: %s", esp_err_to_name(finish_err));
+    cleanupUpdate();
     return INTERNAL_UPDATE_ERROR;
   }
 
   LOG_INF("OTA", "Update completed");
   return OK;
 }
+
+OtaUpdater::OtaUpdaterError OtaUpdater::installUpdate() {
+  const auto beginResult = beginInstallUpdate();
+  if (beginResult != UPDATE_IN_PROGRESS) {
+    return beginResult;
+  }
+
+  OtaUpdaterError result;
+  do {
+    result = performInstallUpdateStep();
+    if (result == UPDATE_IN_PROGRESS) {
+      delay(100);
+    }
+  } while (result == UPDATE_IN_PROGRESS);
+
+  return result;
+}

src/network/OtaUpdater.h

@@ -3,6 +3,8 @@
 #include <functional>
 #include <string>
 
+#include "esp_https_ota.h"
+
 class OtaUpdater {
   bool updateAvailable = false;
   std::string latestVersion;
@@ -11,6 +13,8 @@ class OtaUpdater {
   size_t processedSize = 0;
   size_t totalSize = 0;
   bool render = false;
+  esp_https_ota_handle_t otaHandle = nullptr;
+  bool cancelRequested = false;
 
  public:
   enum OtaUpdaterError {
@@ -21,6 +25,8 @@ class OtaUpdater {
     UPDATE_OLDER_ERROR,
     INTERNAL_UPDATE_ERROR,
     OOM_ERROR,
+    UPDATE_CANCELLED,
+    UPDATE_IN_PROGRESS,
   };
 
   size_t getOtaSize() const { return otaSize; }
@@ -30,10 +36,17 @@ class OtaUpdater {
   size_t getTotalSize() const { return totalSize; }
 
   bool getRender() const { return render; }
+  void clearRender() { render = false; }
+
+  bool isUpdateInProgress() const { return otaHandle != nullptr; }
 
   OtaUpdater() = default;
   bool isUpdateNewer() const;
   const std::string& getLatestVersion() const;
   OtaUpdaterError checkForUpdate();
+  OtaUpdaterError beginInstallUpdate();
+  OtaUpdaterError performInstallUpdateStep();
+  void cancelUpdate();
+  void cleanupUpdate();
   OtaUpdaterError installUpdate();
 };