From 529ff4bd526c258069e388dcdd16bfc1b5c43535 Mon Sep 17 00:00:00 2001
From: Keith Burzinski <kbx81x@gmail.com>
Date: Thu, 26 Sep 2024 17:24:18 -0500
Subject: [PATCH] [wifi] Use custom MAC address if programmed (#7498)

---
 .../wifi/wifi_component_esp_idf.cpp           |  7 ++-
 esphome/core/helpers.cpp                      | 45 ++++++++++++++++++-
 esphome/core/helpers.h                        |  8 ++++
 3 files changed, 57 insertions(+), 3 deletions(-)

diff --git a/esphome/components/wifi/wifi_component_esp_idf.cpp b/esphome/components/wifi/wifi_component_esp_idf.cpp
index 6008acb95d..c430d160f2 100644
--- a/esphome/components/wifi/wifi_component_esp_idf.cpp
+++ b/esphome/components/wifi/wifi_component_esp_idf.cpp
@@ -130,11 +130,16 @@ void event_handler(void *arg, esp_event_base_t event_base, int32_t event_id, voi
 }
 
 void WiFiComponent::wifi_pre_setup_() {
-#ifdef USE_ESP32_IGNORE_EFUSE_MAC_CRC
   uint8_t mac[6];
+#ifdef USE_ESP32_IGNORE_EFUSE_MAC_CRC
   get_mac_address_raw(mac);
   set_mac_address(mac);
   ESP_LOGV(TAG, "Use EFuse MAC without checking CRC: %s", get_mac_address_pretty().c_str());
+#else
+  if (has_custom_mac_address()) {
+    get_mac_address_raw(mac);
+    set_mac_address(mac);
+  }
 #endif
   esp_err_t err = esp_netif_init();
   if (err != ERR_OK) {
diff --git a/esphome/core/helpers.cpp b/esphome/core/helpers.cpp
index e75b06ccd3..2e99b0df70 100644
--- a/esphome/core/helpers.cpp
+++ b/esphome/core/helpers.cpp
@@ -671,9 +671,17 @@ void get_mac_address_raw(uint8_t *mac) {  // NOLINT(readability-non-const-parame
   // match the CRC that goes along with it. For those devices, this
   // work-around reads and uses the MAC address as-is from EFuse,
   // without doing the CRC check.
-  esp_efuse_read_field_blob(ESP_EFUSE_MAC_FACTORY, mac, 48);
+  if (has_custom_mac_address()) {
+    esp_efuse_read_field_blob(ESP_EFUSE_MAC_CUSTOM, mac, 48);
+  } else {
+    esp_efuse_read_field_blob(ESP_EFUSE_MAC_FACTORY, mac, 48);
+  }
 #else
-  esp_efuse_mac_get_default(mac);
+  if (has_custom_mac_address()) {
+    esp_efuse_mac_get_custom(mac);
+  } else {
+    esp_efuse_mac_get_default(mac);
+  }
 #endif
 #elif defined(USE_ESP8266)
   wifi_get_macaddr(STATION_IF, mac);
@@ -685,20 +693,53 @@ void get_mac_address_raw(uint8_t *mac) {  // NOLINT(readability-non-const-parame
 // this should be an error, but that messes with CI checks. #error No mac address method defined
 #endif
 }
+
 std::string get_mac_address() {
   uint8_t mac[6];
   get_mac_address_raw(mac);
   return str_snprintf("%02x%02x%02x%02x%02x%02x", 12, mac[0], mac[1], mac[2], mac[3], mac[4], mac[5]);
 }
+
 std::string get_mac_address_pretty() {
   uint8_t mac[6];
   get_mac_address_raw(mac);
   return str_snprintf("%02X:%02X:%02X:%02X:%02X:%02X", 17, mac[0], mac[1], mac[2], mac[3], mac[4], mac[5]);
 }
+
 #ifdef USE_ESP32
 void set_mac_address(uint8_t *mac) { esp_base_mac_addr_set(mac); }
 #endif
 
+bool has_custom_mac_address() {
+#ifdef USE_ESP32
+  uint8_t mac[6];
+#if defined(CONFIG_SOC_IEEE802154_SUPPORTED) || defined(USE_ESP32_IGNORE_EFUSE_MAC_CRC)
+  return (esp_efuse_read_field_blob(ESP_EFUSE_MAC_CUSTOM, mac, 48) == ESP_OK) && mac_address_is_valid(mac);
+#else
+  return (esp_efuse_mac_get_custom(mac) == ESP_OK) && mac_address_is_valid(mac);
+#endif
+#else
+  return false;
+#endif
+}
+
+bool mac_address_is_valid(const uint8_t *mac) {
+  bool is_all_zeros = true;
+  bool is_all_ones = true;
+
+  for (uint8_t i = 0; i < 6; i++) {
+    if (mac[i] != 0) {
+      is_all_zeros = false;
+    }
+  }
+  for (uint8_t i = 0; i < 6; i++) {
+    if (mac[i] != 0xFF) {
+      is_all_ones = false;
+    }
+  }
+  return !(is_all_zeros || is_all_ones);
+}
+
 void delay_microseconds_safe(uint32_t us) {  // avoids CPU locks that could trigger WDT or affect WiFi/BT stability
   uint32_t start = micros();
 
diff --git a/esphome/core/helpers.h b/esphome/core/helpers.h
index 3e6fe9433e..7df4b84230 100644
--- a/esphome/core/helpers.h
+++ b/esphome/core/helpers.h
@@ -635,6 +635,14 @@ std::string get_mac_address_pretty();
 void set_mac_address(uint8_t *mac);
 #endif
 
+/// Check if a custom MAC address is set (ESP32 & variants)
+/// @return True if a custom MAC address is set (ESP32 & variants), else false
+bool has_custom_mac_address();
+
+/// Check if the MAC address is not all zeros or all ones
+/// @return True if MAC is valid, else false
+bool mac_address_is_valid(const uint8_t *mac);
+
 /// Delay for the given amount of microseconds, possibly yielding to other processes during the wait.
 void delay_microseconds_safe(uint32_t us);