python-zeroconf · bdraco · Jun 14, 2021 · Jun 14, 2021 · Jun 14, 2021 · Jun 14, 2021
diff --git a/tests/test_handlers.py b/tests/test_handlers.py
@@ -360,6 +360,107 @@ def test_unicast_response():
     zc.close()
 
 
+def test_qu_response():
+    """Handle multicast incoming with the QU bit set."""
+    # instantiate a zeroconf instance
+    zc = Zeroconf(interfaces=['127.0.0.1'])
+
+    # service definition
+    type_ = "_test-srvc-type._tcp.local."
+    other_type_ = "_notthesame._tcp.local."
+    name = "xxxyyy"
+    registration_name = "%s.%s" % (name, type_)
+    registration_name2 = "%s.%s" % (name, other_type_)
+    desc = {'path': '/~paulsm/'}
+    info = ServiceInfo(
+        type_, registration_name, 80, 0, 0, desc, "ash-2.local.", addresses=[socket.inet_aton("10.0.1.2")]
+    )
+    info2 = ServiceInfo(
+        other_type_,
+        registration_name2,
+        80,
+        0,
+        0,
+        desc,
+        "ash-other.local.",
+        addresses=[socket.inet_aton("10.0.4.2")],
+    )
+    # register
+    zc.register_service(info)
+
+    def _validate_complete_response(query, out):
+        assert out.id == query.id
+        has_srv = has_txt = has_a = False
+        nbr_additionals = 0
+        nbr_answers = len(out.answers)
+        nbr_authorities = len(out.authorities)
+        for answer in out.additionals:
+            nbr_additionals += 1
+            if answer.type == const._TYPE_SRV:
+                has_srv = True
+            elif answer.type == const._TYPE_TXT:
+                has_txt = True
+            elif answer.type == const._TYPE_A:
+                has_a = True
+        assert nbr_answers == 1 and nbr_additionals == 3 and nbr_authorities == 0
+        assert has_srv and has_txt and has_a
+
+    # With QU should respond to only unicast when the answer has been recently multicast
+    query = r.DNSOutgoing(const._FLAGS_QR_QUERY | const._FLAGS_AA)
+    question = r.DNSQuestion(info.type, const._TYPE_PTR, const._CLASS_IN)
+    question.unique = True  # Set the QU bit
+    assert question.unicast is True
+    query.add_question(question)
+
+    unicast_out, multicast_out = zc.query_handler.response(
+        r.DNSIncoming(query.packets()[0]), "1.2.3.4", const._MDNS_PORT
+    )
+    assert multicast_out is None
+    _validate_complete_response(query, unicast_out)
+
+    _clear_cache(zc)
+    # With QU should respond to only multicast since the response hasn't been seen since 75% of the ttl
+    query = r.DNSOutgoing(const._FLAGS_QR_QUERY | const._FLAGS_AA)
+    question = r.DNSQuestion(info.type, const._TYPE_PTR, const._CLASS_IN)
+    question.unique = True  # Set the QU bit
+    assert question.unicast is True
+    query.add_question(question)
+    unicast_out, multicast_out = zc.query_handler.response(
+        r.DNSIncoming(query.packets()[0]), "1.2.3.4", const._MDNS_PORT
+    )
+    assert unicast_out is None
+    _validate_complete_response(query, multicast_out)
+
+    # With QU set and an authorative answer (probe) should respond to both unitcast and multicast since the response hasn't been seen since 75% of the ttl
+    query = r.DNSOutgoing(const._FLAGS_QR_QUERY | const._FLAGS_AA)
+    question = r.DNSQuestion(info.type, const._TYPE_PTR, const._CLASS_IN)
+    question.unique = True  # Set the QU bit
+    assert question.unicast is True
+    query.add_question(question)
+    query.add_authorative_answer(info2.dns_pointer())
+    unicast_out, multicast_out = zc.query_handler.response(
+        r.DNSIncoming(query.packets()[0]), "1.2.3.4", const._MDNS_PORT
+    )
+    _validate_complete_response(query, unicast_out)
+    _validate_complete_response(query, multicast_out)
+
+    _inject_response(zc, r.DNSIncoming(multicast_out.packets()[0]))
+    # With the cache repopulated; should respond to only unicast when the answer has been recently multicast
+    query = r.DNSOutgoing(const._FLAGS_QR_QUERY | const._FLAGS_AA)
+    question = r.DNSQuestion(info.type, const._TYPE_PTR, const._CLASS_IN)
+    question.unique = True  # Set the QU bit
+    assert question.unicast is True
+    query.add_question(question)
+    unicast_out, multicast_out = zc.query_handler.response(
+        r.DNSIncoming(query.packets()[0]), "1.2.3.4", const._MDNS_PORT
+    )
+    assert multicast_out is None
+    _validate_complete_response(query, unicast_out)
+    # unregister
+    zc.unregister_service(info)
+    zc.close()
+
+
 def test_known_answer_supression():
     zc = Zeroconf(interfaces=['127.0.0.1'])
     type_ = "_knownservice._tcp.local."

