Ensure services are removed from the registry when calling unregister_all_services (#644)

bdraco · web-flow · commit cf0b5b9e2cfa · 2021-06-15T08:45:30.000-10:00
- There was a race condition where a query could be answered for a service
  in the registry while goodbye packets which could result a fresh record
  being broadcast after the goodbye if a query came in at just the right
  time. To avoid this, we now remove the services from the registry right
  after we generate the goodbye packet
diff --git a/tests/test_core.py b/tests/test_core.py
@@ -317,3 +317,32 @@ def test_invalid_packets_ignored_and_does_not_cause_loop_exception():
     time.sleep(0.2)
     zc.close()
     assert zc.cache.get(entry) is not None
+
+
+def test_goodbye_all_services():
+    """Verify generating the goodbye query does not change with time."""
+    zc = Zeroconf(interfaces=['127.0.0.1'])
+    out = zc.generate_unregister_all_services()
+    assert out is None
+    type_ = "_http._tcp.local."
+    registration_name = "xxxyyy.%s" % type_
+    desc = {'path': '/~paulsm/'}
+    info = r.ServiceInfo(
+        type_, registration_name, 80, 0, 0, desc, "ash-2.local.", addresses=[socket.inet_aton("10.0.1.2")]
+    )
+    zc.registry.add(info)
+    out = zc.generate_unregister_all_services()
+    assert out is not None
+    first_packet = out.packets()
+    zc.registry.add(info)
+    out2 = zc.generate_unregister_all_services()
+    assert out2 is not None
+    second_packet = out.packets()
+    assert second_packet == first_packet
+
+    # Verify the registery is empty
+    out3 = zc.generate_unregister_all_services()
+    assert out3 is None
+    assert zc.registry.get_service_infos() == []
+
+    zc.close()
diff --git a/tests/test_init.py b/tests/test_init.py
@@ -184,8 +184,21 @@ def verify_name_change(self, zc, type_, name, number_hosts):
         # verify no name conflict https://tools.ietf.org/html/rfc6762#section-6.6
         zc.register_service(info_service, cooperating_responders=True)
 
-        zc.register_service(info_service, allow_name_change=True)
-        assert info_service.name.split('.')[0] == '%s-%d' % (name, number_hosts + 1)
+        # Create a new object since allow_name_change will mutate the
+        # original object and then we will have the wrong service
+        # in the registry
+        info_service2 = ServiceInfo(
+            type_,
+            '%s.%s' % (name, type_),
+            80,
+            0,
+            0,
+            desc,
+            "ash-2.local.",
+            addresses=[socket.inet_aton("10.0.1.2")],
+        )
+        zc.register_service(info_service2, allow_name_change=True)
+        assert info_service2.name.split('.')[0] == '%s-%d' % (name, number_hosts + 1)
 
     def generate_many_hosts(self, zc, type_, name, number_hosts):
         records_per_server = 2
diff --git a/zeroconf/_core.py b/zeroconf/_core.py
@@ -476,10 +476,22 @@ def unregister_service(self, info: ServiceInfo) -> None:
         self.registry.remove(info)
         self._broadcast_service(info, _UNREGISTER_TIME, 0)
 
-    def unregister_all_services(self) -> None:
-        """Unregister all registered services."""
+    def generate_unregister_all_services(self) -> Optional[DNSOutgoing]:
+        """Generate a DNSOutgoing goodbye for all services and remove them from the registry."""
         service_infos = self.registry.get_service_infos()
         if not service_infos:
+            return None
+        out = DNSOutgoing(_FLAGS_QR_RESPONSE | _FLAGS_AA)
+        for info in service_infos:
+            self._add_broadcast_answer(out, info, 0)
+        self.registry.remove(service_infos)
+        return out
+
+    def unregister_all_services(self) -> None:
+        """Unregister all registered services."""
+        # Send Goodbye packets https://datatracker.ietf.org/doc/html/rfc6762#section-10.1
+        out = self.generate_unregister_all_services()
+        if not out:
             return
         now = current_time_millis()
         next_time = now
@@ -489,9 +501,6 @@ def unregister_all_services(self) -> None:
                 self.wait(next_time - now)
                 now = current_time_millis()
                 continue
-            out = DNSOutgoing(_FLAGS_QR_RESPONSE | _FLAGS_AA)
-            for info in service_infos:
-                self._add_broadcast_answer(out, info, 0)
             self.send(out)
             i += 1
             next_time += _UNREGISTER_TIME
@@ -604,8 +613,8 @@ def close(self) -> None:
         if self._GLOBAL_DONE:
             return
         # remove service listeners
-        self.remove_all_service_listeners()
         self.unregister_all_services()
+        self.remove_all_service_listeners()
         self._GLOBAL_DONE = True
         self.engine.close()
         # shutdown the rest
diff --git a/zeroconf/_services/registry.py b/zeroconf/_services/registry.py
@@ -21,7 +21,7 @@
 """
 
 import threading
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Union
 
 
 from .._exceptions import ServiceNameAlreadyRegistered
@@ -47,21 +47,21 @@ def __init__(
 
     def add(self, info: ServiceInfo) -> None:
         """Add a new service to the registry."""
-
         with self._lock:
             self._add(info)
 
-    def remove(self, info: ServiceInfo) -> None:
+    def remove(self, info: Union[List[ServiceInfo], ServiceInfo]) -> None:
         """Remove a new service from the registry."""
+        infos = info if isinstance(info, list) else [info]
 
         with self._lock:
-            self._remove(info)
+            self._remove(infos)
 
     def update(self, info: ServiceInfo) -> None:
         """Update new service in the registry."""
 
         with self._lock:
-            self._remove(info)
+            self._remove([info])
             self._add(info)
 
     def get_service_infos(self) -> List[ServiceInfo]:
@@ -103,9 +103,10 @@ def _add(self, info: ServiceInfo) -> None:
         self.types.setdefault(info.type.lower(), []).append(info.key)
         self.servers.setdefault(info.server_key, []).append(info.key)
 
-    def _remove(self, info: ServiceInfo) -> None:
-        """Remove a service under the lock."""
-        old_service_info = self._services[info.key]
-        self.types[old_service_info.type.lower()].remove(info.key)
-        self.servers[old_service_info.server_key].remove(info.key)
-        del self._services[info.key]
+    def _remove(self, infos: List[ServiceInfo]) -> None:
+        """Remove a services under the lock."""
+        for info in infos:
+            old_service_info = self._services[info.key]
+            self.types[old_service_info.type.lower()].remove(info.key)
+            self.servers[old_service_info.server_key].remove(info.key)
+            del self._services[info.key]
