[esp32] Add OTA rollback support (#12460)

Co-authored-by: Claude <noreply@anthropic.com>

esphome/components/esp32/__init__.py

@@ -13,6 +13,7 @@ from esphome.const import (
     CONF_ADVANCED,
     CONF_BOARD,
     CONF_COMPONENTS,
+    CONF_DISABLED,
     CONF_ESPHOME,
     CONF_FRAMEWORK,
     CONF_IGNORE_EFUSE_CUSTOM_MAC,
@@ -24,6 +25,7 @@ from esphome.const import (
     CONF_PLATFORMIO_OPTIONS,
     CONF_REF,
     CONF_REFRESH,
+    CONF_SAFE_MODE,
     CONF_SOURCE,
     CONF_TYPE,
     CONF_VARIANT,
@@ -81,6 +83,7 @@ CONF_ASSERTION_LEVEL = "assertion_level"
 CONF_COMPILER_OPTIMIZATION = "compiler_optimization"
 CONF_ENABLE_IDF_EXPERIMENTAL_FEATURES = "enable_idf_experimental_features"
 CONF_ENABLE_LWIP_ASSERT = "enable_lwip_assert"
+CONF_ENABLE_OTA_ROLLBACK = "enable_ota_rollback"
 CONF_EXECUTE_FROM_PSRAM = "execute_from_psram"
 CONF_RELEASE = "release"
 
@@ -571,6 +574,13 @@ def final_validate(config):
                 path=[CONF_FLASH_SIZE],
             )
         )
+    if advanced[CONF_ENABLE_OTA_ROLLBACK]:
+        safe_mode_config = full_config.get(CONF_SAFE_MODE)
+        if safe_mode_config is None or safe_mode_config.get(CONF_DISABLED, False):
+            _LOGGER.warning(
+                "OTA rollback requires safe_mode, disabling rollback support"
+            )
+            advanced[CONF_ENABLE_OTA_ROLLBACK] = False
     if errs:
         raise cv.MultipleInvalid(errs)
 
@@ -691,6 +701,7 @@ FRAMEWORK_SCHEMA = cv.Schema(
                 cv.Optional(CONF_LOOP_TASK_STACK_SIZE, default=8192): cv.int_range(
                     min=8192, max=32768
                 ),
+                cv.Optional(CONF_ENABLE_OTA_ROLLBACK, default=True): cv.boolean,
             }
         ),
         cv.Optional(CONF_COMPONENTS, default=[]): cv.ensure_list(
@@ -1158,6 +1169,11 @@ async def to_code(config):
                 "CONFIG_BOOTLOADER_CACHE_32BIT_ADDR_QUAD_FLASH", True
             )
 
+    # Enable OTA rollback support
+    if advanced[CONF_ENABLE_OTA_ROLLBACK]:
+        add_idf_sdkconfig_option("CONFIG_BOOTLOADER_APP_ROLLBACK_ENABLE", True)
+        cg.add_define("USE_OTA_ROLLBACK")
+
     cg.add_define("ESPHOME_LOOP_TASK_STACK_SIZE", advanced[CONF_LOOP_TASK_STACK_SIZE])
 
     cg.add_define(

esphome/components/esp32/core.cpp

@@ -38,15 +38,11 @@ void arch_init() {
   // Enable the task watchdog only on the loop task (from which we're currently running)
   esp_task_wdt_add(nullptr);
 
-  // If the bootloader was compiled with CONFIG_BOOTLOADER_APP_ROLLBACK_ENABLE the current
-  // partition will get rolled back unless it is marked as valid.
-  esp_ota_img_states_t state;
-  const esp_partition_t *running = esp_ota_get_running_partition();
-  if (esp_ota_get_state_partition(running, &state) == ESP_OK) {
-    if (state == ESP_OTA_IMG_PENDING_VERIFY) {
-      esp_ota_mark_app_valid_cancel_rollback();
-    }
-  }
+  // Handle OTA rollback: mark partition valid immediately unless USE_OTA_ROLLBACK is enabled,
+  // in which case safe_mode will mark it valid after confirming successful boot.
+#ifndef USE_OTA_ROLLBACK
+  esp_ota_mark_app_valid_cancel_rollback();
+#endif
 }
 void IRAM_ATTR HOT arch_feed_wdt() { esp_task_wdt_reset(); }
 

esphome/components/safe_mode/safe_mode.cpp

@@ -9,6 +9,10 @@
 #include <cinttypes>
 #include <cstdio>
 
+#ifdef USE_OTA_ROLLBACK
+#include <esp_ota_ops.h>
+#endif
+
 namespace esphome {
 namespace safe_mode {
 
@@ -32,6 +36,14 @@ void SafeModeComponent::dump_config() {
       ESP_LOGW(TAG, "SAFE MODE IS ACTIVE");
     }
   }
+
+#ifdef USE_OTA_ROLLBACK
+  const esp_partition_t *last_invalid = esp_ota_get_last_invalid_partition();
+  if (last_invalid != nullptr) {
+    ESP_LOGW(TAG, "OTA rollback detected! Rolled back from partition '%s'", last_invalid->label);
+    ESP_LOGW(TAG, "The device reset before the boot was marked successful");
+  }
+#endif
 }
 
 float SafeModeComponent::get_setup_priority() const { return setup_priority::AFTER_WIFI; }
@@ -42,6 +54,10 @@ void SafeModeComponent::loop() {
     ESP_LOGI(TAG, "Boot seems successful; resetting boot loop counter");
     this->clean_rtc();
     this->boot_successful_ = true;
+#ifdef USE_OTA_ROLLBACK
+    // Mark OTA partition as valid to prevent rollback
+    esp_ota_mark_app_valid_cancel_rollback();
+#endif
     // Disable loop since we no longer need to check
     this->disable_loop();
   }

esphome/core/defines.h

@@ -170,6 +170,7 @@
 // ESP32-specific feature flags
 #ifdef USE_ESP32
 #define USE_ESPHOME_TASK_LOG_BUFFER
+#define USE_OTA_ROLLBACK
 
 #define USE_BLUETOOTH_PROXY
 #define BLUETOOTH_PROXY_MAX_CONNECTIONS 3

tests/components/esp32/test.esp32-idf.yaml

@@ -3,6 +3,7 @@ esp32:
   framework:
     type: esp-idf
     advanced:
+      enable_ota_rollback: true
       enable_lwip_mdns_queries: true
       enable_lwip_bridge_interface: true
       disable_libc_locks_in_iram: false  # Test explicit opt-out of RAM optimization