python-zeroconf · bdraco · Dec 20, 2022 · Dec 20, 2022 · Dec 20, 2022 · Dec 20, 2022
diff --git a/bench/outgoing.py b/bench/outgoing.py
@@ -0,0 +1,168 @@
+"""Benchmark for DNSOutgoing."""
+import socket
+import timeit
+
+from zeroconf import DNSAddress, DNSOutgoing, DNSService, DNSText, const
+from zeroconf._protocol.outgoing import State
+
+
+def generate_packets() -> DNSOutgoing:
+    out = DNSOutgoing(const._FLAGS_QR_RESPONSE | const._FLAGS_AA)
+    address = socket.inet_pton(socket.AF_INET, "192.168.208.5")
+
+    additionals = [
+        {
+            "name": "HASS Bridge ZJWH FF5137._hap._tcp.local.",
+            "address": address,
+            "port": 51832,
+            "text": b"\x13md=HASS Bridge"
+            b" ZJWH\x06pv=1.0\x14id=01:6B:30:FF:51:37\x05c#=12\x04s#=1\x04ff=0\x04"
+            b"ci=2\x04sf=0\x0bsh=L0m/aQ==",
+        },
+        {
+            "name": "HASS Bridge 3K9A C2582A._hap._tcp.local.",
+            "address": address,
+            "port": 51834,
+            "text": b"\x13md=HASS Bridge"
+            b" 3K9A\x06pv=1.0\x14id=E2:AA:5B:C2:58:2A\x05c#=12\x04s#=1\x04ff=0\x04"
+            b"ci=2\x04sf=0\x0bsh=b2CnzQ==",
+        },
+        {
+            "name": "Master Bed TV CEDB27._hap._tcp.local.",
+            "address": address,
+            "port": 51830,
+            "text": b"\x10md=Master Bed"
+            b" TV\x06pv=1.0\x14id=9E:B7:44:CE:DB:27\x05c#=18\x04s#=1\x04ff=0\x05"
+            b"ci=31\x04sf=0\x0bsh=CVj1kw==",
+        },
+        {
+            "name": "Living Room TV 921B77._hap._tcp.local.",
+            "address": address,
+            "port": 51833,
+            "text": b"\x11md=Living Room"
+            b" TV\x06pv=1.0\x14id=11:61:E7:92:1B:77\x05c#=17\x04s#=1\x04ff=0\x05"
+            b"ci=31\x04sf=0\x0bsh=qU77SQ==",
+        },
+        {
+            "name": "HASS Bridge ZC8X FF413D._hap._tcp.local.",
+            "address": address,
+            "port": 51829,
+            "text": b"\x13md=HASS Bridge"
+            b" ZC8X\x06pv=1.0\x14id=96:14:45:FF:41:3D\x05c#=12\x04s#=1\x04ff=0\x04"
+            b"ci=2\x04sf=0\x0bsh=b0QZlg==",
+        },
+        {
+            "name": "HASS Bridge WLTF 4BE61F._hap._tcp.local.",
+            "address": address,
+            "port": 51837,
+            "text": b"\x13md=HASS Bridge"
+            b" WLTF\x06pv=1.0\x14id=E0:E7:98:4B:E6:1F\x04c#=2\x04s#=1\x04ff=0\x04"
+            b"ci=2\x04sf=0\x0bsh=ahAISA==",
+        },
+        {
+            "name": "FrontdoorCamera 8941D1._hap._tcp.local.",
+            "address": address,
+            "port": 54898,
+            "text": b"\x12md=FrontdoorCamera\x06pv=1.0\x14id=9F:B7:DC:89:41:D1\x04c#=2\x04"
+            b"s#=1\x04ff=0\x04ci=2\x04sf=0\x0bsh=0+MXmA==",
+        },
+        {
+            "name": "HASS Bridge W9DN 5B5CC5._hap._tcp.local.",
+            "address": address,
+            "port": 51836,
+            "text": b"\x13md=HASS Bridge"
+            b" W9DN\x06pv=1.0\x14id=11:8E:DB:5B:5C:C5\x05c#=12\x04s#=1\x04ff=0\x04"
+            b"ci=2\x04sf=0\x0bsh=6fLM5A==",
+        },
+        {
+            "name": "HASS Bridge Y9OO EFF0A7._hap._tcp.local.",
+            "address": address,
+            "port": 51838,
+            "text": b"\x13md=HASS Bridge"
+            b" Y9OO\x06pv=1.0\x14id=D3:FE:98:EF:F0:A7\x04c#=2\x04s#=1\x04ff=0\x04"
+            b"ci=2\x04sf=0\x0bsh=u3bdfw==",
+        },
+        {
+            "name": "Snooze Room TV 6B89B0._hap._tcp.local.",
+            "address": address,
+            "port": 51835,
+            "text": b"\x11md=Snooze Room"
+            b" TV\x06pv=1.0\x14id=5F:D5:70:6B:89:B0\x05c#=17\x04s#=1\x04ff=0\x05"
+            b"ci=31\x04sf=0\x0bsh=xNTqsg==",
+        },
+        {
+            "name": "AlexanderHomeAssistant 74651D._hap._tcp.local.",
+            "address": address,
+            "port": 54811,
+            "text": b"\x19md=AlexanderHomeAssistant\x06pv=1.0\x14id=59:8A:0B:74:65:1D\x05"
+            b"c#=14\x04s#=1\x04ff=0\x04ci=2\x04sf=0\x0bsh=ccZLPA==",
+        },
+        {
+            "name": "HASS Bridge OS95 39C053._hap._tcp.local.",
+            "address": address,
+            "port": 51831,
+            "text": b"\x13md=HASS Bridge"
+            b" OS95\x06pv=1.0\x14id=7E:8C:E6:39:C0:53\x05c#=12\x04s#=1\x04ff=0\x04ci=2"
+            b"\x04sf=0\x0bsh=Xfe5LQ==",
+        },
+    ]
+
+    out.add_answer_at_time(
+        DNSText(
+            "HASS Bridge W9DN 5B5CC5._hap._tcp.local.",
+            const._TYPE_TXT,
+            const._CLASS_IN | const._CLASS_UNIQUE,
+            const._DNS_OTHER_TTL,
+            b'\x13md=HASS Bridge W9DN\x06pv=1.0\x14id=11:8E:DB:5B:5C:C5\x05c#=12\x04s#=1'
+            b'\x04ff=0\x04ci=2\x04sf=0\x0bsh=6fLM5A==',
+        ),
+        0,
+    )
+
+    for record in additionals:
+        out.add_additional_answer(
+            DNSService(
+                record["name"],  # type: ignore
+                const._TYPE_SRV,
+                const._CLASS_IN | const._CLASS_UNIQUE,
+                const._DNS_HOST_TTL,
+                0,
+                0,
+                record["port"],  # type: ignore
+                record["name"],  # type: ignore
+            )
+        )
+        out.add_additional_answer(
+            DNSText(
+                record["name"],  # type: ignore
+                const._TYPE_TXT,
+                const._CLASS_IN | const._CLASS_UNIQUE,
+                const._DNS_OTHER_TTL,
+                record["text"],  # type: ignore
+            )
+        )
+        out.add_additional_answer(
+            DNSAddress(
+                record["name"],  # type: ignore
+                const._TYPE_A,
+                const._CLASS_IN | const._CLASS_UNIQUE,
+                const._DNS_HOST_TTL,
+                record["address"],  # type: ignore
+            )
+        )
+
+    return out
+
+
+out = generate_packets()
+
+
+def make_outgoing_message() -> None:
+    out.state = State.init
+    out.finished = False
+    out.packets()
+
+
+count = 100000
+time = timeit.Timer(make_outgoing_message).timeit(count)
+print(f"Construction {count} outgoing messages took {time} seconds")
diff --git a/src/zeroconf/_protocol/outgoing.pxd b/src/zeroconf/_protocol/outgoing.pxd
@@ -0,0 +1,68 @@
+
+import cython
+
+from .incoming cimport DNSIncoming
+
+
+cdef cython.uint _CLASS_UNIQUE
+cdef cython.uint _DNS_PACKET_HEADER_LEN
+cdef cython.uint _FLAGS_QR_MASK
+cdef cython.uint _FLAGS_QR_QUERY
+cdef cython.uint _FLAGS_QR_RESPONSE
+cdef cython.uint _FLAGS_TC
+cdef cython.uint _MAX_MSG_ABSOLUTE
+cdef cython.uint _MAX_MSG_TYPICAL
+
+
+cdef class DNSOutgoing:
+
+    cdef public unsigned int flags
+    cdef public object finished
+    cdef public object id
+    cdef public bint multicast
+    cdef public cython.list packets_data
+    cdef public object names
+    cdef public cython.list data
+    cdef public unsigned int size
+    cdef public object allow_long
+    cdef public object state
+    cdef public cython.list questions
+    cdef public cython.list answers
+    cdef public cython.list authorities
+    cdef public cython.list additionals
+
+    cdef _reset_for_next_packet(self)
+
+    cdef _write_byte(self, object value)
+
+    cdef _insert_short_at_start(self, object value)
+
+    cdef _replace_short(self, object index, object value)
+
+    cdef _write_int(self, object value)
+
+    cdef _write_question(self, object question)
+
+    cdef _write_record_class(self, object record)
+
+    cdef _check_data_limit_or_rollback(self, object start_data_length, object start_size)
+
+    cdef _write_questions_from_offset(self, object questions_offset)
+
+    cdef _write_answers_from_offset(self, object answer_offset)
+
+    cdef _write_records_from_offset(self, object records, object offset)
+
+    cdef _has_more_to_add(self, object questions_offset, object answer_offset, object authority_offset, object additional_offset)
+
+    @cython.locals(
+        questions_offset=cython.uint,
+        answer_offset=cython.uint,
+        authority_offset=cython.uint,
+        additional_offset=cython.uint,
+        questions_written=cython.uint,
+        answers_written=cython.uint,
+        authorities_written=cython.uint,
+        additionals_written=cython.uint,
+    )
+    cdef _packets(self)
diff --git a/src/zeroconf/_protocol/outgoing.py b/src/zeroconf/_protocol/outgoing.py
@@ -21,6 +21,7 @@
 """
 
 import enum
+import logging
 from typing import Dict, List, Optional, Sequence, Tuple, Union
 
 from .._cache import DNSCache
@@ -40,6 +41,11 @@
 from .incoming import DNSIncoming
 
 
+class State(enum.Enum):
+    init = 0
+    finished = 1
+
+
 class DNSOutgoing:
 
     """Object representation of an outgoing packet"""
@@ -74,7 +80,7 @@ def __init__(self, flags: int, multicast: bool = True, id_: int = 0) -> None:
         self.size: int = _DNS_PACKET_HEADER_LEN
         self.allow_long: bool = True
 
-        self.state = self.State.init
+        self.state = State.init
 
         self.questions: List[DNSQuestion] = []
         self.answers: List[Tuple[DNSRecord, float]] = []
@@ -107,10 +113,6 @@ def __repr__(self) -> str:
             ]
         )
 
-    class State(enum.Enum):
-        init = 0
-        finished = 1
-
     def add_question(self, record: DNSQuestion) -> None:
         """Adds a question"""
         self.questions.append(record)
@@ -373,8 +375,10 @@ def packets(self) -> List[bytes]:
         will be written out to a single oversized packet no more than
         _MAX_MSG_ABSOLUTE in length (and hence will be subject to IP
         fragmentation potentially)."""
