ota: Better usage of App rollback feature

- Mark new app firmware as valid only if MQTT connection is successful.
- Report OTA success to RainMaker backend only after successful
  connection after booting into new firmware.
- If the new firmware cannot connect to MQTT within specified time
 (Default 90 seconds, configuratble), it will roll back into older
 firmware. Even if the timer logic fails for some reason, a reboot
 will also trigger the rollback.

Author: shahpiyushv

Diff for: components/esp_rainmaker/Kconfig.projbuild

@@ -199,6 +199,16 @@ menu "ESP RainMaker Config"
                 However, please ensure that your application has enough memory headroom to allow this,
                 else, the OTA may fail.
 
+        config ESP_RMAKER_OTA_ROLLBACK_WAIT_PERIOD
+            int "OTA Rollback Wait Period (Seconds)"
+            default 90
+            range 30 600
+            help
+                After an OTA Update, if CONFIG_BOOTLOADER_APP_ROLLBACK_ENABLE is set, then the firmware will wait for MQTT
+                connection to mark the new firmware as valid. However, if it is not able to do so within
+                this wait period (in seconds), the firmware will be marked as invalid and the older
+                firmware will be booted into.
+
     endmenu
 
     menu "ESP RainMaker Scheduling"

Diff for: components/esp_rainmaker/src/ota/esp_rmaker_ota.c

@@ -14,18 +14,24 @@
 
 #include <string.h>
 #include <freertos/FreeRTOS.h>
+#include <freertos/timers.h>
 #include <freertos/task.h>
+#include <esp_event.h>
 #include <esp_log.h>
 #include <esp_ota_ops.h>
 #include <esp_partition.h>
 #include <esp_https_ota.h>
 #include <esp_wifi_types.h>
 #include <esp_wifi.h>
+#include <nvs.h>
 #if CONFIG_BT_ENABLED
 #include <esp_bt.h>
 #endif /* CONFIG_BT_ENABLED */
 
 #include <esp_rmaker_utils.h>
+#include <esp_rmaker_common_events.h>
+
+#include "esp_rmaker_internal.h"
 #include "esp_rmaker_ota_internal.h"
 
 #if ESP_IDF_VERSION >= ESP_IDF_VERSION_VAL(4, 4, 0)
@@ -44,11 +50,12 @@
 
 #endif /* !IDF4.4 */
 static const char *TAG = "esp_rmaker_ota";
+static TimerHandle_t s_ota_rollback_timer;
 
 #define OTA_REBOOT_TIMER_SEC    10
 #define DEF_HTTP_TX_BUFFER_SIZE    1024
 #define DEF_HTTP_RX_BUFFER_SIZE    CONFIG_ESP_RMAKER_OTA_HTTP_RX_BUFFER_SIZE
-
+#define RMAKER_OTA_ROLLBACK_WAIT_PERIOD    CONFIG_ESP_RMAKER_OTA_ROLLBACK_WAIT_PERIOD
 extern const char esp_rmaker_ota_def_cert[] asm("_binary_rmaker_ota_server_crt_start");
 const char *ESP_RMAKER_OTA_DEFAULT_SERVER_CERT = esp_rmaker_ota_def_cert;
 char *esp_rmaker_ota_status_to_string(ota_status_t status)
