From e2f2feafdd912f5e249ccb1b47502e4201665cf2 Mon Sep 17 00:00:00 2001
From: Rene Guca <45061891+rguca@users.noreply.github.com>
Date: Thu, 18 Jan 2024 08:30:58 +0100
Subject: [PATCH] WiFi fast_connect: save/load BSSID and channel for faster
 connect from sleep (#5931)

---
 esphome/components/wifi/wifi_component.cpp | 38 ++++++++++++++++++++++
 esphome/components/wifi/wifi_component.h   |  9 +++++
 2 files changed, 47 insertions(+)

diff --git a/esphome/components/wifi/wifi_component.cpp b/esphome/components/wifi/wifi_component.cpp
index 519489097a..05938d87a2 100644
--- a/esphome/components/wifi/wifi_component.cpp
+++ b/esphome/components/wifi/wifi_component.cpp
@@ -55,6 +55,9 @@ void WiFiComponent::start() {
   uint32_t hash = this->has_sta() ? fnv1_hash(App.get_compilation_time()) : 88491487UL;
 
   this->pref_ = global_preferences->make_preference<wifi::SavedWifiSettings>(hash, true);
+  if (this->fast_connect_) {
+    this->fast_connect_pref_ = global_preferences->make_preference<wifi::SavedWifiFastConnectSettings>(hash, false);
+  }
 
   SavedWifiSettings save{};
   if (this->pref_.load(&save)) {
@@ -78,6 +81,7 @@ void WiFiComponent::start() {
 
     if (this->fast_connect_) {
       this->selected_ap_ = this->sta_[0];
+      this->load_fast_connect_settings_();
       this->start_connecting(this->selected_ap_, false);
     } else {
       this->start_scanning();
@@ -604,6 +608,11 @@ void WiFiComponent::check_connecting_finished() {
 
     this->state_ = WIFI_COMPONENT_STATE_STA_CONNECTED;
     this->num_retried_ = 0;
+
+    if (this->fast_connect_) {
+      this->save_fast_connect_settings_();
+    }
+
     return;
   }
 
@@ -705,6 +714,35 @@ bool WiFiComponent::is_esp32_improv_active_() {
 #endif
 }
 
+void WiFiComponent::load_fast_connect_settings_() {
+  SavedWifiFastConnectSettings fast_connect_save{};
+
+  if (this->fast_connect_pref_.load(&fast_connect_save)) {
+    bssid_t bssid{};
+    std::copy(fast_connect_save.bssid, fast_connect_save.bssid + 6, bssid.begin());
+    this->selected_ap_.set_bssid(bssid);
+    this->selected_ap_.set_channel(fast_connect_save.channel);
+
+    ESP_LOGD(TAG, "Loaded saved fast_connect wifi settings");
+  }
+}
+
+void WiFiComponent::save_fast_connect_settings_() {
+  bssid_t bssid = wifi_bssid();
+  uint8_t channel = wifi_channel_();
+
+  if (bssid != this->selected_ap_.get_bssid() || channel != this->selected_ap_.get_channel()) {
+    SavedWifiFastConnectSettings fast_connect_save{};
+
+    memcpy(fast_connect_save.bssid, bssid.data(), 6);
+    fast_connect_save.channel = channel;
+
+    this->fast_connect_pref_.save(&fast_connect_save);
+
+    ESP_LOGD(TAG, "Saved fast_connect wifi settings");
+  }
+}
+
 void WiFiAP::set_ssid(const std::string &ssid) { this->ssid_ = ssid; }
 void WiFiAP::set_bssid(bssid_t bssid) { this->bssid_ = bssid; }
 void WiFiAP::set_bssid(optional<bssid_t> bssid) { this->bssid_ = bssid; }
diff --git a/esphome/components/wifi/wifi_component.h b/esphome/components/wifi/wifi_component.h
index 6cbdc51caf..be5095105c 100644
--- a/esphome/components/wifi/wifi_component.h
+++ b/esphome/components/wifi/wifi_component.h
@@ -48,6 +48,11 @@ struct SavedWifiSettings {
   char password[65];
 } PACKED;  // NOLINT
 
+struct SavedWifiFastConnectSettings {
+  uint8_t bssid[6];
+  uint8_t channel;
+} PACKED;  // NOLINT
+
 enum WiFiComponentState {
   /** Nothing has been initialized yet. Internal AP, if configured, is disabled at this point. */
   WIFI_COMPONENT_STATE_OFF = 0,
@@ -334,6 +339,9 @@ class WiFiComponent : public Component {
   bool is_captive_portal_active_();
   bool is_esp32_improv_active_();
 
+  void load_fast_connect_settings_();
+  void save_fast_connect_settings_();
+
 #ifdef USE_ESP8266
   static void wifi_event_callback(System_Event_t *event);
   void wifi_scan_done_callback_(void *arg, STATUS status);
@@ -381,6 +389,7 @@ class WiFiComponent : public Component {
   optional<float> output_power_;
   bool passive_scan_{false};
   ESPPreferenceObject pref_;
+  ESPPreferenceObject fast_connect_pref_;
   bool has_saved_wifi_settings_{false};
 #ifdef USE_WIFI_11KV_SUPPORT
   bool btm_{false};