From 6d0da631bf0052bc10d55c60ea22f2862c46339c Mon Sep 17 00:00:00 2001
From: Paulus Schoutsen <balloob@gmail.com>
Date: Fri, 8 Oct 2021 22:12:06 -0700
Subject: [PATCH] Handle prepare timeout in websocket API (#55989)

---
 homeassistant/components/websocket_api/http.py | 13 +++++++++----
 tests/components/websocket_api/test_http.py    | 14 +++++++++++++-
 2 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/homeassistant/components/websocket_api/http.py b/homeassistant/components/websocket_api/http.py
index bce6713403a..8d75d50e59a 100644
--- a/homeassistant/components/websocket_api/http.py
+++ b/homeassistant/components/websocket_api/http.py
@@ -61,7 +61,7 @@ class WebSocketHandler:
         """Initialize an active connection."""
         self.hass = hass
         self.request = request
-        self.wsock: web.WebSocketResponse | None = None
+        self.wsock = web.WebSocketResponse(heartbeat=55)
         self._to_write: asyncio.Queue = asyncio.Queue(maxsize=MAX_PENDING_MSG)
         self._handle_task: asyncio.Task | None = None
         self._writer_task: asyncio.Task | None = None
@@ -71,7 +71,6 @@ class WebSocketHandler:
     async def _writer(self) -> None:
         """Write outgoing messages."""
         # Exceptions if Socket disconnected or cancelled by connection handler
-        assert self.wsock is not None
         with suppress(RuntimeError, ConnectionResetError, *CANCELLATION_ERRORS):
             while not self.wsock.closed:
                 message = await self._to_write.get()
@@ -143,8 +142,14 @@ class WebSocketHandler:
     async def async_handle(self) -> web.WebSocketResponse:
         """Handle a websocket response."""
         request = self.request
-        wsock = self.wsock = web.WebSocketResponse(heartbeat=55)
-        await wsock.prepare(request)
+        wsock = self.wsock
+        try:
+            async with async_timeout.timeout(10):
+                await wsock.prepare(request)
+        except asyncio.TimeoutError:
+            self._logger.warning("Timeout preparing request from %s", request.remote)
+            return wsock
+
         self._logger.debug("Connected from %s", request.remote)
         self._handle_task = asyncio.current_task()
 
diff --git a/tests/components/websocket_api/test_http.py b/tests/components/websocket_api/test_http.py
index f3952f1dc4b..336c79d22b8 100644
--- a/tests/components/websocket_api/test_http.py
+++ b/tests/components/websocket_api/test_http.py
@@ -1,8 +1,9 @@
 """Test Websocket API http module."""
+import asyncio
 from datetime import timedelta
 from unittest.mock import patch
 
-from aiohttp import WSMsgType
+from aiohttp import ServerDisconnectedError, WSMsgType, web
 import pytest
 
 from homeassistant.components.websocket_api import const, http
@@ -80,3 +81,14 @@ async def test_non_json_message(hass, websocket_client, caplog):
         f"Unable to serialize to JSON. Bad data found at $.result[0](State: test_domain.entity).attributes.bad={bad_data}(<class 'object'>"
         in caplog.text
     )
+
+
+async def test_prepare_fail(hass, hass_ws_client, caplog):
+    """Test failing to prepare."""
+    with patch(
+        "homeassistant.components.websocket_api.http.web.WebSocketResponse.prepare",
+        side_effect=(asyncio.TimeoutError, web.WebSocketResponse.prepare),
+    ), pytest.raises(ServerDisconnectedError):
+        await hass_ws_client(hass)
+
+    assert "Timeout preparing request" in caplog.text