Small speed up to the plaintext frame helper (#625)

This commit is contained in:
J. Nick Koston 2023-11-09 10:20:42 -06:00 committed by GitHub
parent 622d6060ee
commit 94157a8ac1
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 15 additions and 13 deletions

View File

@ -142,8 +142,7 @@ class APINoiseFrameHelper(APIFrameHelper):
self._add_to_buffer(data) self._add_to_buffer(data)
while self._buffer: while self._buffer:
self._pos = 0 self._pos = 0
header = self._read_exactly(3) if (header := self._read_exactly(3)) is None:
if header is None:
return return
preamble = header[0] preamble = header[0]
msg_size_high = header[1] msg_size_high = header[1]

View File

@ -78,15 +78,16 @@ class APIPlaintextFrameHelper(APIFrameHelper):
f"{self._log_name}: Error while writing data: {err}" f"{self._log_name}: Error while writing data: {err}"
) from err ) from err
def data_received(self, data: bytes) -> None: # pylint: disable=too-many-branches def data_received( # pylint: disable=too-many-branches,too-many-return-statements
self, data: bytes
) -> None:
self._add_to_buffer(data) self._add_to_buffer(data)
while self._buffer: while self._buffer:
# Read preamble, which should always 0x00 # Read preamble, which should always 0x00
# Also try to get the length and msg type # Also try to get the length and msg type
# to avoid multiple calls to _read_exactly # to avoid multiple calls to _read_exactly
self._pos = 0 self._pos = 0
init_bytes = self._read_exactly(3) if (init_bytes := self._read_exactly(3)) is None:
if init_bytes is None:
return return
msg_type_int: int | None = None msg_type_int: int | None = None
length_int = 0 length_int = 0
@ -125,22 +126,25 @@ class APIPlaintextFrameHelper(APIFrameHelper):
length = init_bytes[1:3] length = init_bytes[1:3]
# If the message is long, we need to read the rest of the length # If the message is long, we need to read the rest of the length
while length[-1] & 0x80 == 0x80: while length[-1] & 0x80 == 0x80:
add_length = self._read_exactly(1) if (add_length := self._read_exactly(1)) is None:
if add_length is None:
return return
length += add_length length += add_length
length_int = bytes_to_varuint(length) or 0 length_int = bytes_to_varuint(length) or 0
# Since the length is longer than 1 byte we do not have the # Since the length is longer than 1 byte we do not have the
# message type yet. # message type yet.
msg_type = b"" if (msg_type_byte := self._read_exactly(1)) is None:
return
msg_type = msg_type_byte
if msg_type[-1] & 0x80 != 0x80:
# Message type is only 1 byte
msg_type_int = msg_type[0]
# If the we do not have the message type yet because the message # If the we do not have the message type yet because the message
# length was so long it did not fit into the first byte we need # length was so long it did not fit into the first byte we need
# to read the (rest) of the message type # to read the (rest) of the message type
if msg_type_int is None: if msg_type_int is None:
while not msg_type or msg_type[-1] & 0x80 == 0x80: while msg_type[-1] & 0x80 == 0x80:
add_msg_type = self._read_exactly(1) if (add_msg_type := self._read_exactly(1)) is None:
if add_msg_type is None:
return return
msg_type += add_msg_type msg_type += add_msg_type
msg_type_int = bytes_to_varuint(msg_type) msg_type_int = bytes_to_varuint(msg_type)
@ -151,13 +155,12 @@ class APIPlaintextFrameHelper(APIFrameHelper):
if length_int == 0: if length_int == 0:
packet_data = b"" packet_data = b""
else: else:
maybe_packet_data = self._read_exactly(length_int)
# The packet data is not yet available, wait for more data # The packet data is not yet available, wait for more data
# to arrive before continuing, since callback_packet has not # to arrive before continuing, since callback_packet has not
# been called yet the buffer will not be cleared and the next # been called yet the buffer will not be cleared and the next
# call to data_received will continue processing the packet # call to data_received will continue processing the packet
# at the start of the frame. # at the start of the frame.
if maybe_packet_data is None: if (maybe_packet_data := self._read_exactly(length_int)) is None:
return return
packet_data = maybe_packet_data packet_data = maybe_packet_data