@@ -71,6 +78,8 @@ char *esp_rmaker_ota_status_to_string(ota_status_t status)
 }
 esp_err_t esp_rmaker_ota_report_status(esp_rmaker_ota_handle_t ota_handle, ota_status_t status, char *additional_info)
 {
+    ESP_LOGI(TAG, "Reporting %s: %s", esp_rmaker_ota_status_to_string(status), additional_info);
+
     if (!ota_handle) {
         return ESP_FAIL;
     }
@@ -267,7 +276,19 @@ esp_err_t esp_rmaker_ota_default_cb(esp_rmaker_ota_handle_t ota_handle, esp_rmak
     ota_finish_err = esp_https_ota_finish(https_ota_handle);
     if ((err == ESP_OK) && (ota_finish_err == ESP_OK)) {
         ESP_LOGI(TAG, "OTA upgrade successful. Rebooting in %d seconds...", OTA_REBOOT_TIMER_SEC);
+#ifdef CONFIG_BOOTLOADER_APP_ROLLBACK_ENABLE
+        nvs_handle handle;
+        esp_err_t err = nvs_open_from_partition(ESP_RMAKER_NVS_PART_NAME, RMAKER_OTA_NVS_NAMESPACE, NVS_READWRITE, &handle);
+        if (err == ESP_OK) {
+            uint8_t ota_update = 1;
+            nvs_set_blob(handle, RMAKER_OTA_UPDATE_FLAG_NVS_NAME, &ota_update, sizeof(ota_update));
+            nvs_close(handle);
+        }
+        /* Success will be reported after a reboot since Rollback is enabled */
+        esp_rmaker_ota_report_status(ota_handle, OTA_STATUS_IN_PROGRESS, "Rebooting into new firmware");
+#else
         esp_rmaker_ota_report_status(ota_handle, OTA_STATUS_SUCCESS, "OTA Upgrade finished successfully");
+#endif
         esp_rmaker_reboot(OTA_REBOOT_TIMER_SEC);
         return ESP_OK;
     } else {
@@ -284,6 +305,101 @@ esp_err_t esp_rmaker_ota_default_cb(esp_rmaker_ota_handle_t ota_handle, esp_rmak
     return ESP_FAIL;
 }
 
+
+static void event_handler(void* arg, esp_event_base_t event_base,
+                          int32_t event_id, void* event_data)
+{
+    esp_rmaker_ota_t *ota = (esp_rmaker_ota_t *)arg;
+    esp_event_handler_unregister(RMAKER_COMMON_EVENT, RMAKER_MQTT_EVENT_CONNECTED, &event_handler);
+    esp_rmaker_ota_report_status((esp_rmaker_ota_handle_t )ota, OTA_STATUS_SUCCESS, "OTA Upgrade finished and verified successfully");
+    esp_ota_mark_app_valid_cancel_rollback();
+    ota->ota_in_progress = false;
+    if (s_ota_rollback_timer) {
+        xTimerStop(s_ota_rollback_timer, portMAX_DELAY);
+        xTimerDelete(s_ota_rollback_timer, portMAX_DELAY);
+        s_ota_rollback_timer = NULL;
+    }
+    if (ota->type == OTA_USING_TOPICS) {
+        esp_rmaker_ota_fetch();
+    }
+}
+
+static void esp_ota_rollback(TimerHandle_t handle)
+{
+    ESP_LOGE(TAG, "Could not verify firmware even after %d seconds since boot-up. Rolling back.",
+            RMAKER_OTA_ROLLBACK_WAIT_PERIOD);
+    esp_ota_mark_app_invalid_rollback_and_reboot();
+}
+
+static esp_err_t esp_ota_check_for_mqtt(esp_rmaker_ota_t *ota)
+{
+    s_ota_rollback_timer = xTimerCreate("ota_rollback_tm", (RMAKER_OTA_ROLLBACK_WAIT_PERIOD * 1000) / portTICK_PERIOD_MS,
+                            pdTRUE, NULL, esp_ota_rollback);
+    if (s_ota_rollback_timer) {
+        xTimerStart(s_ota_rollback_timer, 0);
+    } else {
+        ESP_LOGW(TAG, "Could not create rollback timer. Will require manual reboot if firmware verification fails");
+    }
+
+    return esp_event_handler_register(RMAKER_COMMON_EVENT, RMAKER_MQTT_EVENT_CONNECTED, &event_handler, ota);
+}
+
+static void esp_rmaker_ota_manage_rollback(esp_rmaker_ota_config_t *ota_config, esp_rmaker_ota_t *ota)
+{
+    const esp_partition_t *running = esp_ota_get_running_partition();
+    esp_ota_img_states_t ota_state;
+    if (esp_ota_get_state_partition(running, &ota_state) == ESP_OK) {
+        ESP_LOGI(TAG, "OTA state = %d", ota_state);
+        /* Not checking for CONFIG_BOOTLOADER_APP_ROLLBACK_ENABLE here because the firmware may have
+         * it disabled, but bootloader may have it enabled, in which case, we will have to
+         * handle this state.
+         */
+        if (ota_state == ESP_OTA_IMG_PENDING_VERIFY) {
+            ESP_LOGI(TAG, "First Boot after an OTA");
+            /* Run diagnostic function */
+            bool diagnostic_is_ok = true;
+            if (ota_config->ota_diag) {
+                diagnostic_is_ok = ota_config->ota_diag();
+            }
+            if (diagnostic_is_ok) {
+                ESP_LOGI(TAG, "Diagnostics completed successfully! Continuing execution ...");
+                /* Will not mark the image valid here immediately, but instead will wait for
+                 * MQTT connection. The below flag will tell the OTA functions that the earlier
+                 * OTA is still in progress.
+                 */
+                ota->ota_in_progress = true;
+                esp_ota_check_for_mqtt(ota);
+            } else {
+                ESP_LOGE(TAG, "Diagnostics failed! Start rollback to the previous version ...");
+                esp_ota_mark_app_invalid_rollback_and_reboot();
+            }
+#ifdef CONFIG_BOOTLOADER_APP_ROLLBACK_ENABLE
+        } else {
+            /* If rollback is enabled, and the ota update flag is found, it means that the firmware was rolled back
+            */
+            nvs_handle handle;
+            esp_err_t err = nvs_open_from_partition(ESP_RMAKER_NVS_PART_NAME, RMAKER_OTA_NVS_NAMESPACE, NVS_READWRITE, &handle);
+            if (err == ESP_OK) {
+                uint8_t ota_update = 0;
+                size_t len = sizeof(ota_update);
+                if ((err = nvs_get_blob(handle, RMAKER_OTA_UPDATE_FLAG_NVS_NAME, &ota_update, &len)) == ESP_OK) {
+                    ota->rolled_back = true;
+                    nvs_erase_key(handle, RMAKER_OTA_UPDATE_FLAG_NVS_NAME);
+                    if (ota->type == OTA_USING_PARAMS) {
+                        /* Calling this only for OTA_USING_PARAMS, because for OTA_USING_TOPICS,
+                         * the work queue function will manage the status reporting later.
+                         */
+                        esp_rmaker_ota_report_status((esp_rmaker_ota_handle_t )ota,
+                                OTA_STATUS_REJECTED, "Firmware rolled back");
+                    }
+                }
+                nvs_close(handle);
+            }
+#endif
+        }
+    }
+}
+
 static const esp_rmaker_ota_config_t ota_default_config = {
     .server_cert = esp_rmaker_ota_def_cert,
 };
@@ -307,26 +423,6 @@ esp_err_t esp_rmaker_ota_enable(esp_rmaker_ota_config_t *ota_config, esp_rmaker_
         ESP_LOGE(TAG, "Failed to allocate memory for esp_rmaker_ota_t");
         return ESP_ERR_NO_MEM;
     }
-    const esp_partition_t *running = esp_ota_get_running_partition();
-    esp_ota_img_states_t ota_state;
-    if (esp_ota_get_state_partition(running, &ota_state) == ESP_OK) {
-            ESP_LOGI(TAG, "OTA state = %d", ota_state);
-        if (ota_state == ESP_OTA_IMG_PENDING_VERIFY) {
-            ESP_LOGI(TAG, "First Boot after an OTA");
-            /* Run diagnostic function */
-            bool diagnostic_is_ok = true;
-            if (ota_config->ota_diag) {
-                diagnostic_is_ok = ota_config->ota_diag();
-            }
-            if (diagnostic_is_ok) {
-                ESP_LOGI(TAG, "Diagnostics completed successfully! Continuing execution ...");
-                esp_ota_mark_app_valid_cancel_rollback();
-            } else {
-                ESP_LOGE(TAG, "Diagnostics failed! Start rollback to the previous version ...");
-                esp_ota_mark_app_invalid_rollback_and_reboot();
-            }
-        }
-    }
     if (ota_config->ota_cb) {
         ota->ota_cb = ota_config->ota_cb;
     } else {
@@ -342,6 +438,7 @@ esp_err_t esp_rmaker_ota_enable(esp_rmaker_ota_config_t *ota_config, esp_rmaker_
         err = esp_rmaker_ota_enable_using_topics(ota);
     }
     if (err == ESP_OK) {
+        esp_rmaker_ota_manage_rollback(ota_config, ota);
         ota_init_done = true;
     } else {
         free(ota);

Diff for: components/esp_rainmaker/src/ota/esp_rmaker_ota_internal.h

@@ -17,6 +17,11 @@
 #include <stdint.h>
 #include <esp_err.h>
 #include <esp_rmaker_ota.h>
+
+#define RMAKER_OTA_NVS_NAMESPACE            "rmaker_ota"
+#define RMAKER_OTA_JOB_ID_NVS_NAME          "rmaker_ota_id"
+#define RMAKER_OTA_UPDATE_FLAG_NVS_NAME     "ota_update"
+
 typedef struct {
     esp_rmaker_ota_type_t type;
     esp_rmaker_ota_cb_t ota_cb;
@@ -25,6 +30,7 @@ typedef struct {
     char *url;
     int filesize;
     bool ota_in_progress;
+    bool rolled_back;
     ota_status_t last_reported_status;
     void *transient_priv;
 } esp_rmaker_ota_t;

Diff for: components/esp_rainmaker/src/ota/esp_rmaker_ota_using_params.c

@@ -79,11 +79,10 @@ static esp_err_t esp_rmaker_ota_service_cb(const esp_rmaker_device_t *device, co
 
 esp_err_t esp_rmaker_ota_report_status_using_params(esp_rmaker_ota_handle_t ota_handle, ota_status_t status, char *additional_info)
 {
-    if (!ota_handle) {
+    const esp_rmaker_device_t *device = esp_rmaker_node_get_device_by_name(esp_rmaker_get_node(), ESP_RMAKER_OTA_SERV_NAME);
+    if (!device) {
         return ESP_FAIL;
     }
-    esp_rmaker_ota_t *ota = (esp_rmaker_ota_t *)ota_handle;
-    const esp_rmaker_device_t *device = (esp_rmaker_device_t *)ota->transient_priv;
     esp_rmaker_param_t *info_param = esp_rmaker_device_get_param_by_type(device, ESP_RMAKER_PARAM_OTA_INFO);
     esp_rmaker_param_t *status_param = esp_rmaker_device_get_param_by_type(device, ESP_RMAKER_PARAM_OTA_STATUS);
 

Diff for: components/esp_rainmaker/src/ota/esp_rmaker_ota_using_topics.c

@@ -17,10 +17,12 @@
 #include <json_generator.h>
 #include <esp_log.h>
 #include <esp_system.h>
+#include <nvs.h>
 #include <esp_rmaker_work_queue.h>
 #include <esp_rmaker_core.h>
 #include <esp_rmaker_ota.h>
 
+#include "esp_rmaker_internal.h"
 #include "esp_rmaker_ota_internal.h"
 #include "esp_rmaker_mqtt.h"
 
@@ -53,13 +55,29 @@ esp_err_t esp_rmaker_ota_report_status_using_topics(esp_rmaker_ota_handle_t ota_
     json_gen_start_object(&jstr);
     if (ota->transient_priv) {
         json_gen_obj_set_string(&jstr, "ota_job_id", (char *)ota->transient_priv);
+    } else {
+        /* This will get executed only when the OTA status is being reported after a reboot, either to
+         * indicate successful verification of new firmware, or to indicate that firmware was rolled back
+         */
+        nvs_handle handle;
+        esp_err_t err = nvs_open_from_partition(ESP_RMAKER_NVS_PART_NAME, RMAKER_OTA_NVS_NAMESPACE, NVS_READWRITE, &handle);
+        if (err == ESP_OK) {
+            char job_id[64] = {0};
+            size_t len = sizeof(job_id);
+            if ((err = nvs_get_blob(handle, RMAKER_OTA_JOB_ID_NVS_NAME, job_id, &len)) == ESP_OK) {
+                json_gen_obj_set_string(&jstr, "ota_job_id", job_id);
+                nvs_erase_key(handle, RMAKER_OTA_JOB_ID_NVS_NAME);
+            }
+            nvs_close(handle);
+        }
     }
     json_gen_obj_set_string(&jstr, "status", esp_rmaker_ota_status_to_string(status));
     json_gen_obj_set_string(&jstr, "additional_info", additional_info);
     json_gen_end_object(&jstr);
     json_gen_str_end(&jstr);
 
     char publish_topic[100];
+    ESP_LOGI(TAG, "%s",publish_payload);
     snprintf(publish_topic, sizeof(publish_topic), "node/%s/%s", node_id, OTASTATUS_TOPIC_SUFFIX);
     esp_err_t err = esp_rmaker_mqtt_publish(publish_topic, publish_payload, strlen(publish_payload),
                         RMAKER_MQTT_QOS1, NULL);
@@ -129,6 +147,12 @@ static void ota_url_handler(const char *topic, void *payload, size_t payload_len
         goto end;
     }
     json_obj_get_string(&jctx, "ota_job_id", ota_job_id, len);
+    nvs_handle handle;
+    esp_err_t err = nvs_open_from_partition(ESP_RMAKER_NVS_PART_NAME, RMAKER_OTA_NVS_NAMESPACE, NVS_READWRITE, &handle);
+    if (err == ESP_OK) {
+        nvs_set_blob(handle, RMAKER_OTA_JOB_ID_NVS_NAME, ota_job_id, strlen(ota_job_id));
+        nvs_close(handle);
+    }
     ESP_LOGI(TAG, "OTA Job ID: %s", ota_job_id);
     ota->transient_priv = ota_job_id;
     len = 0;
@@ -171,6 +195,7 @@ static void ota_url_handler(const char *topic, void *payload, size_t payload_len
 
 esp_err_t esp_rmaker_ota_fetch(void)
 {
+    ESP_LOGI(TAG, "Fetching OTA details, if any.");
     esp_rmaker_node_info_t *info = esp_rmaker_node_get_info(esp_rmaker_get_node());
     if (!info) {
         ESP_LOGE(TAG, "Node info not found. Cant send otafetch request");
@@ -217,9 +242,17 @@ static esp_err_t esp_rmaker_ota_subscribe(void *priv_data)
 
 static void esp_rmaker_ota_work_fn(void *priv_data)
 {
+    esp_rmaker_ota_t *ota = (esp_rmaker_ota_t *)priv_data;
+    /* If the firmware was rolled back, indicate that first */
+    if (ota->rolled_back) {
+        esp_rmaker_ota_report_status((esp_rmaker_ota_handle_t )ota, OTA_STATUS_REJECTED, "Firmware rolled back");
+        ota->rolled_back = false;
+    }
     esp_rmaker_ota_subscribe(priv_data);
 #ifdef CONFIG_ESP_RMAKER_OTA_AUTOFETCH
-    esp_rmaker_ota_fetch();
+    if (ota->ota_in_progress != true) {
+        esp_rmaker_ota_fetch();
+    }
     if (ota_autofetch_period > 0) {
         esp_timer_create_args_t autofetch_timer_conf = {
             .callback = esp_rmaker_ota_timer_cb_fetch,