+        return self._packets()
 
-        if self.state == self.State.finished:
+    def _packets(self) -> List[bytes]:
+        if self.state == State.finished:
             return self.packets_data
 
         questions_offset = 0
@@ -383,25 +387,27 @@ def packets(self) -> List[bytes]:
         additional_offset = 0
         # we have to at least write out the question
         first_time = True
+        debug_enable = log.isEnabledFor(logging.DEBUG)
 
         while first_time or self._has_more_to_add(
             questions_offset, answer_offset, authority_offset, additional_offset
         ):
             first_time = False
-            log.debug(
-                "offsets = questions=%d, answers=%d, authorities=%d, additionals=%d",
-                questions_offset,
-                answer_offset,
-                authority_offset,
-                additional_offset,
-            )
-            log.debug(
-                "lengths = questions=%d, answers=%d, authorities=%d, additionals=%d",
-                len(self.questions),
-                len(self.answers),
-                len(self.authorities),
-                len(self.additionals),
-            )
+            if debug_enable:
+                log.debug(
+                    "offsets = questions=%d, answers=%d, authorities=%d, additionals=%d",
+                    questions_offset,
+                    answer_offset,
+                    authority_offset,
+                    additional_offset,
+                )
+                log.debug(
+                    "lengths = questions=%d, answers=%d, authorities=%d, additionals=%d",
+                    len(self.questions),
+                    len(self.answers),
+                    len(self.authorities),
+                    len(self.additionals),
+                )
 
             questions_written = self._write_questions_from_offset(questions_offset)
             answers_written = self._write_answers_from_offset(answer_offset)
