Add additional recommended records to PTR responses (#184)

inwotep · jstasiak · commit ea6426547f79 · 2019-08-08T21:01:23.000+02:00
RFC6763 indicates a server should include the SRV/TXT/A/AAAA records
when responding to a PTR record request.  This optimization ensures
the client doesn't have to then query for these additional records.

It has been observed that when multiple Windows 10 machines are monitoring
for the same service, this unoptimized response to the PTR record
request can cause extremely high CPU usage in both the DHCP Client
&amp; Device Association service (I suspect due to all clients having to
then sending/receiving the additional queries/responses).
diff --git a/test_zeroconf.py b/test_zeroconf.py
@@ -621,7 +621,7 @@ def send(out, addr=r._MDNS_ADDR, port=r._MDNS_PORT):
         query.add_question(r.DNSQuestion(info.name, r._TYPE_TXT, r._CLASS_IN))
         query.add_question(r.DNSQuestion(info.server, r._TYPE_A, r._CLASS_IN))
         zc.handle_query(r.DNSIncoming(query.packet()), r._MDNS_ADDR, r._MDNS_PORT)
-        assert nbr_answers == 4 and nbr_additionals == 1 and nbr_authorities == 0
+        assert nbr_answers == 4 and nbr_additionals == 4 and nbr_authorities == 0
         nbr_answers = nbr_additionals = nbr_authorities = 0
 
         # unregister
@@ -644,7 +644,7 @@ def send(out, addr=r._MDNS_ADDR, port=r._MDNS_PORT):
         query.add_question(r.DNSQuestion(info.name, r._TYPE_TXT, r._CLASS_IN))
         query.add_question(r.DNSQuestion(info.server, r._TYPE_A, r._CLASS_IN))
         zc.handle_query(r.DNSIncoming(query.packet()), r._MDNS_ADDR, r._MDNS_PORT)
-        assert nbr_answers == 4 and nbr_additionals == 1 and nbr_authorities == 0
+        assert nbr_answers == 4 and nbr_additionals == 4 and nbr_authorities == 0
         nbr_answers = nbr_additionals = nbr_authorities = 0
 
         # unregister
@@ -1036,3 +1036,58 @@ def test_multiple_addresses():
     )
 
     assert info.addresses == [address, address]
+
+
+def test_ptr_optimization():
+
+    # instantiate a zeroconf instance
+    zc = Zeroconf(interfaces=['127.0.0.1'])
+
+    # service definition
+    type_ = "_test-srvc-type._tcp.local."
+    name = "xxxyyy"
+    registration_name = "%s.%s" % (name, type_)
+
+    desc = {'path': '/~paulsm/'}
+    info = ServiceInfo(type_, registration_name, socket.inet_aton("10.0.1.2"), 80, 0, 0, desc, "ash-2.local.")
+
+    # we are going to monkey patch the zeroconf send to check packet sizes
+    old_send = zc.send
+
+    nbr_answers = nbr_additionals = nbr_authorities = 0
+    has_srv = has_txt = has_a = False
+
+    def send(out, addr=r._MDNS_ADDR, port=r._MDNS_PORT):
+        """Sends an outgoing packet."""
+        nonlocal nbr_answers, nbr_additionals, nbr_authorities
+        nonlocal has_srv, has_txt, has_a
+
+        nbr_answers += len(out.answers)
+        nbr_authorities += len(out.authorities)
+        for answer in out.additionals:
+            nbr_additionals += 1
+            if answer.type == r._TYPE_SRV:
+                has_srv = True
+            elif answer.type == r._TYPE_TXT:
+                has_txt = True
+            elif answer.type == r._TYPE_A:
+                has_a = True
+
+        old_send(out, addr=addr, port=port)
+
+    # monkey patch the zeroconf send
+    setattr(zc, "send", send)
+
+    # register
+    zc.register_service(info)
+    nbr_answers = nbr_additionals = nbr_authorities = 0
+
+    # query
+    query = r.DNSOutgoing(r._FLAGS_QR_QUERY | r._FLAGS_AA)
+    query.add_question(r.DNSQuestion(info.type, r._TYPE_PTR, r._CLASS_IN))
+    zc.handle_query(r.DNSIncoming(query.packet()), r._MDNS_ADDR, r._MDNS_PORT)
+    assert nbr_answers == 1 and nbr_additionals == 3 and nbr_authorities == 0
+    assert has_srv and has_txt and has_a
+
+    # unregister
+    zc.unregister_service(info)
diff --git a/zeroconf.py b/zeroconf.py
@@ -2338,6 +2338,40 @@ def handle_query(self, msg: DNSIncoming, addr: Optional[str], port: int) -> None
                             msg,
                             DNSPointer(service.type, _TYPE_PTR, _CLASS_IN, service.other_ttl, service.name),
                         )
+
+                        # Add recommended additional answers according to
+                        # https://tools.ietf.org/html/rfc6763#section-12.1.
+                        out.add_additional_answer(
+                            DNSService(
+                                service.name,
+                                _TYPE_SRV,
+                                _CLASS_IN | _CLASS_UNIQUE,
+                                service.host_ttl,
+                                service.priority,
+                                service.weight,
+                                service.port,
+                                service.server,
+                            )
+                        )
+                        out.add_additional_answer(
+                            DNSText(
+                                service.name,
+                                _TYPE_TXT,
+                                _CLASS_IN | _CLASS_UNIQUE,
+                                service.other_ttl,
+                                service.text,
+                            )
+                        )
+                        for address in service.addresses:
+                            out.add_additional_answer(
+                                DNSAddress(
+                                    service.server,
+                                    _TYPE_A,
+                                    _CLASS_IN | _CLASS_UNIQUE,
+                                    service.host_ttl,
+                                    address,
+                                )
+                            )
             else:
                 try:
                     if out is None:
