Uncouple safe_mode from OTA (#6759)

esphome/components/safe_mode/__init__.py

@@ -1,5 +1,56 @@
+from esphome.cpp_generator import RawExpression
 import esphome.codegen as cg
+import esphome.config_validation as cv
+from esphome.const import (
+    CONF_DISABLED,
+    CONF_ID,
+    CONF_NUM_ATTEMPTS,
+    CONF_REBOOT_TIMEOUT,
+    CONF_SAFE_MODE,
+    KEY_PAST_SAFE_MODE,
+)
+from esphome.core import CORE, coroutine_with_priority
 
-CODEOWNERS = ["@paulmonigatti", "@jsuanet"]
+
+CODEOWNERS = ["@paulmonigatti", "@jsuanet", "@kbx81"]
 
 safe_mode_ns = cg.esphome_ns.namespace("safe_mode")
+SafeModeComponent = safe_mode_ns.class_("SafeModeComponent", cg.Component)
+
+
+def _remove_id_if_disabled(value):
+    value = value.copy()
+    if value[CONF_DISABLED]:
+        value.pop(CONF_ID)
+    return value
+
+
+CONFIG_SCHEMA = cv.All(
+    cv.Schema(
+        {
+            cv.GenerateID(): cv.declare_id(SafeModeComponent),
+            cv.Optional(CONF_DISABLED, default=False): cv.boolean,
+            cv.Optional(CONF_NUM_ATTEMPTS, default="10"): cv.positive_not_null_int,
+            cv.Optional(
+                CONF_REBOOT_TIMEOUT, default="5min"
+            ): cv.positive_time_period_milliseconds,
+        }
+    ).extend(cv.COMPONENT_SCHEMA),
+    _remove_id_if_disabled,
+)
+
+
+@coroutine_with_priority(50.0)
+async def to_code(config):
+    if config[CONF_DISABLED]:
+        return
+
+    var = cg.new_Pvariable(config[CONF_ID])
+    await cg.register_component(var, config)
+
+    condition = var.should_enter_safe_mode(
+        config[CONF_NUM_ATTEMPTS], config[CONF_REBOOT_TIMEOUT]
+    )
+    cg.add(RawExpression(f"if ({condition}) return"))
+    CORE.data[CONF_SAFE_MODE] = {}
+    CORE.data[CONF_SAFE_MODE][KEY_PAST_SAFE_MODE] = True

esphome/components/safe_mode/button/__init__.py

@@ -1,16 +1,15 @@
 import esphome.codegen as cg
 import esphome.config_validation as cv
 from esphome.components import button
-from esphome.components.esphome.ota import ESPHomeOTAComponent
 from esphome.const import (
-    CONF_ESPHOME,
+    CONF_SAFE_MODE,
     DEVICE_CLASS_RESTART,
     ENTITY_CATEGORY_CONFIG,
     ICON_RESTART_ALERT,
 )
-from .. import safe_mode_ns
+from .. import safe_mode_ns, SafeModeComponent
 
-DEPENDENCIES = ["ota.esphome"]
+DEPENDENCIES = ["safe_mode"]
 
 SafeModeButton = safe_mode_ns.class_("SafeModeButton", button.Button, cg.Component)
 
@@ -21,7 +20,7 @@ CONFIG_SCHEMA = (
         entity_category=ENTITY_CATEGORY_CONFIG,
         icon=ICON_RESTART_ALERT,
     )
-    .extend({cv.GenerateID(CONF_ESPHOME): cv.use_id(ESPHomeOTAComponent)})
+    .extend({cv.GenerateID(CONF_SAFE_MODE): cv.use_id(SafeModeComponent)})
     .extend(cv.COMPONENT_SCHEMA)
 )
 
@@ -30,5 +29,5 @@ async def to_code(config):
     var = await button.new_button(config)
     await cg.register_component(var, config)
 
-    ota = await cg.get_variable(config[CONF_ESPHOME])
-    cg.add(var.set_ota(ota))
+    safe_mode_component = await cg.get_variable(config[CONF_SAFE_MODE])
+    cg.add(var.set_safe_mode(safe_mode_component))

