From 81fc15d6ac866605e8253cb08789e2b2752e2ecf Mon Sep 17 00:00:00 2001
From: Stefan Agner <stefan@agner.ch>
Date: Wed, 26 Mar 2025 22:13:59 +0100
Subject: [PATCH] Handle unexpected WebSocket messages during auth (#5788)

* Handle unexpected WebSocket messages during auth

When an add-on does not respond or closes the WebSocket connection
during the authentication phase Supervisor does not handle errors
gracefully. Simply log such unexpected authentication to avoid
unnecessary stack traces in the log and make such cases no longer
appear on Sentry.

* Add pytest

* Introduce a timeout of 10s
---
 supervisor/api/proxy.py | 14 +++++++++++---
 tests/api/test_proxy.py | 19 +++++++++++++++++++
 2 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/supervisor/api/proxy.py b/supervisor/api/proxy.py
index 000df3c53..d47efa8eb 100644
--- a/supervisor/api/proxy.py
+++ b/supervisor/api/proxy.py
@@ -6,7 +6,7 @@ from contextlib import asynccontextmanager
 import logging
 
 import aiohttp
-from aiohttp import web
+from aiohttp import WSMessageTypeError, web
 from aiohttp.client_exceptions import ClientConnectorError
 from aiohttp.client_ws import ClientWebSocketResponse
 from aiohttp.hdrs import AUTHORIZATION, CONTENT_TYPE
@@ -215,8 +215,8 @@ class APIProxy(CoreSysAttributes):
                 dumps=json_dumps,
             )
 
-            # Check API access
-            response = await server.receive_json()
+            # Check API access, wait up to 10s just like _async_handle_auth_phase in Core
+            response = await server.receive_json(timeout=10)
             supervisor_token = response.get("api_password") or response.get(
                 "access_token"
             )
@@ -237,6 +237,14 @@ class APIProxy(CoreSysAttributes):
                 {"type": "auth_ok", "ha_version": self.sys_homeassistant.version},
                 dumps=json_dumps,
             )
+        except TimeoutError:
+            _LOGGER.error("Timeout during authentication for WebSocket API")
+            return server
+        except WSMessageTypeError as err:
+            _LOGGER.error(
+                "Unexpected message during authentication for WebSocket API: %s", err
+            )
+            return server
         except (RuntimeError, ValueError) as err:
             _LOGGER.error("Can't initialize handshake: %s", err)
             return server
diff --git a/tests/api/test_proxy.py b/tests/api/test_proxy.py
index 033585a2a..fb9c0b53b 100644
--- a/tests/api/test_proxy.py
+++ b/tests/api/test_proxy.py
@@ -5,6 +5,7 @@ from __future__ import annotations
 import asyncio
 from collections.abc import Awaitable, Callable, Coroutine, Generator
 from json import dumps
+import logging
 from typing import Any, cast
 from unittest.mock import patch
 
@@ -175,3 +176,21 @@ async def test_proxy_invalid_auth(
     auth_not_ok = await websocket.receive_json()
     assert auth_not_ok["type"] == "auth_invalid"
     assert auth_not_ok["message"] == "Invalid access"
+
+
+async def test_proxy_auth_abort_log(
+    api_client: TestClient,
+    install_addon_example: Addon,
+    caplog: pytest.LogCaptureFixture,
+):
+    """Test WebSocket closed during authentication gets logged."""
+    install_addon_example.persist[ATTR_ACCESS_TOKEN] = "abc123"
+    websocket = await api_client.ws_connect("/core/websocket")
+    auth_resp = await websocket.receive_json()
+    assert auth_resp["type"] == "auth_required"
+    caplog.clear()
+    with caplog.at_level(logging.ERROR):
+        await websocket.close()
+        assert (
+            "Unexpected message during authentication for WebSocket API" in caplog.text
+        )