From 108a54a4a873e749bbb9e45f7c19c8097a79a7c9 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Mon, 19 Aug 2024 15:28:05 -0500
Subject: [PATCH] Handle WebSocket client disconnect during prepare (#124173)

---
 .../components/websocket_api/http.py          |  7 ++++++
 tests/components/websocket_api/test_http.py   | 22 +++++++++++++++++--
 2 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/homeassistant/components/websocket_api/http.py b/homeassistant/components/websocket_api/http.py
index e33da9a8b4a..1ad8d909ce8 100644
--- a/homeassistant/components/websocket_api/http.py
+++ b/homeassistant/components/websocket_api/http.py
@@ -306,6 +306,13 @@ class WebSocketHandler:
         try:
             async with asyncio.timeout(10):
                 await wsock.prepare(request)
+        except ConnectionResetError:
+            # Likely the client disconnected before we prepared the websocket
+            logger.debug(
+                "%s: Connection reset by peer while preparing WebSocket",
+                self.description,
+            )
+            return wsock
         except TimeoutError:
             logger.warning("Timeout preparing request from %s", request.remote)
             return wsock
diff --git a/tests/components/websocket_api/test_http.py b/tests/components/websocket_api/test_http.py
index 11665da11b4..2530d885942 100644
--- a/tests/components/websocket_api/test_http.py
+++ b/tests/components/websocket_api/test_http.py
@@ -363,12 +363,12 @@ async def test_non_json_message(
     assert "bad=<object" in caplog.text
 
 
-async def test_prepare_fail(
+async def test_prepare_fail_timeout(
     hass: HomeAssistant,
     hass_ws_client: WebSocketGenerator,
     caplog: pytest.LogCaptureFixture,
 ) -> None:
-    """Test failing to prepare."""
+    """Test failing to prepare due to timeout."""
     with (
         patch(
             "homeassistant.components.websocket_api.http.web.WebSocketResponse.prepare",
@@ -381,6 +381,24 @@ async def test_prepare_fail(
     assert "Timeout preparing request" in caplog.text
 
 
+async def test_prepare_fail_connection_reset(
+    hass: HomeAssistant,
+    hass_ws_client: WebSocketGenerator,
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """Test failing to prepare due to connection reset."""
+    with (
+        patch(
+            "homeassistant.components.websocket_api.http.web.WebSocketResponse.prepare",
+            side_effect=(ConnectionResetError, web.WebSocketResponse.prepare),
+        ),
+        pytest.raises(WSServerHandshakeError),
+    ):
+        await hass_ws_client(hass)
+
+    assert "Connection reset by peer while preparing WebSocket" in caplog.text
+
+
 async def test_enable_coalesce(
     hass: HomeAssistant,
     hass_ws_client: WebSocketGenerator,