Handle prepare timeout in websocket API (#55989)

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()
 

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