From 2b44bc47df46557415869aaa1d1b1c8f23275bd0 Mon Sep 17 00:00:00 2001
From: Hrudaynath Dhabe <hrudaynath.dhabe@espressif.com>
Date: Thu, 15 Aug 2019 21:10:00 +0800
Subject: [PATCH] esp_https_ota: Added error checking functionalities.

Current implimentation of esp_http_ota does not perform any error-checking in the data writing phase calls `esp_ota_get_next_update_partition()` irrespetive of the received state of the image. A few additional error checking mechanism have now been added inside the esp_https_ota which returns the control in case an invalid header is received and a wrapper to the function `esp_http_client_is_complete_data_received()` of `esp_http_client` has been added.
---
 .../esp_https_ota/include/esp_https_ota.h     | 13 ++++++++++++
 components/esp_https_ota/src/esp_https_ota.c  | 20 ++++++++++++++-----
 .../main/advanced_https_ota_example.c         |  5 +++++
 .../main/native_ota_example.c                 |  5 +++++
 4 files changed, 38 insertions(+), 5 deletions(-)

diff --git a/components/esp_https_ota/include/esp_https_ota.h b/components/esp_https_ota/include/esp_https_ota.h
index 0c8fd0693..7feba8f7e 100644
--- a/components/esp_https_ota/include/esp_https_ota.h
+++ b/components/esp_https_ota/include/esp_https_ota.h
@@ -110,6 +110,19 @@ esp_err_t esp_https_ota_begin(esp_https_ota_config_t *ota_config, esp_https_ota_
  */
 esp_err_t esp_https_ota_perform(esp_https_ota_handle_t https_ota_handle);
 
+/**
+ * @brief   Checks if complete data was received or not
+ *
+ * @note    This API can be called just before esp_https_ota_finish() to validate if the complete image was indeed received.
+ *
+ * @param[in]   https_ota_handle pointer to esp_https_ota_handle_t structure
+ *
+ * @return
+ *    - false
+ *    - true
+ */
+bool esp_https_ota_is_complete_data_received(esp_https_ota_handle_t https_ota_handle);
+
 /**
  * @brief    Clean-up HTTPS OTA Firmware upgrade and close HTTPS connection
  *
diff --git a/components/esp_https_ota/src/esp_https_ota.c b/components/esp_https_ota/src/esp_https_ota.c
index baba0bbaa..d2ca50103 100644
--- a/components/esp_https_ota/src/esp_https_ota.c
+++ b/components/esp_https_ota/src/esp_https_ota.c
@@ -85,17 +85,21 @@ static esp_err_t _http_handle_response_code(esp_http_client_handle_t http_client
 static esp_err_t _http_connect(esp_http_client_handle_t http_client)
 {
     esp_err_t err = ESP_FAIL;
-    int status_code;
+    int status_code, header_ret;
     do {
         err = esp_http_client_open(http_client, 0);
         if (err != ESP_OK) {
             ESP_LOGE(TAG, "Failed to open HTTP connection: %s", esp_err_to_name(err));
             return err;
         }
-        esp_http_client_fetch_headers(http_client);
+        header_ret = esp_http_client_fetch_headers(http_client);
+        if (header_ret < 0) {
+            return header_ret;
+        }
         status_code = esp_http_client_get_status_code(http_client);
-        if (_http_handle_response_code(http_client, status_code) != ESP_OK) {
-            return ESP_FAIL;
+        err = _http_handle_response_code(http_client, status_code);
+        if (err != ESP_OK) {
+            return err;
         }
     } while (process_again(status_code));
     return err;
@@ -276,6 +280,12 @@ esp_err_t esp_https_ota_perform(esp_https_ota_handle_t https_ota_handle)
     return ESP_OK;
 }
 
+bool esp_https_ota_is_complete_data_received(esp_https_ota_handle_t https_ota_handle)
+{
+    esp_https_ota_t *handle = (esp_https_ota_t *)https_ota_handle;
+    return esp_http_client_is_complete_data_received(handle->http_client);
+}
+
 esp_err_t esp_https_ota_finish(esp_https_ota_handle_t https_ota_handle)
 {
     esp_https_ota_t *handle = (esp_https_ota_t *)https_ota_handle;
@@ -361,4 +371,4 @@ esp_err_t esp_https_ota(const esp_http_client_config_t *config)
         return ota_finish_err;
     }
     return ESP_OK;
-}
\ No newline at end of file
+}
diff --git a/examples/system/ota/advanced_https_ota/main/advanced_https_ota_example.c b/examples/system/ota/advanced_https_ota/main/advanced_https_ota_example.c
index 43fd16243..241495be3 100644
--- a/examples/system/ota/advanced_https_ota/main/advanced_https_ota_example.c
+++ b/examples/system/ota/advanced_https_ota/main/advanced_https_ota_example.c
@@ -142,6 +142,11 @@ void advanced_ota_example_task(void * pvParameter)
         ESP_LOGD(TAG, "Image bytes read: %d", esp_https_ota_get_image_len_read(https_ota_handle));
     }
 
+    if (esp_https_ota_is_complete_data_received(&https_ota_handle) != true) {
+        // the OTA image was not completely received and user can customise the response to this situation.
+        ESP_LOGE(TAG, "Complete data was not received.");
+    }
+
 ota_end:
     ota_finish_err = esp_https_ota_finish(https_ota_handle);
     if ((err == ESP_OK) && (ota_finish_err == ESP_OK)) {
diff --git a/examples/system/ota/native_ota_example/main/native_ota_example.c b/examples/system/ota/native_ota_example/main/native_ota_example.c
index 854e77fac..b3b34d57d 100644
--- a/examples/system/ota/native_ota_example/main/native_ota_example.c
+++ b/examples/system/ota/native_ota_example/main/native_ota_example.c
@@ -243,6 +243,11 @@ static void ota_example_task(void *pvParameter)
         }
     }
     ESP_LOGI(TAG, "Total Write binary data length : %d", binary_file_length);
+    if (esp_http_client_is_complete_data_received(client) != true) {
+        ESP_LOGE(TAG, "Error in receiving complete file");
+        http_cleanup(client);
+        task_fatal_error();
+    }
 
     if (esp_ota_end(update_handle) != ESP_OK) {
         ESP_LOGE(TAG, "esp_ota_end failed!");