From ac6693f17768d519e7f2fe0206602b7dea2f2050 Mon Sep 17 00:00:00 2001
From: Jesse Hills <3060199+jesserockz@users.noreply.github.com>
Date: Thu, 9 Feb 2023 10:25:44 +1300
Subject: [PATCH] mDNS updates (#4399)

---
 esphome/components/mdns/mdns_component.h      |  2 ++
 .../mdns/{mdns_esp_idf.cpp => mdns_esp32.cpp} | 14 +++++++---
 .../components/mdns/mdns_esp32_arduino.cpp    | 26 -------------------
 esphome/components/mdns/mdns_esp8266.cpp      | 12 ++++++---
 esphome/components/mdns/mdns_rp2040.cpp       |  5 ++++
 esphome/zeroconf.py                           |  5 ++++
 6 files changed, 31 insertions(+), 33 deletions(-)
 rename esphome/components/mdns/{mdns_esp_idf.cpp => mdns_esp32.cpp} (87%)
 delete mode 100644 esphome/components/mdns/mdns_esp32_arduino.cpp

diff --git a/esphome/components/mdns/mdns_component.h b/esphome/components/mdns/mdns_component.h
index 5f9179b799..b2cb10db62 100644
--- a/esphome/components/mdns/mdns_component.h
+++ b/esphome/components/mdns/mdns_component.h
@@ -35,6 +35,8 @@ class MDNSComponent : public Component {
 
   void add_extra_service(MDNSService service) { services_extra_.push_back(std::move(service)); }
 
+  void on_shutdown() override;
+
  protected:
   std::vector<MDNSService> services_extra_{};
   std::vector<MDNSService> services_{};
diff --git a/esphome/components/mdns/mdns_esp_idf.cpp b/esphome/components/mdns/mdns_esp32.cpp
similarity index 87%
rename from esphome/components/mdns/mdns_esp_idf.cpp
rename to esphome/components/mdns/mdns_esp32.cpp
index 40d9f1d5f3..6081c96637 100644
--- a/esphome/components/mdns/mdns_esp_idf.cpp
+++ b/esphome/components/mdns/mdns_esp32.cpp
@@ -1,9 +1,10 @@
-#ifdef USE_ESP_IDF
+#ifdef USE_ESP32
 
-#include "mdns_component.h"
-#include "esphome/core/log.h"
 #include <mdns.h>
 #include <cstring>
+#include "esphome/core/hal.h"
+#include "esphome/core/log.h"
+#include "mdns_component.h"
 
 namespace esphome {
 namespace mdns {
@@ -47,7 +48,12 @@ void MDNSComponent::setup() {
   }
 }
 
+void MDNSComponent::on_shutdown() {
+  mdns_free();
+  delay(40);  // Allow the mdns packets announcing service removal to be sent
+}
+
 }  // namespace mdns
 }  // namespace esphome
 
-#endif
+#endif  // USE_ESP32
diff --git a/esphome/components/mdns/mdns_esp32_arduino.cpp b/esphome/components/mdns/mdns_esp32_arduino.cpp
deleted file mode 100644
index 6a66beef92..0000000000
--- a/esphome/components/mdns/mdns_esp32_arduino.cpp
+++ /dev/null
@@ -1,26 +0,0 @@
-#ifdef USE_ESP32_FRAMEWORK_ARDUINO
-
-#include "mdns_component.h"
-#include "esphome/core/log.h"
-#include <ESPmDNS.h>
-
-namespace esphome {
-namespace mdns {
-
-void MDNSComponent::setup() {
-  this->compile_records_();
-
-  MDNS.begin(this->hostname_.c_str());
-
-  for (const auto &service : this->services_) {
-    MDNS.addService(service.service_type.c_str(), service.proto.c_str(), service.port);
-    for (const auto &record : service.txt_records) {
-      MDNS.addServiceTxt(service.service_type.c_str(), service.proto.c_str(), record.key.c_str(), record.value.c_str());
-    }
-  }
-}
-
-}  // namespace mdns
-}  // namespace esphome
-
-#endif  // USE_ESP32_FRAMEWORK_ARDUINO
diff --git a/esphome/components/mdns/mdns_esp8266.cpp b/esphome/components/mdns/mdns_esp8266.cpp
index b91f60cb6b..4ccfe42baa 100644
--- a/esphome/components/mdns/mdns_esp8266.cpp
+++ b/esphome/components/mdns/mdns_esp8266.cpp
@@ -1,10 +1,11 @@
 #if defined(USE_ESP8266) && defined(USE_ARDUINO)
 
-#include "mdns_component.h"
-#include "esphome/core/log.h"
+#include <ESP8266mDNS.h>
 #include "esphome/components/network/ip_address.h"
 #include "esphome/components/network/util.h"
-#include <ESP8266mDNS.h>
+#include "esphome/core/hal.h"
+#include "esphome/core/log.h"
+#include "mdns_component.h"
 
 namespace esphome {
 namespace mdns {
@@ -37,6 +38,11 @@ void MDNSComponent::setup() {
 
 void MDNSComponent::loop() { MDNS.update(); }
 
+void MDNSComponent::on_shutdown() {
+  MDNS.close();
+  delay(10);
+}
+
 }  // namespace mdns
 }  // namespace esphome
 
diff --git a/esphome/components/mdns/mdns_rp2040.cpp b/esphome/components/mdns/mdns_rp2040.cpp
index e443a78c48..b30a3a4ee7 100644
--- a/esphome/components/mdns/mdns_rp2040.cpp
+++ b/esphome/components/mdns/mdns_rp2040.cpp
@@ -38,6 +38,11 @@ void MDNSComponent::setup() {
 
 void MDNSComponent::loop() { MDNS.update(); }
 
+void MDNSComponent::on_shutdown() {
+  MDNS.close();
+  delay(40);
+}
+
 }  // namespace mdns
 }  // namespace esphome
 
diff --git a/esphome/zeroconf.py b/esphome/zeroconf.py
index 1ba8397e14..b0dddfd152 100644
--- a/esphome/zeroconf.py
+++ b/esphome/zeroconf.py
@@ -157,6 +157,11 @@ class DashboardImportDiscovery:
             return
         if state_change == ServiceStateChange.Removed:
             self.import_state.pop(name, None)
+            return
+
+        if state_change == ServiceStateChange.Updated and name not in self.import_state:
+            # Ignore updates for devices that are not in the import state
+            return
 
         info = zeroconf.get_service_info(service_type, name)
         _LOGGER.debug("-> resolved info: %s", info)