Skip to content

Commit

Permalink
feat: speed up reading incoming packets
Browse files Browse the repository at this point in the history
  • Loading branch information
bdraco committed Nov 13, 2023
1 parent bf2cfde commit 1c7a4d3
Show file tree
Hide file tree
Showing 2 changed files with 35 additions and 24 deletions.
11 changes: 6 additions & 5 deletions src/zeroconf/_protocol/incoming.pxd
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down
48 changes: 29 additions & 19 deletions src/zeroconf/_protocol/incoming.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)

Expand All @@ -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:
Expand Down Expand Up @@ -284,8 +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)
self.offset += 6
priority = self._read_short()
weight = self._read_short()
port = self._read_short()
return DNSService(
domain,
type_,
Expand Down

0 comments on commit 1c7a4d3

Please sign in to comment.