[safe_mode] Allow recovering soft-bricked devices via reboot to recovery partition (#16339)

Co-authored-by: J. Nick Koston <nick@koston.org>
Co-authored-by: J. Nick Koston <nick@home-assistant.io>

esphome/components/safe_mode/safe_mode.cpp

@@ -15,6 +15,7 @@
 #elif defined(USE_ESP32)
 #include <esp_ota_ops.h>
 #include <esp_system.h>
+#include <esp_image_format.h>
 #endif
 #endif
 
@@ -22,6 +23,37 @@ namespace esphome::safe_mode {
 
 static const char *const TAG = "safe_mode";
 
+#if defined(USE_ESP32) && defined(USE_OTA_ROLLBACK) && !defined(USE_OTA_PARTITIONS)
+// Find a non-running app partition. If verify is true, only returns a partition
+// whose image passes verification (expensive: reads flash). Returns nullptr if none found.
+static const esp_partition_t *find_alternate_app_partition(bool verify) {
+  const esp_partition_t *running = esp_ota_get_running_partition();
+  const esp_partition_t *result = nullptr;
+  esp_partition_iterator_t it = esp_partition_find(ESP_PARTITION_TYPE_APP, ESP_PARTITION_SUBTYPE_ANY, nullptr);
+  while (it != nullptr) {
+    const esp_partition_t *p = esp_partition_get(it);
+    if (p->address != running->address) {
+      if (!verify) {
+        result = p;
+        break;
+      }
+      esp_image_metadata_t data = {};
+      const esp_partition_pos_t part_pos = {
+          .offset = p->address,
+          .size = p->size,
+      };
+      if (esp_image_verify(ESP_IMAGE_VERIFY_SILENT, &part_pos, &data) == ESP_OK) {
+        result = p;
+        break;
+      }
+    }
+    it = esp_partition_next(it);
+  }
+  esp_partition_iterator_release(it);
+  return result;
+}
+#endif
+
 void SafeModeComponent::dump_config() {
   ESP_LOGCONFIG(TAG,
                 "Safe Mode:\n"
@@ -34,7 +66,11 @@ void SafeModeComponent::dump_config() {
 #if defined(USE_ESP32) && defined(USE_OTA_ROLLBACK)
   const char *state_str;
   if (this->ota_state_ == ESP_OTA_IMG_NEW) {
+#ifdef USE_OTA_PARTITIONS
+    state_str = "support unknown";
+#else
     state_str = "not supported";
+#endif
   } else if (this->ota_state_ == ESP_OTA_IMG_PENDING_VERIFY) {
     state_str = "supported";
   } else {
@@ -64,6 +100,18 @@ void SafeModeComponent::dump_config() {
                     " See https://esphome.io/guides/faq.html#brownout-detector-was-triggered");
     }
   }
+  if (!this->app_ota_possible_) {
+    ESP_LOGW(TAG, "OTA updates are impossible.");
+#ifdef USE_OTA_PARTITIONS
+    ESP_LOGW(TAG, " OTA partition table update or serial flashing is required.");
+#else
+    if (find_alternate_app_partition(false) != nullptr) {
+      ESP_LOGW(TAG, " Activate safe mode to reboot to the recovery partition.");
+    } else {
+      ESP_LOGE(TAG, " No recovery partition available; serial flashing is required.");
+    }
+#endif
+  }
 #endif
 }
 
@@ -124,8 +172,10 @@ bool SafeModeComponent::should_enter_safe_mode(uint8_t num_attempts, uint32_t en
 
 #if defined(USE_ESP32) && defined(USE_OTA_ROLLBACK)
   // Check partition state to detect if bootloader supports rollback
-  const esp_partition_t *running = esp_ota_get_running_partition();
-  esp_ota_get_state_partition(running, &this->ota_state_);
+  const esp_partition_t *running_part = esp_ota_get_running_partition();
+  esp_ota_get_state_partition(running_part, &this->ota_state_);
+  const esp_partition_t *next_part = esp_ota_get_next_update_partition(nullptr);
+  this->app_ota_possible_ = (next_part != nullptr && next_part != running_part);
 #endif
 
   uint32_t rtc_val = this->read_rtc_();
@@ -151,6 +201,28 @@ bool SafeModeComponent::should_enter_safe_mode(uint8_t num_attempts, uint32_t en
     ESP_LOGE(TAG, "Boot loop detected");
   }
 
+#if defined(USE_ESP32) && defined(USE_OTA_ROLLBACK) && !defined(USE_OTA_PARTITIONS)
+  // Allow recovery of soft-bricked devices
+  // Instead of starting safe_mode, reboot to the other app partition if all conditions are met:
+  // - app OTA is impossible (for example because the other app partition has type 'factory')
+  // - the other app partition contains a valid app (for example Tasmota safeboot image or ESPHome)
+  // - allow_partition_access is not configured making recovery via partition table update impossible
+  // Image verification is deferred until here so the cost is only paid when entering safe mode,
+  // not on every boot.
+  if (!this->app_ota_possible_) {
+    const esp_partition_t *rollback_part = find_alternate_app_partition(true);
+    if (rollback_part != nullptr) {
+      esp_err_t err = esp_ota_set_boot_partition(rollback_part);
+      if (err == ESP_OK) {
+        ESP_LOGW(TAG, "OTA updates are impossible. Rebooting to recovery app.");
+        App.reboot();
+      } else {
+        ESP_LOGE(TAG, "Failed to set recovery boot partition: %s", esp_err_to_name(err));
+      }
+    }
+  }
+#endif
+
   this->status_set_error();
   this->set_timeout(enable_time, []() {
     ESP_LOGW(TAG, "Timeout, restarting");

esphome/components/safe_mode/safe_mode.h

@@ -48,11 +48,14 @@ class SafeModeComponent final : public Component {
   uint32_t safe_mode_enable_time_{60000};         ///< The time safe mode should remain active for
   uint32_t safe_mode_rtc_value_{0};
   uint32_t safe_mode_start_time_{0};  ///< stores when safe mode was enabled
+#if defined(USE_ESP32) && defined(USE_OTA_ROLLBACK)
+  esp_ota_img_states_t ota_state_{ESP_OTA_IMG_UNDEFINED};  // 4-byte enum
+#endif
   // Group 1-byte members together to minimize padding
   bool boot_successful_{false};  ///< set to true after boot is considered successful
   uint8_t safe_mode_num_attempts_{0};
 #if defined(USE_ESP32) && defined(USE_OTA_ROLLBACK)
-  esp_ota_img_states_t ota_state_{ESP_OTA_IMG_UNDEFINED};
+  bool app_ota_possible_{true};
 #endif
   // Larger objects at the end
   ESPPreferenceObject rtc_;