From 17c02ff252bbf548842b8e8566251a43a255b490 Mon Sep 17 00:00:00 2001
From: Ivan Grokhotkov <ivan@espressif.com>
Date: Sun, 12 Jun 2016 14:17:27 +0800
Subject: [PATCH] ESP8266mDNS: restart listening when WiFi STA is
 connected/disconnected (#1828)

---
 libraries/ESP8266mDNS/ESP8266mDNS.cpp | 32 +++++++++++++++++++++++----
 libraries/ESP8266mDNS/ESP8266mDNS.h   |  4 ++++
 2 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/libraries/ESP8266mDNS/ESP8266mDNS.cpp b/libraries/ESP8266mDNS/ESP8266mDNS.cpp
index 9b63de82e..71dca7326 100644
--- a/libraries/ESP8266mDNS/ESP8266mDNS.cpp
+++ b/libraries/ESP8266mDNS/ESP8266mDNS.cpp
@@ -143,8 +143,6 @@ MDNSResponder::~MDNSResponder() {
 }
 
 bool MDNSResponder::begin(const char* hostname){
-  // Open the MDNS socket if it isn't already open.
-
   size_t n = strlen(hostname);
   if (n > 63) { // max size for a single label.
     return false;
@@ -155,15 +153,41 @@ bool MDNSResponder::begin(const char* hostname){
   _hostName.toLowerCase();
 
   // If instance name is not already set copy hostname to instance name
-  if (_instanceName.equals("") ) _instanceName=hostname; 
+  if (_instanceName.equals("") ) _instanceName=hostname;
 
+  _gotIPHandler = WiFi.onStationModeGotIP([this](const WiFiEventStationModeGotIP& event){
+    _restart();
+  });
+
+  _disconnectedHandler = WiFi.onStationModeDisconnected([this](const WiFiEventStationModeDisconnected& event) {
+    _restart();
+  });
+
+  return _listen();
+}
+
+void MDNSResponder::_restart() {
+  if (_conn) {
+    _conn->unref();
+    _conn = nullptr;
+  }
+  _listen();
+}
+
+bool MDNSResponder::_listen() {
   // Open the MDNS socket if it isn't already open.
   if (!_conn) {
     uint32_t ourIp = _getOurIp();
     if(ourIp == 0){
+      #ifdef MDNS_DEBUG_RX
+      Serial.println("MDNS: no IP address to listen on");
+      #endif
       return false;
     }
-
+    #ifdef MDNS_DEBUG_RX
+    Serial.print("MDNS listening on IP: ");
+    Serial.println(IPAddress(ourIp));
+    #endif
     ip_addr_t ifaddr;
     ifaddr.addr = ourIp;
     ip_addr_t multicast_addr;
diff --git a/libraries/ESP8266mDNS/ESP8266mDNS.h b/libraries/ESP8266mDNS/ESP8266mDNS.h
index 7d6cb110f..a32c18337 100644
--- a/libraries/ESP8266mDNS/ESP8266mDNS.h
+++ b/libraries/ESP8266mDNS/ESP8266mDNS.h
@@ -114,6 +114,8 @@ private:
   struct MDNSQuery * _query;
   bool _newQuery;
   bool _waitingForAnswers;
+  WiFiEventHandler _disconnectedHandler;
+  WiFiEventHandler _gotIPHandler;
   
 
   uint32_t _getOurIp();
@@ -125,6 +127,8 @@ private:
   size_t advertiseServices(); // advertise all hosted services
   MDNSAnswer* _getAnswerFromIdx(int idx);
   int _getNumAnswers();
+  bool _listen();
+  void _restart();
 };
 
 extern MDNSResponder MDNS;