Avoid including additionals when the answer is suppressed by known-answer supression

bdraco · bdraco · commit 575097e2b02c · 2021-06-14T12:00:47.000-10:00
- If the PTR or SRV record was suppressed we should not include the additionals per https://tools.ietf.org/html/rfc6763#section-12.1
diff --git a/tests/test_handlers.py b/tests/test_handlers.py
@@ -12,7 +12,7 @@
 import unittest.mock
 
 import zeroconf as r
-from zeroconf import ServiceInfo, Zeroconf
+from zeroconf import ServiceInfo, Zeroconf, current_time_millis
 from zeroconf import const
 
 from . import _clear_cache
@@ -292,3 +292,113 @@ def test_unicast_response():
     # unregister
     zc.unregister_service(info)
     zc.close()
+
+
+def test_known_answer_supression():
+    zc = Zeroconf(interfaces=['127.0.0.1'])
+    type_ = "_knownservice._tcp.local."
+    name = "knownname"
+    registration_name = "%s.%s" % (name, type_)
+    desc = {'path': '/~paulsm/'}
+    server_name = "ash-2.local."
+    info = ServiceInfo(
+        type_, registration_name, 80, 0, 0, desc, server_name, addresses=[socket.inet_aton("10.0.1.2")]
+    )
+    zc.register_service(info)
+
+    now = current_time_millis()
+
+    # Test PTR supression
+    generated = r.DNSOutgoing(const._FLAGS_QR_QUERY)
+    question = r.DNSQuestion(type_, const._TYPE_PTR, const._CLASS_IN)
+    generated.add_question(question)
+    packets = generated.packets()
+    unicast_out, multicast_out = zc.query_handler.response(
+        r.DNSIncoming(packets[0]), "1.2.3.4", const._MDNS_PORT
+    )
+    assert unicast_out is None
+    assert multicast_out is not None and multicast_out.answers
+
+    generated = r.DNSOutgoing(const._FLAGS_QR_QUERY)
+    question = r.DNSQuestion(type_, const._TYPE_PTR, const._CLASS_IN)
+    generated.add_question(question)
+    generated.add_answer_at_time(info.dns_pointer(), now)
+    packets = generated.packets()
+    unicast_out, multicast_out = zc.query_handler.response(
+        r.DNSIncoming(packets[0]), "1.2.3.4", const._MDNS_PORT
+    )
+    assert unicast_out is None
+    # If the answer is suppressed, the additional should be suppresed as well
+    assert not multicast_out or not multicast_out.answers
+
+    # Test A supression
+    generated = r.DNSOutgoing(const._FLAGS_QR_QUERY)
+    question = r.DNSQuestion(server_name, const._TYPE_A, const._CLASS_IN)
+    generated.add_question(question)
+    packets = generated.packets()
+    unicast_out, multicast_out = zc.query_handler.response(
+        r.DNSIncoming(packets[0]), "1.2.3.4", const._MDNS_PORT
+    )
+    assert unicast_out is None
+    assert multicast_out is not None and multicast_out.answers
+
+    generated = r.DNSOutgoing(const._FLAGS_QR_QUERY)
+    question = r.DNSQuestion(server_name, const._TYPE_A, const._CLASS_IN)
+    generated.add_question(question)
+    for dns_address in info.dns_addresses():
+        generated.add_answer_at_time(dns_address, now)
+    packets = generated.packets()
+    unicast_out, multicast_out = zc.query_handler.response(
+        r.DNSIncoming(packets[0]), "1.2.3.4", const._MDNS_PORT
+    )
+    assert unicast_out is None
+    assert not multicast_out or not multicast_out.answers
+
+    # Test SRV supression
+    generated = r.DNSOutgoing(const._FLAGS_QR_QUERY)
+    question = r.DNSQuestion(registration_name, const._TYPE_SRV, const._CLASS_IN)
+    generated.add_question(question)
+    packets = generated.packets()
+    unicast_out, multicast_out = zc.query_handler.response(
+        r.DNSIncoming(packets[0]), "1.2.3.4", const._MDNS_PORT
+    )
+    assert unicast_out is None
+    assert multicast_out is not None and multicast_out.answers
+
+    generated = r.DNSOutgoing(const._FLAGS_QR_QUERY)
+    question = r.DNSQuestion(registration_name, const._TYPE_SRV, const._CLASS_IN)
+    generated.add_question(question)
+    generated.add_answer_at_time(info.dns_service(), now)
+    packets = generated.packets()
+    unicast_out, multicast_out = zc.query_handler.response(
+        r.DNSIncoming(packets[0]), "1.2.3.4", const._MDNS_PORT
+    )
+    assert unicast_out is None
+    # If the answer is suppressed, the additional should be suppresed as well
+    assert not multicast_out or not multicast_out.answers
+
+    # Test TXT supression
+    generated = r.DNSOutgoing(const._FLAGS_QR_QUERY)
+    question = r.DNSQuestion(registration_name, const._TYPE_TXT, const._CLASS_IN)
+    generated.add_question(question)
+    packets = generated.packets()
+    unicast_out, multicast_out = zc.query_handler.response(
+        r.DNSIncoming(packets[0]), "1.2.3.4", const._MDNS_PORT
+    )
+    assert unicast_out is None
+    assert multicast_out is not None and multicast_out.answers
+
+    generated = r.DNSOutgoing(const._FLAGS_QR_QUERY)
+    question = r.DNSQuestion(registration_name, const._TYPE_TXT, const._CLASS_IN)
+    generated.add_question(question)
+    generated.add_answer_at_time(info.dns_text(), now)
+    packets = generated.packets()
+    unicast_out, multicast_out = zc.query_handler.response(
+        r.DNSIncoming(packets[0]), "1.2.3.4", const._MDNS_PORT
+    )
+    assert unicast_out is None
+    assert not multicast_out or not multicast_out.answers
+
+    # unregister
+    zc.unregister_service(info)
+    zc.close()
diff --git a/zeroconf/_handlers.py b/zeroconf/_handlers.py
@@ -75,6 +75,9 @@ def _answer_service_type_enumeration_query(self, msg: DNSIncoming, out: DNSOutgo
     def _answer_ptr_query(self, msg: DNSIncoming, out: DNSOutgoing, question: DNSQuestion) -> None:
         """Answer a PTR query."""
         for service in self.registry.get_infos_type(question.name):
+            dns_pointer = service.dns_pointer()
+            if dns_pointer.suppressed_by(msg):
+                continue
             out.add_answer(msg, service.dns_pointer())
             # Add recommended additional answers according to
             # https://tools.ietf.org/html/rfc6763#section-12.1.
@@ -103,8 +106,12 @@ def _answer_non_ptr_query(self, msg: DNSIncoming, out: DNSOutgoing, question: DN
         if question.type in (_TYPE_TXT, _TYPE_ANY):
             out.add_answer(msg, service.dns_text())
         if question.type == _TYPE_SRV:
-            for dns_address in service.dns_addresses():
-                out.add_additional_answer(dns_address)
+            dns_service = service.dns_service()
+            if not dns_service.suppressed_by(msg):
+                # Add recommended additional answers according to
+                # https://datatracker.ietf.org/doc/html/rfc6763#section-12.2
+                for dns_address in service.dns_addresses():
+                    out.add_additional_answer(dns_address)
 
     def response(  # pylint: disable=unused-argument
         self, msg: DNSIncoming, addr: Optional[str], port: int
