From 3752b8280e2a6837dcb226faecf458ae2282a387 Mon Sep 17 00:00:00 2001
From: Otto Winter <otto@otto-winter.com>
Date: Wed, 9 Feb 2022 16:29:50 +0100
Subject: [PATCH] Fix no timeout for handshake (#176)

---
 aioesphomeapi/_frame_helper.py |  9 ++++++++-
 aioesphomeapi/connection.py    | 10 ++++++++--
 2 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/aioesphomeapi/_frame_helper.py b/aioesphomeapi/_frame_helper.py
index f142c62..19a2049 100644
--- a/aioesphomeapi/_frame_helper.py
+++ b/aioesphomeapi/_frame_helper.py
@@ -179,7 +179,7 @@ class APINoiseFrameHelper(APIFrameHelper):
         _LOGGER.debug("Received frame %s", frame.hex())
         return frame
 
-    async def perform_handshake(self, expected_name: Optional[str]) -> None:
+    async def _perform_handshake(self, expected_name: Optional[str]) -> None:
         await self._write_frame(b"")  # ClientHello
         prologue = b"NoiseAPIInit" + b"\x00\x00"
 
@@ -236,6 +236,13 @@ class APINoiseFrameHelper(APIFrameHelper):
         _LOGGER.debug("Handshake complete!")
         self._ready_event.set()
 
+    async def perform_handshake(self, expected_name: Optional[str]) -> None:
+        # Allow up to 60 seconds for handhsake
+        try:
+            await asyncio.wait_for(self._perform_handshake(expected_name), timeout=60.0)
+        except asyncio.TimeoutError as err:
+            raise HandshakeAPIError("Timeout during handshake") from err
+
     async def write_packet(self, packet: Packet) -> None:
         # Wait for handshake to complete
         await self._ready_event.wait()
diff --git a/aioesphomeapi/connection.py b/aioesphomeapi/connection.py
index d8b0761..3a083e2 100644
--- a/aioesphomeapi/connection.py
+++ b/aioesphomeapi/connection.py
@@ -164,7 +164,7 @@ class APIConnection:
         except asyncio.TimeoutError as err:
             raise SocketAPIError(f"Timeout while connecting to {sockaddr}") from err
 
-        _LOGGER.debug("%s: Opened socket for", self._params.address)
+        _LOGGER.debug("%s: Opened socket", self._params.address)
 
     async def _connect_init_frame_helper(self) -> None:
         """Step 3 in connect process: initialize the frame helper and init read loop."""
@@ -266,7 +266,7 @@ class APIConnection:
                 "Connection can only be used once, connection is not in init state"
             )
 
-        try:
+        async def _do_connect() -> None:
             addr = await self._connect_resolve_host()
             await self._connect_socket_connect(addr)
             await self._connect_init_frame_helper()
@@ -274,6 +274,12 @@ class APIConnection:
             await self._connect_start_ping()
             if login:
                 await self.login()
+
+        try:
+            # Allow 2 minutes for connect; this is only as a last measure
+            # to protect from issues if some part of the connect process mistakenly
+            # does not have a timeout
+            await asyncio.wait_for(_do_connect(), timeout=120.0)
         except Exception:  # pylint: disable=broad-except
             # Always clean up the connection if an error occured during connect
             self._connection_state = ConnectionState.CLOSED