From 2d5162dc3c375332657264b39f13e4031e7a8765 Mon Sep 17 00:00:00 2001
From: Angus Gratton <angus@espressif.com>
Date: Fri, 3 Feb 2017 10:07:30 +1100
Subject: [PATCH] OTA: Always clean up OTA handle regardless of esp_ota_end()
 result

As reported on forum: http://esp32.com/viewtopic.php?f=14&t=1093
---
 components/app_update/esp_ota_ops.c         | 70 +++++++++++----------
 components/app_update/include/esp_ota_ops.h | 13 ++--
 2 files changed, 47 insertions(+), 36 deletions(-)

diff --git a/components/app_update/esp_ota_ops.c b/components/app_update/esp_ota_ops.c
index 233875389..c702a203e 100644
--- a/components/app_update/esp_ota_ops.c
+++ b/components/app_update/esp_ota_ops.c
@@ -182,40 +182,10 @@ esp_err_t esp_ota_end(esp_ota_handle_t handle)
 {
     ota_ops_entry_t *it;
     size_t image_size;
-    esp_err_t __attribute__((unused)) ret;
+    esp_err_t ret = ESP_OK;
 
     for (it = LIST_FIRST(&s_ota_ops_entries_head); it != NULL; it = LIST_NEXT(it, entries)) {
         if (it->handle == handle) {
-            // an ota handle need to be ended after erased and wrote data in it
-            if ((it->erased_size == 0) || (it->wrote_size == 0)) {
-                return ESP_ERR_INVALID_ARG;
-            }
-
-#ifdef CONFIG_FLASH_ENCRYPTION_ENABLED
-            if (it->partial_bytes > 0 && esp_flash_encryption_enabled()) {
-                /* Write out last 16 bytes, if necessary */
-                ret = esp_partition_write(&it->part, it->wrote_size, it->partial_data, 16);
-                if (ret != ESP_OK) {
-                    return ret;
-                }
-                it->wrote_size += 16;
-                it->partial_bytes = 0;
-            }
-#endif
-
-            if (esp_image_basic_verify(it->part.address, true, &image_size) != ESP_OK) {
-                return ESP_ERR_OTA_VALIDATE_FAILED;
-            }
-
-#ifdef CONFIG_SECURE_BOOT_ENABLED
-            esp_err_t ret;
-            ret = esp_secure_boot_verify_signature(it->part.address, image_size);
-            if (ret != ESP_OK) {
-                return ESP_ERR_OTA_VALIDATE_FAILED;
-            }
-#endif
-
-            LIST_REMOVE(it, entries);
             break;
         }
     }
@@ -224,8 +194,44 @@ esp_err_t esp_ota_end(esp_ota_handle_t handle)
         return ESP_ERR_NOT_FOUND;
     }
 
+    /* 'it' holds the ota_ops_entry_t for 'handle' */
+
+    // esp_ota_end() is only valid if some data was written to this handle
+    if ((it->erased_size == 0) || (it->wrote_size == 0)) {
+        ret = ESP_ERR_INVALID_ARG;
+        goto cleanup;
+    }
+
+#ifdef CONFIG_FLASH_ENCRYPTION_ENABLED
+    if (it->partial_bytes > 0 && esp_flash_encryption_enabled()) {
+        /* Write out last 16 bytes, if necessary */
+        ret = esp_partition_write(&it->part, it->wrote_size, it->partial_data, 16);
+        if (ret != ESP_OK) {
+            ret = ESP_ERR_INVALID_STATE;
+            goto cleanup;
+        }
+        it->wrote_size += 16;
+        it->partial_bytes = 0;
+    }
+#endif
+
+    if (esp_image_basic_verify(it->part.address, true, &image_size) != ESP_OK) {
+        ret = ESP_ERR_OTA_VALIDATE_FAILED;
+        goto cleanup;
+    }
+
+#ifdef CONFIG_SECURE_BOOT_ENABLED
+    ret = esp_secure_boot_verify_signature(it->part.address, image_size);
+    if (ret != ESP_OK) {
+        ret = ESP_ERR_OTA_VALIDATE_FAILED;
+        goto cleanup;
+    }
+#endif
+
+ cleanup:
+    LIST_REMOVE(it, entries);
     free(it);
-    return ESP_OK;
+    return ret;
 }
 
 static uint32_t ota_select_crc(const ota_select *s)
diff --git a/components/app_update/include/esp_ota_ops.h b/components/app_update/include/esp_ota_ops.h
index 846aa2b2c..fe3307763 100755
--- a/components/app_update/include/esp_ota_ops.h
+++ b/components/app_update/include/esp_ota_ops.h
@@ -73,11 +73,16 @@ esp_err_t esp_ota_write(esp_ota_handle_t handle, const void* data, size_t size);
 /**
  * @brief   Finish the update and validate written data
  *
- * @param   handle  Handle obtained from esp_ota_begin 
+ * @param   handle  Handle obtained from esp_ota_begin.
  *
- * @return: 
- *    - ESP_OK: if validate ota image pass
- *    - ESP_ERR_OTA_VALIDATE_FAILED: validate the ota image is invalid
+ * @note After calling esp_ota_end(), the handle is no longer valid and any memory associated with it is freed (regardless of result).
+ *
+ * @return:
+ *    - ESP_OK: Newly written OTA app image is valid.
+ *    - ESP_ERR_NOT_FOUND: OTA handle was not found.
+ *    - ESP_ERR_INVALID_ARG: Handle was never written to.
+ *    - ESP_ERR_OTA_VALIDATE_FAILED: OTA image is invalid (either not a valid app image, or - if secure boot is enabled - signature failed to verify.)
+ *    - ESP_ERR_INVALID_STATE: If flash encryption is enabled, this result indicates an internal error writing the final encrypted bytes to flash.
  */
 esp_err_t esp_ota_end(esp_ota_handle_t handle);