@@ -417,13 +423,14 @@ def packets(self) -> List[bytes]:
             answer_offset += answers_written
             authority_offset += authorities_written
             additional_offset += additionals_written
-            log.debug(
-                "now offsets = questions=%d, answers=%d, authorities=%d, additionals=%d",
-                questions_offset,
-                answer_offset,
-                authority_offset,
-                additional_offset,
-            )
+            if debug_enable:
+                log.debug(
+                    "now offsets = questions=%d, answers=%d, authorities=%d, additionals=%d",
+                    questions_offset,
+                    answer_offset,
+                    authority_offset,
+                    additional_offset,
+                )
 
             if self.is_query() and self._has_more_to_add(
                 questions_offset, answer_offset, authority_offset, additional_offset
@@ -447,5 +454,5 @@ def packets(self) -> List[bytes]:
             ) > 0:
                 log.warning("packets() made no progress adding records; returning")
                 break
-        self.state = self.State.finished
+        self.state = State.finished
         return self.packets_data
diff --git a/tests/test_core.py b/tests/test_core.py
@@ -714,12 +714,17 @@ def test_guard_against_oversized_packets():
             0,
         )
 
-    # We are patching to generate an oversized packet
-    with patch.object(outgoing, "_MAX_MSG_ABSOLUTE", 100000), patch.object(
-        outgoing, "_MAX_MSG_TYPICAL", 100000
-    ):
-        over_sized_packet = generated.packets()[0]
-        assert len(over_sized_packet) > const._MAX_MSG_ABSOLUTE
+    try:
+        # We are patching to generate an oversized packet
+        with patch.object(outgoing, "_MAX_MSG_ABSOLUTE", 100000), patch.object(
+            outgoing, "_MAX_MSG_TYPICAL", 100000
+        ):
+            over_sized_packet = generated.packets()[0]
+            assert len(over_sized_packet) > const._MAX_MSG_ABSOLUTE
+    except AttributeError:
+        # cannot patch with cython
+        zc.close()
+        return
 
     generated = r.DNSOutgoing(const._FLAGS_QR_RESPONSE)
     okpacket_record = r.DNSText(