esphome/components/safe_mode/button/safe_mode_button.cpp

@@ -8,11 +8,13 @@ namespace safe_mode {
 
 static const char *const TAG = "safe_mode.button";
 
-void SafeModeButton::set_ota(esphome::ESPHomeOTAComponent *ota) { this->ota_ = ota; }
+void SafeModeButton::set_safe_mode(SafeModeComponent *safe_mode_component) {
+  this->safe_mode_component_ = safe_mode_component;
+}
 
 void SafeModeButton::press_action() {
   ESP_LOGI(TAG, "Restarting device in safe mode...");
-  this->ota_->set_safe_mode_pending(true);
+  this->safe_mode_component_->set_safe_mode_pending(true);
 
   // Let MQTT settle a bit
   delay(100);  // NOLINT

esphome/components/safe_mode/button/safe_mode_button.h

@@ -1,7 +1,7 @@
 #pragma once
 
 #include "esphome/components/button/button.h"
-#include "esphome/components/esphome/ota/ota_esphome.h"
+#include "esphome/components/safe_mode/safe_mode.h"
 #include "esphome/core/component.h"
 
 namespace esphome {
@@ -10,10 +10,10 @@ namespace safe_mode {
 class SafeModeButton : public button::Button, public Component {
  public:
   void dump_config() override;
-  void set_ota(esphome::ESPHomeOTAComponent *ota);
+  void set_safe_mode(SafeModeComponent *safe_mode_component);
 
  protected:
-  esphome::ESPHomeOTAComponent *ota_;
+  SafeModeComponent *safe_mode_component_;
   void press_action() override;
 };
 

esphome/components/safe_mode/safe_mode.cpp

@@ -0,0 +1,125 @@
+#include "safe_mode.h"
+
+#include "esphome/core/application.h"
+#include "esphome/core/hal.h"
+#include "esphome/core/log.h"
+#include "esphome/core/util.h"
+
+#include <cerrno>
+#include <cinttypes>
+#include <cstdio>
+
+namespace esphome {
+namespace safe_mode {
+
+static const char *const TAG = "safe_mode";
+
+void SafeModeComponent::dump_config() {
+  ESP_LOGCONFIG(TAG, "Safe Mode:");
+  ESP_LOGCONFIG(TAG, "  Invoke after %u boot attempts", this->safe_mode_num_attempts_);
+  ESP_LOGCONFIG(TAG, "  Remain in safe mode for %" PRIu32 " seconds",
+                this->safe_mode_enable_time_ / 1000);  // because milliseconds
+
+  if (this->safe_mode_rtc_value_ > 1 && this->safe_mode_rtc_value_ != SafeModeComponent::ENTER_SAFE_MODE_MAGIC) {
+    auto remaining_restarts = this->safe_mode_num_attempts_ - this->safe_mode_rtc_value_;
+    if (remaining_restarts) {
+      ESP_LOGW(TAG, "Last reset occurred too quickly; safe mode will be invoked in %" PRIu32 " restarts",
+               remaining_restarts);
+    } else {
+      ESP_LOGW(TAG, "SAFE MODE IS ACTIVE");
+    }
+  }
+}
+
+float SafeModeComponent::get_setup_priority() const { return setup_priority::AFTER_WIFI; }
+
+void SafeModeComponent::loop() {
+  if (!this->boot_successful_ && (millis() - this->safe_mode_start_time_) > this->safe_mode_enable_time_) {
+    // successful boot, reset counter
+    ESP_LOGI(TAG, "Boot seems successful; resetting boot loop counter");
+    this->clean_rtc();
+    this->boot_successful_ = true;
+  }
+}
+
+void SafeModeComponent::set_safe_mode_pending(const bool &pending) {
+  uint32_t current_rtc = this->read_rtc_();
+
+  if (pending && current_rtc != SafeModeComponent::ENTER_SAFE_MODE_MAGIC) {
+    ESP_LOGI(TAG, "Device will enter safe mode on next boot");
+    this->write_rtc_(SafeModeComponent::ENTER_SAFE_MODE_MAGIC);
+  }
+
+  if (!pending && current_rtc == SafeModeComponent::ENTER_SAFE_MODE_MAGIC) {
+    ESP_LOGI(TAG, "Safe mode pending has been cleared");
+    this->clean_rtc();
+  }
+}
+
+bool SafeModeComponent::get_safe_mode_pending() {
+  return this->read_rtc_() == SafeModeComponent::ENTER_SAFE_MODE_MAGIC;
+}
+
+bool SafeModeComponent::should_enter_safe_mode(uint8_t num_attempts, uint32_t enable_time) {
+  this->safe_mode_start_time_ = millis();
+  this->safe_mode_enable_time_ = enable_time;
+  this->safe_mode_num_attempts_ = num_attempts;
+  this->rtc_ = global_preferences->make_preference<uint32_t>(233825507UL, false);
+  this->safe_mode_rtc_value_ = this->read_rtc_();
+
+  bool is_manual_safe_mode = this->safe_mode_rtc_value_ == SafeModeComponent::ENTER_SAFE_MODE_MAGIC;
+
+  if (is_manual_safe_mode) {
+    ESP_LOGI(TAG, "Safe mode invoked manually");
+  } else {
+    ESP_LOGCONFIG(TAG, "There have been %" PRIu32 " suspected unsuccessful boot attempts", this->safe_mode_rtc_value_);
+  }
+
+  if (this->safe_mode_rtc_value_ >= num_attempts || is_manual_safe_mode) {
+    this->clean_rtc();
+
+    if (!is_manual_safe_mode) {
+      ESP_LOGE(TAG, "Boot loop detected. Proceeding to safe mode");
+    }
+
+    this->status_set_error();
+    this->set_timeout(enable_time, []() {
+      ESP_LOGW(TAG, "Safe mode enable time has elapsed -- restarting");
+      App.reboot();
+    });
+
+    // Delay here to allow power to stabilize before Wi-Fi/Ethernet is initialised
+    delay(300);  // NOLINT
+    App.setup();
+
+    ESP_LOGW(TAG, "SAFE MODE IS ACTIVE");
+
+    return true;
+  } else {
+    // increment counter
+    this->write_rtc_(this->safe_mode_rtc_value_ + 1);
+    return false;
+  }
+}
+
+void SafeModeComponent::write_rtc_(uint32_t val) {
+  this->rtc_.save(&val);
+  global_preferences->sync();
+}
+
+uint32_t SafeModeComponent::read_rtc_() {
+  uint32_t val;
+  if (!this->rtc_.load(&val))
+    return 0;
+  return val;
+}
+
+void SafeModeComponent::clean_rtc() { this->write_rtc_(0); }
+
+void SafeModeComponent::on_safe_shutdown() {
+  if (this->read_rtc_() != SafeModeComponent::ENTER_SAFE_MODE_MAGIC)
+    this->clean_rtc();
+}
+
+}  // namespace safe_mode
+}  // namespace esphome

esphome/components/safe_mode/safe_mode.h

@@ -0,0 +1,44 @@
+#pragma once
+
+#include "esphome/core/component.h"
+#include "esphome/core/defines.h"
+#include "esphome/core/helpers.h"
+#include "esphome/core/preferences.h"
+
+namespace esphome {
+namespace safe_mode {
+
+/// SafeModeComponent provides a safe way to recover from repeated boot failures
+class SafeModeComponent : public Component {
+ public:
+  bool should_enter_safe_mode(uint8_t num_attempts, uint32_t enable_time);
+
+  /// Set to true if the next startup will enter safe mode
+  void set_safe_mode_pending(const bool &pending);
+  bool get_safe_mode_pending();
+
+  void dump_config() override;
+  float get_setup_priority() const override;
+  void loop() override;
+
+  void clean_rtc();
+
+  void on_safe_shutdown() override;
+
+ protected:
+  void write_rtc_(uint32_t val);
+  uint32_t read_rtc_();
+
+  bool boot_successful_{false};            ///< set to true after boot is considered successful
+  uint32_t safe_mode_start_time_;          ///< stores when safe mode was enabled
+  uint32_t safe_mode_enable_time_{60000};  ///< The time safe mode should remain active for
+  uint32_t safe_mode_rtc_value_;
+  uint8_t safe_mode_num_attempts_;
+  ESPPreferenceObject rtc_;
+
+  static const uint32_t ENTER_SAFE_MODE_MAGIC =
+      0x5afe5afe;  ///< a magic number to indicate that safe mode should be entered on next boot
+};
+
+}  // namespace safe_mode
+}  // namespace esphome

esphome/components/safe_mode/switch/__init__.py

@@ -1,15 +1,14 @@
 import esphome.codegen as cg
 import esphome.config_validation as cv
 from esphome.components import switch
-from esphome.components.esphome.ota import ESPHomeOTAComponent
 from esphome.const import (
-    CONF_ESPHOME,
+    CONF_SAFE_MODE,
     ENTITY_CATEGORY_CONFIG,
     ICON_RESTART_ALERT,
 )
-from .. import safe_mode_ns
+from .. import safe_mode_ns, SafeModeComponent
 
-DEPENDENCIES = ["ota.esphome"]
+DEPENDENCIES = ["safe_mode"]
 
 SafeModeSwitch = safe_mode_ns.class_("SafeModeSwitch", switch.Switch, cg.Component)
 
@@ -20,7 +19,7 @@ CONFIG_SCHEMA = (
         entity_category=ENTITY_CATEGORY_CONFIG,
         icon=ICON_RESTART_ALERT,
     )
-    .extend({cv.GenerateID(CONF_ESPHOME): cv.use_id(ESPHomeOTAComponent)})
+    .extend({cv.GenerateID(CONF_SAFE_MODE): cv.use_id(SafeModeComponent)})
     .extend(cv.COMPONENT_SCHEMA)
 )
 
@@ -29,5 +28,5 @@ async def to_code(config):
     var = await switch.new_switch(config)
     await cg.register_component(var, config)
 
-    ota = await cg.get_variable(config[CONF_ESPHOME])
-    cg.add(var.set_ota(ota))
+    safe_mode_component = await cg.get_variable(config[CONF_SAFE_MODE])
+    cg.add(var.set_safe_mode(safe_mode_component))

esphome/components/safe_mode/switch/safe_mode_switch.cpp

@@ -6,9 +6,11 @@
 namespace esphome {
 namespace safe_mode {
 
-static const char *const TAG = "safe_mode_switch";
+static const char *const TAG = "safe_mode.switch";
 
-void SafeModeSwitch::set_ota(esphome::ESPHomeOTAComponent *ota) { this->ota_ = ota; }
+void SafeModeSwitch::set_safe_mode(SafeModeComponent *safe_mode_component) {
+  this->safe_mode_component_ = safe_mode_component;
+}
 
 void SafeModeSwitch::write_state(bool state) {
   // Acknowledge
@@ -16,13 +18,14 @@ void SafeModeSwitch::write_state(bool state) {
 
   if (state) {
     ESP_LOGI(TAG, "Restarting device in safe mode...");
-    this->ota_->set_safe_mode_pending(true);
+    this->safe_mode_component_->set_safe_mode_pending(true);
 
     // Let MQTT settle a bit
     delay(100);  // NOLINT
     App.safe_reboot();
   }
 }
+
 void SafeModeSwitch::dump_config() { LOG_SWITCH("", "Safe Mode Switch", this); }
 
 }  // namespace safe_mode

esphome/components/safe_mode/switch/safe_mode_switch.h

@@ -1,6 +1,6 @@
 #pragma once
 
-#include "esphome/components/esphome/ota/ota_esphome.h"
+#include "esphome/components/safe_mode/safe_mode.h"
 #include "esphome/components/switch/switch.h"
 #include "esphome/core/component.h"
 
@@ -10,10 +10,10 @@ namespace safe_mode {
 class SafeModeSwitch : public switch_::Switch, public Component {
  public:
   void dump_config() override;
-  void set_ota(esphome::ESPHomeOTAComponent *ota);
+  void set_safe_mode(SafeModeComponent *safe_mode_component);
 
  protected:
-  esphome::ESPHomeOTAComponent *ota_;
+  SafeModeComponent *safe_mode_component_;
   void write_state(bool state) override;
 };