feat: speed up reading incoming packets

bdraco · bdraco · commit 6280ad4cfa1d · 2023-11-13T15:23:15.000-06:00
diff --git a/src/zeroconf/_protocol/incoming.pxd b/src/zeroconf/_protocol/incoming.pxd
@@ -21,11 +21,6 @@ cdef cython.uint _FLAGS_TC
 cdef cython.uint _FLAGS_QR_QUERY
 cdef cython.uint _FLAGS_QR_RESPONSE
 
-cdef object UNPACK_3H
-cdef object UNPACK_6H
-cdef object UNPACK_HH
-cdef object UNPACK_HHiH
-
 cdef object DECODE_EXCEPTIONS
 
 cdef object IncomingDecodeError
@@ -79,6 +74,12 @@ cdef class DNSIncoming:
 
     cpdef is_response(self)
 
+    @cython.locals(offset="unsigned char")
+    cdef unsigned int _read_short(self)
+
+    @cython.locals(offset="unsigned char")
+    cdef unsigned int _read_long(self)
+
     @cython.locals(
         off=cython.uint,
         label_idx=cython.uint,
diff --git a/src/zeroconf/_protocol/incoming.py b/src/zeroconf/_protocol/incoming.py
@@ -60,10 +60,6 @@
 
 DECODE_EXCEPTIONS = (IndexError, struct.error, IncomingDecodeError)
 
-UNPACK_3H = struct.Struct(b'!3H').unpack_from
-UNPACK_6H = struct.Struct(b'!6H').unpack_from
-UNPACK_HH = struct.Struct(b'!HH').unpack_from
-UNPACK_HHiH = struct.Struct(b'!HHiH').unpack_from
 
 _seen_logs: Dict[str, Union[int, tuple]] = {}
 _str = str
@@ -207,24 +203,35 @@ def __repr__(self) -> str:
             ]
         )
 
+    def _read_short(self) -> _int:
+        """Reads an unsigned short in network order from the packet."""
+        offset = self.offset
+        view = self.view
+        self.offset += 2
+        return view[offset] << 8 | view[offset + 1]
+
+    def _read_long(self) -> _int:
+        """Reads an unsigned long in network order from the packet."""
+        offset = self.offset
+        view = self.view
+        self.offset += 4
+        return view[offset] << 24 | view[offset + 1] << 16 | view[offset + 2] << 8 | view[offset + 3]
+
     def _read_header(self) -> None:
         """Reads header portion of packet"""
-        (
-            self.id,
-            self.flags,
-            self.num_questions,
-            self.num_answers,
-            self.num_authorities,
-            self.num_additionals,
-        ) = UNPACK_6H(self.data)
-        self.offset += 12
+        self.id = self._read_short()
+        self.flags = self._read_short()
+        self.num_questions = self._read_short()
+        self.num_answers = self._read_short()
+        self.num_authorities = self._read_short()
+        self.num_additionals = self._read_short()
 
     def _read_questions(self) -> None:
         """Reads questions section of packet"""
         for _ in range(self.num_questions):
             name = self._read_name()
-            type_, class_ = UNPACK_HH(self.data, self.offset)
-            self.offset += 4
+            type_ = self._read_short()
+            class_ = self._read_short()
             question = DNSQuestion(name, type_, class_)
             self.questions.append(question)
 
@@ -249,8 +256,10 @@ def _read_others(self) -> None:
         n = self.num_answers + self.num_authorities + self.num_additionals
         for _ in range(n):
             domain = self._read_name()
-            type_, class_, ttl, length = UNPACK_HHiH(self.data, self.offset)
-            self.offset += 10
+            type_ = self._read_short()
+            class_ = self._read_short()
+            ttl = self._read_long()
+            length = self._read_short()
             end = self.offset + length
             rec = None
             try:
@@ -284,7 +293,9 @@ def _read_record(
         if type_ == _TYPE_TXT:
             return DNSText(domain, type_, class_, ttl, self._read_string(length), self.now)
         if type_ == _TYPE_SRV:
-            priority, weight, port = UNPACK_3H(self.data, self.offset)
+            priority = self._read_short()
+            weight = self._read_short()
+            port = self._read_short()
             self.offset += 6
             return DNSService(
                 domain,
