From 56028d093268db5590c46847b9f55409f20818e8 Mon Sep 17 00:00:00 2001 From: Jesse Hills <3060199+jesserockz@users.noreply.github.com> Date: Tue, 28 Jul 2026 08:30:52 +1200 Subject: [PATCH] [safe_mode] Prevent OTA rollback when entering deep sleep (#17699) --- esphome/components/safe_mode/__init__.py | 5 +++ esphome/components/safe_mode/safe_mode.cpp | 35 ++++++++++++--- esphome/components/safe_mode/safe_mode.h | 3 ++ esphome/core/defines.h | 1 + tests/component_tests/safe_mode/__init__.py | 0 .../safe_mode/test_safe_mode.py | 45 +++++++++++++++++++ .../safe_mode/test_safe_mode_default.yaml | 8 ++++ .../safe_mode/test_safe_mode_disabled.yaml | 9 ++++ .../test_safe_mode_no_shutdown_confirm.yaml | 9 ++++ .../test-ota-rollback.esp32-idf.yaml | 19 ++++++++ .../test-ota-rollback.nrf52-mcumgr.yaml | 16 +++++++ .../test-no-shutdown-confirm.esp32-idf.yaml | 11 +++++ 12 files changed, 154 insertions(+), 7 deletions(-) create mode 100644 tests/component_tests/safe_mode/__init__.py create mode 100644 tests/component_tests/safe_mode/test_safe_mode.py create mode 100644 tests/component_tests/safe_mode/test_safe_mode_default.yaml create mode 100644 tests/component_tests/safe_mode/test_safe_mode_disabled.yaml create mode 100644 tests/component_tests/safe_mode/test_safe_mode_no_shutdown_confirm.yaml create mode 100644 tests/components/deep_sleep/test-ota-rollback.esp32-idf.yaml create mode 100644 tests/components/deep_sleep/test-ota-rollback.nrf52-mcumgr.yaml create mode 100644 tests/components/safe_mode/test-no-shutdown-confirm.esp32-idf.yaml diff --git a/esphome/components/safe_mode/__init__.py b/esphome/components/safe_mode/__init__.py index c11447e604..70096a56bc 100644 --- a/esphome/components/safe_mode/__init__.py +++ b/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)) diff --git a/esphome/components/safe_mode/safe_mode.cpp b/esphome/components/safe_mode/safe_mode.cpp index 2eb1085ee5..ce029b4f55 100644 --- a/esphome/components/safe_mode/safe_mode.cpp +++ b/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(); + // 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 diff --git a/esphome/components/safe_mode/safe_mode.h b/esphome/components/safe_mode/safe_mode.h index d81b8a42d1..0633c92a78 100644 --- a/esphome/components/safe_mode/safe_mode.h +++ b/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 diff --git a/esphome/core/defines.h b/esphome/core/defines.h index b794254e12..2b84c72a3c 100644 --- a/esphome/core/defines.h +++ b/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 diff --git a/tests/component_tests/safe_mode/__init__.py b/tests/component_tests/safe_mode/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/component_tests/safe_mode/test_safe_mode.py b/tests/component_tests/safe_mode/test_safe_mode.py new file mode 100644 index 0000000000..617a3b4855 --- /dev/null +++ b/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) diff --git a/tests/component_tests/safe_mode/test_safe_mode_default.yaml b/tests/component_tests/safe_mode/test_safe_mode_default.yaml new file mode 100644 index 0000000000..07c38250aa --- /dev/null +++ b/tests/component_tests/safe_mode/test_safe_mode_default.yaml @@ -0,0 +1,8 @@ +--- +esphome: + name: test + +esp32: + board: nodemcu-32s + +safe_mode: diff --git a/tests/component_tests/safe_mode/test_safe_mode_disabled.yaml b/tests/component_tests/safe_mode/test_safe_mode_disabled.yaml new file mode 100644 index 0000000000..ac9b224191 --- /dev/null +++ b/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 diff --git a/tests/component_tests/safe_mode/test_safe_mode_no_shutdown_confirm.yaml b/tests/component_tests/safe_mode/test_safe_mode_no_shutdown_confirm.yaml new file mode 100644 index 0000000000..64df87b381 --- /dev/null +++ b/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 diff --git a/tests/components/deep_sleep/test-ota-rollback.esp32-idf.yaml b/tests/components/deep_sleep/test-ota-rollback.esp32-idf.yaml new file mode 100644 index 0000000000..bb24377675 --- /dev/null +++ b/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" diff --git a/tests/components/deep_sleep/test-ota-rollback.nrf52-mcumgr.yaml b/tests/components/deep_sleep/test-ota-rollback.nrf52-mcumgr.yaml new file mode 100644 index 0000000000..485490576d --- /dev/null +++ b/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 diff --git a/tests/components/safe_mode/test-no-shutdown-confirm.esp32-idf.yaml b/tests/components/safe_mode/test-no-shutdown-confirm.esp32-idf.yaml new file mode 100644 index 0000000000..5ee1f308a3 --- /dev/null +++ b/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