diff --git a/zeroconf/_handlers.py b/zeroconf/_handlers.py
@@ -76,6 +76,25 @@ def __init__(self, cache: DNSCache, msg: DNSIncoming, ucast_source: bool) -> Non
         self._ucast: _RecordSetType = {RecordSetKeys.Answers: set(), RecordSetKeys.Additionals: set()}
         self._mcast: _RecordSetType = {RecordSetKeys.Answers: set(), RecordSetKeys.Additionals: set()}
 
+    def add_qu_question_response(
+        self,
+        answers: Set[DNSRecord],
+        additionals: Set[DNSRecord],
+    ) -> None:
+        """Generate a response to a multicast QU query."""
+        self._add_qu_question_response_to_target(answers, RecordSetKeys.Answers)
+        self._add_qu_question_response_to_target(additionals, RecordSetKeys.Additionals)
+
+    def _add_qu_question_response_to_target(self, target: Set[DNSRecord], answer_type: RecordSetKeys) -> None:
+        """Add part of the QU response."""
+        for record in target:
+            if self._is_probe:
+                self._ucast[answer_type].add(record)
+            if not self._has_mcast_within_one_quarter_ttl(record):
+                self._mcast[answer_type].add(record)
+            elif not self._is_probe:
+                self._ucast[answer_type].add(record)
+
     def add_ucast_question_response(self, answers: Set[DNSRecord], additionals: Set[DNSRecord]) -> None:
         """Generate a response to a unicast query."""
         self._ucast[RecordSetKeys.Answers].update(answers)
@@ -119,8 +138,23 @@ def _construct_outgoing_from_record_set(
             out.add_answer_at_time(answer, 0)
         for additional in rrset[RecordSetKeys.Additionals]:
             out.add_additional_answer(additional)
+
         return out
 
+    def _has_mcast_within_one_quarter_ttl(self, record: DNSRecord) -> bool:
+        """Check to see if a record has been mcasted recently.
+
+        https://datatracker.ietf.org/doc/html/rfc6762#section-5.4
+        When receiving a question with the unicast-response bit set, a
+        responder SHOULD usually respond with a unicast packet directed back
+        to the querier.  However, if the responder has not multicast that
+        record recently (within one quarter of its TTL), then the responder
+        SHOULD instead multicast the response so as to keep all the peer
+        caches up to date
+        """
+        maybe_entry = self._cache.get(record)
+        return bool(maybe_entry and maybe_entry.is_recent(self._now))
+
     def _suppress_mcasts_from_last_second(self, records: Set[DNSRecord]) -> None:
         """Remove any records that were already sent in the last second."""
         records -= set(record for record in records if self._has_mcast_record_in_last_second(record))
@@ -226,11 +260,14 @@ def response(  # pylint: disable=unused-argument
 
         for question in msg.questions:
             all_answers = self._answer_any_question(msg, question)
-            if ucast_source:
-                query_res.add_ucast_question_response(*all_answers)
-            # We always multicast as well even if its a unicast
-            # source as long as we haven't done it recently (75% of ttl)
-            query_res.add_mcast_question_response(*all_answers)
+            if not ucast_source and question.unicast:
+                query_res.add_qu_question_response(*all_answers)
+            else:
+                if ucast_source:
+                    query_res.add_ucast_question_response(*all_answers)
+                # We always multicast as well even if its a unicast
+                # source as long as we haven't done it recently (75% of ttl)
+                query_res.add_mcast_question_response(*all_answers)
 
         return query_res.outgoing_unicast(), query_res.outgoing_multicast()