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