Add ESP32 OTA rollback support

mongoose.c

@@ -10330,7 +10330,7 @@ void mg_ota_poll(struct mg_mgr *mgr) {
   static uint64_t t = 5000;  // Fire first time 5 sec after boot
 
   static uint64_t feed_timer;  // Advances 500ms per tick; tracks elapsed time
-  if (MG_OTA_STATE_GET() == MG_OTA_TESTING &&
+  if (MG_OTA_STATE_GET() == MG_OTA_FAILED &&
       mg_timer_expired(&feed_timer, 500, mg_millis())) {
     if (feed_timer < (uint64_t) MG_OTA_ROLLBACK_TIMEOUT_SECONDS * 1000) {
       MG_OTA_ROLLBACK_TIMER_FEED();  // Feed watchdog every 500ms
@@ -10537,6 +10537,10 @@ bool mg_ota_end(void) {
   }
   MG_DEBUG(("Finished ESP32 OTA, success: %d", s_ota_success));
   s_ota_update_partition = NULL;
+  if (s_ota_success) {
+    MG_OTA_STATE_SET(MG_OTA_TESTING);
+    esp_restart();
+  }
   return s_ota_success;
 }
 

mongoose.h

@@ -314,6 +314,11 @@ extern "C" {
 #include <esp_timer.h>     // amalgamation ditching them
 #include <esp_system.h>
 #include <esp_mac.h>
+#include <esp_idf_version.h>
+#include <esp_ota_ops.h>
+#include <esp_task_wdt.h>
+#include <soc/soc.h>
+#include <soc/rtc_cntl_reg.h>
 
 #define MG_PATH_MAX 128
 
@@ -325,6 +330,57 @@ extern "C" {
 #define MG_ENABLE_DIRLIST 1
 #endif
 
+#ifndef MG_OTA_ESP32_WDT_TIMEOUT_MS
+#define MG_OTA_ESP32_WDT_TIMEOUT_MS 10000U
+#endif
+
+static inline void mg_ota_esp32_timer_start(void) {
+#if ESP_IDF_VERSION_MAJOR >= 5
+  esp_task_wdt_config_t cfg = {0};
+  esp_err_t err;
+
+  cfg.timeout_ms = MG_OTA_ESP32_WDT_TIMEOUT_MS;
+  cfg.idle_core_mask = 0;
+  cfg.trigger_panic = true;
+
+  err = esp_task_wdt_init(&cfg);
+  if (err == ESP_ERR_INVALID_STATE) esp_task_wdt_reconfigure(&cfg);
+#else
+  esp_task_wdt_init(MG_OTA_ESP32_WDT_TIMEOUT_MS / 1000U, true);
+#endif
+  esp_task_wdt_add(NULL);
+  esp_task_wdt_reset();
+}
+
+static inline uint32_t mg_ota_esp32_state_set(uint32_t state) {
+  if (state == 0) esp_ota_mark_app_valid_cancel_rollback();
+  REG_WRITE(RTC_CNTL_STORE0_REG, state);
+  return REG_READ(RTC_CNTL_STORE0_REG);
+}
+
+#define MG_OTA_STATE_GET() \
+  REG_READ(RTC_CNTL_STORE0_REG)
+
+#define MG_OTA_STATE_SET(v) mg_ota_esp32_state_set((uint32_t) (v))
+
+#ifndef MG_OTA_ROLLBACK_TIMER_START
+#define MG_OTA_ROLLBACK_TIMER_START() mg_ota_esp32_timer_start()
+#endif
+
+#ifndef MG_OTA_ROLLBACK_TIMER_FEED
+#define MG_OTA_ROLLBACK_TIMER_FEED() \
+  do {                               \
+    (void) esp_task_wdt_reset();     \
+  } while (0)
+#endif
+
+#ifndef MG_OTA_ROLLBACK
+#define MG_OTA_ROLLBACK()                                  \
+  do {                                                     \
+    (void) esp_ota_mark_app_invalid_rollback_and_reboot(); \
+  } while (0)
+#endif
+
 #endif
 
 

src/arch_esp32.h

@@ -21,6 +21,11 @@
 #include <esp_timer.h>     // amalgamation ditching them
 #include <esp_system.h>
 #include <esp_mac.h>
+#include <esp_idf_version.h>
+#include <esp_ota_ops.h>
+#include <esp_task_wdt.h>
+#include <soc/soc.h>
+#include <soc/rtc_cntl_reg.h>
 
 #define MG_PATH_MAX 128
 
@@ -32,4 +37,55 @@
 #define MG_ENABLE_DIRLIST 1
 #endif
 
+#ifndef MG_OTA_ESP32_WDT_TIMEOUT_MS
+#define MG_OTA_ESP32_WDT_TIMEOUT_MS 10000U
+#endif
+
+static inline void mg_ota_esp32_timer_start(void) {
+#if ESP_IDF_VERSION_MAJOR >= 5
+  esp_task_wdt_config_t cfg = {0};
+  esp_err_t err;
+
+  cfg.timeout_ms = MG_OTA_ESP32_WDT_TIMEOUT_MS;
+  cfg.idle_core_mask = 0;
+  cfg.trigger_panic = true;
+
+  err = esp_task_wdt_init(&cfg);
+  if (err == ESP_ERR_INVALID_STATE) esp_task_wdt_reconfigure(&cfg);
+#else
+  esp_task_wdt_init(MG_OTA_ESP32_WDT_TIMEOUT_MS / 1000U, true);
+#endif
+  esp_task_wdt_add(NULL);
+  esp_task_wdt_reset();
+}
+
+static inline uint32_t mg_ota_esp32_state_set(uint32_t state) {
+  if (state == 0) esp_ota_mark_app_valid_cancel_rollback();
+  REG_WRITE(RTC_CNTL_STORE0_REG, state);
+  return REG_READ(RTC_CNTL_STORE0_REG);
+}
+
+#define MG_OTA_STATE_GET() \
+  REG_READ(RTC_CNTL_STORE0_REG)
+
+#define MG_OTA_STATE_SET(v) mg_ota_esp32_state_set((uint32_t) (v))
+
+#ifndef MG_OTA_ROLLBACK_TIMER_START
+#define MG_OTA_ROLLBACK_TIMER_START() mg_ota_esp32_timer_start()
+#endif
+
+#ifndef MG_OTA_ROLLBACK_TIMER_FEED
+#define MG_OTA_ROLLBACK_TIMER_FEED() \
+  do {                               \
+    (void) esp_task_wdt_reset();     \
+  } while (0)
+#endif
+
+#ifndef MG_OTA_ROLLBACK
+#define MG_OTA_ROLLBACK()                                  \
+  do {                                                     \
+    (void) esp_ota_mark_app_invalid_rollback_and_reboot(); \
+  } while (0)
+#endif
+
 #endif

src/ota.c

@@ -185,7 +185,7 @@ void mg_ota_poll(struct mg_mgr *mgr) {
   static uint64_t t = 5000;  // Fire first time 5 sec after boot
 
   static uint64_t feed_timer;  // Advances 500ms per tick; tracks elapsed time
-  if (MG_OTA_STATE_GET() == MG_OTA_TESTING &&
+  if (MG_OTA_STATE_GET() == MG_OTA_FAILED &&
       mg_timer_expired(&feed_timer, 500, mg_millis())) {
     if (feed_timer < (uint64_t) MG_OTA_ROLLBACK_TIMEOUT_SECONDS * 1000) {
       MG_OTA_ROLLBACK_TIMER_FEED();  // Feed watchdog every 500ms

src/ota_esp32.c

@@ -46,6 +46,10 @@ bool mg_ota_end(void) {
   }
   MG_DEBUG(("Finished ESP32 OTA, success: %d", s_ota_success));
   s_ota_update_partition = NULL;
+  if (s_ota_success) {
+    MG_OTA_STATE_SET(MG_OTA_TESTING);
+    esp_restart();
+  }
   return s_ota_success;
 }