[safe_mode] Prevent OTA rollback when entering deep sleep (#17699)

esphome/components/safe_mode/__init__.py

@@ -16,6 +16,7 @@ from esphome.cpp_generator import RawExpression
 CODEOWNERS = ["@paulmonigatti", "@jsuanet", "@kbx81"]
 
 CONF_BOOT_IS_GOOD_AFTER = "boot_is_good_after"
+CONF_BOOT_IS_GOOD_ON_SHUTDOWN = "boot_is_good_on_shutdown"
 CONF_ON_SAFE_MODE = "on_safe_mode"
 
 safe_mode_ns = cg.esphome_ns.namespace("safe_mode")
@@ -37,6 +38,7 @@ CONFIG_SCHEMA = cv.All(
             cv.Optional(
                 CONF_BOOT_IS_GOOD_AFTER, default="1min"
             ): cv.positive_time_period_milliseconds,
+            cv.Optional(CONF_BOOT_IS_GOOD_ON_SHUTDOWN, default=True): cv.boolean,
             cv.Optional(CONF_DISABLED, default=False): cv.boolean,
             cv.Optional(CONF_NUM_ATTEMPTS, default="10"): cv.positive_not_null_int,
             cv.Optional(
@@ -78,6 +80,9 @@ async def to_code(config):
         var = cg.new_Pvariable(config[CONF_ID])
         await cg.register_component(var, config)
 
+        if config[CONF_BOOT_IS_GOOD_ON_SHUTDOWN]:
+            cg.add_define("USE_SAFE_MODE_BOOT_IS_GOOD_ON_SHUTDOWN")
+
         if on_safe_mode := config.get(CONF_ON_SAFE_MODE):
             cg.add_define("USE_SAFE_MODE_CALLBACK")
             cg.add_define("ESPHOME_SAFE_MODE_CALLBACK_COUNT", len(on_safe_mode))

esphome/components/safe_mode/safe_mode.cpp

@@ -117,19 +117,31 @@ void SafeModeComponent::dump_config() {
 
 float SafeModeComponent::get_setup_priority() const { return setup_priority::AFTER_WIFI; }
 
-void SafeModeComponent::mark_successful() {
-  this->clean_rtc();
-  this->boot_successful_ = true;
-#if defined(USE_OTA_ROLLBACK)
-// Mark OTA partition as valid to prevent rollback
+#ifdef USE_OTA_ROLLBACK
+void SafeModeComponent::confirm_app_image_() {
+  // Mark the running app image as valid so the bootloader will not roll back
+  // to the previously flashed image
 #if defined(USE_ZEPHYR)
   if (!boot_is_img_confirmed()) {
     boot_write_img_confirmed();
   }
 #elif defined(USE_ESP32)
-  // Mark OTA partition as valid to prevent rollback
+  // esp_ota_mark_app_valid_cancel_rollback() acts on the partition selected
+  // for the next boot, not the running one. After an OTA update those differ:
+  // the new image is already selected, and marking it valid before it has ever
+  // booted would disable rollback protection for that update.
+  if (esp_ota_get_running_partition() == esp_ota_get_boot_partition()) {
     esp_ota_mark_app_valid_cancel_rollback();
+  }
 #endif
+}
+#endif
+
+void SafeModeComponent::mark_successful() {
+  this->clean_rtc();
+  this->boot_successful_ = true;
+#ifdef USE_OTA_ROLLBACK
+  this->confirm_app_image_();
 #endif
   // Disable loop since we no longer need to check
   this->disable_loop();
@@ -266,6 +278,15 @@ void SafeModeComponent::clean_rtc() {
 void SafeModeComponent::on_safe_shutdown() {
   if (this->read_rtc_() != SafeModeComponent::ENTER_SAFE_MODE_MAGIC)
     this->clean_rtc();
+#if defined(USE_OTA_ROLLBACK) && defined(USE_SAFE_MODE_BOOT_IS_GOOD_ON_SHUTDOWN)
+  // An orderly shutdown (deep sleep, restart, power off) means the firmware is
+  // functional, so confirm the running app image even if boot_is_good_after has
+  // not elapsed yet. Without this, a device that enters deep sleep shortly
+  // after waking would have every OTA update rolled back by the bootloader on
+  // the next wake. Can be turned off with boot_is_good_on_shutdown: false for
+  // strict rollback semantics.
+  this->confirm_app_image_();
+#endif
 }
 
 }  // namespace esphome::safe_mode

esphome/components/safe_mode/safe_mode.h

@@ -42,6 +42,9 @@ class SafeModeComponent final : public Component {
  protected:
   void write_rtc_(uint32_t val);
   uint32_t read_rtc_();
+#ifdef USE_OTA_ROLLBACK
+  void confirm_app_image_();
+#endif
 
   // Group all 4-byte aligned members together to avoid padding
   uint32_t safe_mode_boot_is_good_after_{60000};  ///< The amount of time after which the boot is considered successful

esphome/core/defines.h

@@ -159,6 +159,7 @@
 #define USE_PREFERENCES_SYNC_EVERY_LOOP
 #define USE_PROVISIONING
 #define USE_QR_CODE
+#define USE_SAFE_MODE_BOOT_IS_GOOD_ON_SHUTDOWN
 #define USE_SAFE_MODE_CALLBACK
 #define ESPHOME_SAFE_MODE_CALLBACK_COUNT 1
 #define USE_SELECT

tests/component_tests/safe_mode/test_safe_mode.py

@@ -0,0 +1,45 @@
+"""Tests for the safe_mode component."""
+
+from collections.abc import Callable
+
+from esphome.core import CORE
+
+SHUTDOWN_DEFINE = "USE_SAFE_MODE_BOOT_IS_GOOD_ON_SHUTDOWN"
+
+
+def _has_define(name: str) -> bool:
+    return any(define.name == name for define in CORE.defines)
+
+
+def test_boot_is_good_on_shutdown_default(
+    generate_main: Callable[[str], str],
+) -> None:
+    """By default, an orderly shutdown confirms the app image."""
+    main_cpp = generate_main(
+        "tests/component_tests/safe_mode/test_safe_mode_default.yaml"
+    )
+
+    assert "safe_mode::SafeModeComponent" in main_cpp
+    assert _has_define(SHUTDOWN_DEFINE)
+
+
+def test_boot_is_good_on_shutdown_disabled(
+    generate_main: Callable[[str], str],
+) -> None:
+    """With boot_is_good_on_shutdown: false, the define is not added."""
+    main_cpp = generate_main(
+        "tests/component_tests/safe_mode/test_safe_mode_no_shutdown_confirm.yaml"
+    )
+
+    assert "safe_mode::SafeModeComponent" in main_cpp
+    assert not _has_define(SHUTDOWN_DEFINE)
+
+
+def test_safe_mode_disabled(generate_main: Callable[[str], str]) -> None:
+    """With safe_mode disabled, no component and no define are generated."""
+    main_cpp = generate_main(
+        "tests/component_tests/safe_mode/test_safe_mode_disabled.yaml"
+    )
+
+    assert "safe_mode::SafeModeComponent" not in main_cpp
+    assert not _has_define(SHUTDOWN_DEFINE)

tests/component_tests/safe_mode/test_safe_mode_default.yaml

@@ -0,0 +1,8 @@
+---
+esphome:
+  name: test
+
+esp32:
+  board: nodemcu-32s
+
+safe_mode:

tests/component_tests/safe_mode/test_safe_mode_disabled.yaml

@@ -0,0 +1,9 @@
+---
+esphome:
+  name: test
+
+esp32:
+  board: nodemcu-32s
+
+safe_mode:
+  disabled: true

tests/component_tests/safe_mode/test_safe_mode_no_shutdown_confirm.yaml

@@ -0,0 +1,9 @@
+---
+esphome:
+  name: test
+
+esp32:
+  board: nodemcu-32s
+
+safe_mode:
+  boot_is_good_on_shutdown: false

tests/components/deep_sleep/test-ota-rollback.esp32-idf.yaml

@@ -0,0 +1,19 @@
+# Deep sleep combined with OTA while bootloader rollback support is enabled
+# (the default on ESP-IDF). Entering deep sleep runs the safe shutdown hooks,
+# where safe_mode confirms the running app image so the bootloader does not
+# roll back a fresh OTA update when the device goes to sleep before
+# boot_is_good_after has elapsed.
+substitutions:
+  wakeup_pin: GPIO4
+
+packages:
+  deep_sleep: !include common.yaml
+  deep_sleep_esp32: !include common-esp32.yaml
+
+wifi:
+  ssid: MySSID
+  password: password1
+
+ota:
+  - platform: esphome
+    password: "superlongpasswordthatnoonewillknow"

tests/components/deep_sleep/test-ota-rollback.nrf52-mcumgr.yaml

@@ -0,0 +1,16 @@
+# Deep sleep combined with mcumgr OTA while MCUboot image rollback is enabled
+# (the default on nRF52). Entering system-off deep sleep runs the safe
+# shutdown hooks, where safe_mode confirms the running image so MCUboot does
+# not revert a fresh OTA update on the next wake.
+packages:
+  deep_sleep: !include common.yaml
+
+deep_sleep:
+  run_duration: 10s
+
+zephyr_ble_server:
+
+ota:
+  - platform: zephyr_mcumgr
+    transport:
+      ble: true

tests/components/safe_mode/test-no-shutdown-confirm.esp32-idf.yaml

@@ -0,0 +1,11 @@
+# Compile with OTA rollback support active (ota + safe_mode on ESP-IDF, the
+# default) but boot_is_good_on_shutdown disabled, so an orderly shutdown does
+# not confirm the app image; only boot_is_good_after / mark_successful do.
+packages:
+  safe_mode: !include common-enabled.yaml
+
+safe_mode:
+  boot_is_good_on_shutdown: false
+
+ota:
+  - platform: esphome