Upstream/aioesphomeapi
1
0
Fork 0
mirror of https://github.com/esphome/aioesphomeapi.git synced 2025-01-12 20:10:42 +01:00
Code Issues Projects Releases Wiki Activity

Reduce lookups for noise in the frame helper (#465)

Browse Source
This commit is contained in:
J. Nick Koston 2023-07-10 21:49:34 -10:00 committed by GitHub
parent 7a80e3529b
commit 5f82e90b70
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 11 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
aioesphomeapi/_frame_helper.py
View File

@ -263,6 +263,7 @@ class APINoiseFrameHelper(APIFrameHelper):
"_noise_psk", "_noise_psk",
"_expected_name", "_expected_name",
"_state", "_state",
"_dispatch",
"_server_name", "_server_name",
"_proto", "_proto",
"_decrypt", "_decrypt",
@ -282,7 +283,7 @@ class APINoiseFrameHelper(APIFrameHelper):
self._ready_future = asyncio.get_event_loop().create_future() self._ready_future = asyncio.get_event_loop().create_future()
self._noise_psk = noise_psk self._noise_psk = noise_psk
self._expected_name = expected_name self._expected_name = expected_name
self._state = NoiseConnectionState.HELLO self._set_state(NoiseConnectionState.HELLO)
self._server_name: Optional[str] = None self._server_name: Optional[str] = None
self._decrypt: Optional[Callable[[bytes], bytes]] = None self._decrypt: Optional[Callable[[bytes], bytes]] = None
self._encrypt: Optional[Callable[[bytes], bytes]] = None self._encrypt: Optional[Callable[[bytes], bytes]] = None
@ -292,13 +293,18 @@ class APINoiseFrameHelper(APIFrameHelper):
if not self._ready_future.done(): if not self._ready_future.done():
self._ready_future.set_exception(exc) self._ready_future.set_exception(exc)
def _set_state(self, state: NoiseConnectionState) -> None:
"""Set the current state."""
self._state = state
self._dispatch = self.STATE_TO_CALLABLE[state]
def close(self) -> None: def close(self) -> None:
"""Close the connection.""" """Close the connection."""
# Make sure we set the ready event if its not already set # Make sure we set the ready event if its not already set
# so that we don't block forever on the ready event if we # so that we don't block forever on the ready event if we
# are waiting for the handshake to complete. # are waiting for the handshake to complete.
self._set_ready_future_exception(APIConnectionError("Connection closed")) self._set_ready_future_exception(APIConnectionError("Connection closed"))
self._state = NoiseConnectionState.CLOSED self._set_state(NoiseConnectionState.CLOSED)
super().close() super().close()
def _handle_error_and_close(self, exc: Exception) -> None: def _handle_error_and_close(self, exc: Exception) -> None:
@ -375,7 +381,7 @@ class APINoiseFrameHelper(APIFrameHelper):
return return
try: try:
self.STATE_TO_CALLABLE[self._state](self, frame) self._dispatch(self, frame)
except Exception as err: # pylint: disable=broad-except except Exception as err: # pylint: disable=broad-except
self._handle_error_and_close(err) self._handle_error_and_close(err)
finally: finally:
@ -420,7 +426,7 @@ class APINoiseFrameHelper(APIFrameHelper):
) )
return return
self._state = NoiseConnectionState.HANDSHAKE self._set_state(NoiseConnectionState.HANDSHAKE)
self._send_handshake() self._send_handshake()
def _setup_proto(self) -> None: def _setup_proto(self) -> None:
@ -462,7 +468,7 @@ class APINoiseFrameHelper(APIFrameHelper):
self._handle_error_and_close(ex) self._handle_error_and_close(ex)
return return
_LOGGER.debug("Handshake complete") _LOGGER.debug("Handshake complete")
self._state = NoiseConnectionState.READY self._set_state(NoiseConnectionState.READY)
noise_protocol = self._proto.noise_protocol noise_protocol = self._proto.noise_protocol
self._decrypt = partial( self._decrypt = partial(
noise_protocol.cipher_state_decrypt.decrypt_with_ad, # pylint: disable=no-member noise_protocol.cipher_state_decrypt.decrypt_with_ad, # pylint: disable=no-member