503 lines
19 KiB
Python
503 lines
19 KiB
Python
"""View to accept incoming websocket connection."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from collections import deque
|
|
from collections.abc import Callable, Coroutine
|
|
import datetime as dt
|
|
from functools import partial
|
|
import logging
|
|
from typing import TYPE_CHECKING, Any, Final
|
|
|
|
from aiohttp import WSMsgType, web
|
|
|
|
from homeassistant.components.http import KEY_HASS, HomeAssistantView
|
|
from homeassistant.const import EVENT_HOMEASSISTANT_STOP
|
|
from homeassistant.core import Event, HomeAssistant, callback
|
|
from homeassistant.helpers.dispatcher import async_dispatcher_send
|
|
from homeassistant.helpers.event import async_call_later
|
|
from homeassistant.util.async_ import create_eager_task
|
|
from homeassistant.util.json import json_loads
|
|
|
|
from .auth import AUTH_REQUIRED_MESSAGE, AuthPhase
|
|
from .const import (
|
|
DATA_CONNECTIONS,
|
|
MAX_PENDING_MSG,
|
|
PENDING_MSG_MAX_FORCE_READY,
|
|
PENDING_MSG_PEAK,
|
|
PENDING_MSG_PEAK_TIME,
|
|
SIGNAL_WEBSOCKET_CONNECTED,
|
|
SIGNAL_WEBSOCKET_DISCONNECTED,
|
|
URL,
|
|
)
|
|
from .error import Disconnect
|
|
from .messages import message_to_json_bytes
|
|
from .util import describe_request
|
|
|
|
if TYPE_CHECKING:
|
|
from .connection import ActiveConnection
|
|
|
|
|
|
_WS_LOGGER: Final = logging.getLogger(f"{__name__}.connection")
|
|
|
|
|
|
class WebsocketAPIView(HomeAssistantView):
|
|
"""View to serve a websockets endpoint."""
|
|
|
|
name: str = "websocketapi"
|
|
url: str = URL
|
|
requires_auth: bool = False
|
|
|
|
async def get(self, request: web.Request) -> web.WebSocketResponse:
|
|
"""Handle an incoming websocket connection."""
|
|
return await WebSocketHandler(request.app[KEY_HASS], request).async_handle()
|
|
|
|
|
|
class WebSocketAdapter(logging.LoggerAdapter):
|
|
"""Add connection id to websocket messages."""
|
|
|
|
def process(self, msg: str, kwargs: Any) -> tuple[str, Any]:
|
|
"""Add connid to websocket log messages."""
|
|
assert self.extra is not None
|
|
return f'[{self.extra["connid"]}] {msg}', kwargs
|
|
|
|
|
|
class WebSocketHandler:
|
|
"""Handle an active websocket client connection."""
|
|
|
|
__slots__ = (
|
|
"_hass",
|
|
"_loop",
|
|
"_request",
|
|
"_wsock",
|
|
"_handle_task",
|
|
"_writer_task",
|
|
"_closing",
|
|
"_authenticated",
|
|
"_logger",
|
|
"_peak_checker_unsub",
|
|
"_connection",
|
|
"_message_queue",
|
|
"_ready_future",
|
|
"_release_ready_queue_size",
|
|
)
|
|
|
|
def __init__(self, hass: HomeAssistant, request: web.Request) -> None:
|
|
"""Initialize an active connection."""
|
|
self._hass = hass
|
|
self._loop = hass.loop
|
|
self._request: web.Request = request
|
|
self._wsock = web.WebSocketResponse(heartbeat=55)
|
|
self._handle_task: asyncio.Task | None = None
|
|
self._writer_task: asyncio.Task | None = None
|
|
self._closing: bool = False
|
|
self._authenticated: bool = False
|
|
self._logger = WebSocketAdapter(_WS_LOGGER, {"connid": id(self)})
|
|
self._peak_checker_unsub: Callable[[], None] | None = None
|
|
self._connection: ActiveConnection | None = None
|
|
|
|
# The WebSocketHandler has a single consumer and path
|
|
# to where messages are queued. This allows the implementation
|
|
# to use a deque and an asyncio.Future to avoid the overhead of
|
|
# an asyncio.Queue.
|
|
self._message_queue: deque[bytes] = deque()
|
|
self._ready_future: asyncio.Future[int] | None = None
|
|
self._release_ready_queue_size: int = 0
|
|
|
|
def __repr__(self) -> str:
|
|
"""Return the representation."""
|
|
return (
|
|
"<WebSocketHandler "
|
|
f"closing={self._closing} "
|
|
f"authenticated={self._authenticated} "
|
|
f"description={self.description}>"
|
|
)
|
|
|
|
@property
|
|
def description(self) -> str:
|
|
"""Return a description of the connection."""
|
|
if connection := self._connection:
|
|
return connection.get_description(self._request)
|
|
if request := self._request:
|
|
return describe_request(request)
|
|
return "finished connection"
|
|
|
|
async def _writer(
|
|
self, send_bytes_text: Callable[[bytes], Coroutine[Any, Any, None]]
|
|
) -> None:
|
|
"""Write outgoing messages."""
|
|
# Variables are set locally to avoid lookups in the loop
|
|
message_queue = self._message_queue
|
|
logger = self._logger
|
|
wsock = self._wsock
|
|
loop = self._loop
|
|
is_debug_log_enabled = partial(logger.isEnabledFor, logging.DEBUG)
|
|
debug = logger.debug
|
|
can_coalesce = self._connection and self._connection.can_coalesce
|
|
ready_message_count = len(message_queue)
|
|
# Exceptions if Socket disconnected or cancelled by connection handler
|
|
try:
|
|
while not wsock.closed:
|
|
if not message_queue:
|
|
self._ready_future = loop.create_future()
|
|
ready_message_count = await self._ready_future
|
|
|
|
if self._closing:
|
|
return
|
|
|
|
if not can_coalesce:
|
|
# coalesce may be enabled later in the connection
|
|
can_coalesce = self._connection and self._connection.can_coalesce
|
|
|
|
if not can_coalesce or ready_message_count == 1:
|
|
message = message_queue.popleft()
|
|
if is_debug_log_enabled():
|
|
debug("%s: Sending %s", self.description, message)
|
|
await send_bytes_text(message)
|
|
continue
|
|
|
|
coalesced_messages = b"".join((b"[", b",".join(message_queue), b"]"))
|
|
message_queue.clear()
|
|
if is_debug_log_enabled():
|
|
debug("%s: Sending %s", self.description, coalesced_messages)
|
|
await send_bytes_text(coalesced_messages)
|
|
except asyncio.CancelledError:
|
|
debug("%s: Writer cancelled", self.description)
|
|
raise
|
|
except (RuntimeError, ConnectionResetError) as ex:
|
|
debug("%s: Unexpected error in writer: %s", self.description, ex)
|
|
finally:
|
|
debug("%s: Writer done", self.description)
|
|
# Clean up the peak checker when we shut down the writer
|
|
self._cancel_peak_checker()
|
|
|
|
@callback
|
|
def _cancel_peak_checker(self) -> None:
|
|
"""Cancel the peak checker."""
|
|
if self._peak_checker_unsub is not None:
|
|
self._peak_checker_unsub()
|
|
self._peak_checker_unsub = None
|
|
|
|
@callback
|
|
def _send_message(self, message: str | bytes | dict[str, Any]) -> None:
|
|
"""Queue sending a message to the client.
|
|
|
|
Closes connection if the client is not reading the messages.
|
|
|
|
Async friendly.
|
|
"""
|
|
if self._closing:
|
|
# Connection is cancelled, don't flood logs about exceeding
|
|
# max pending messages.
|
|
return
|
|
|
|
if type(message) is not bytes: # noqa: E721
|
|
if isinstance(message, dict):
|
|
message = message_to_json_bytes(message)
|
|
elif isinstance(message, str):
|
|
message = message.encode("utf-8")
|
|
|
|
message_queue = self._message_queue
|
|
message_queue.append(message)
|
|
if (queue_size_after_add := len(message_queue)) >= MAX_PENDING_MSG:
|
|
self._logger.error(
|
|
(
|
|
"%s: Client unable to keep up with pending messages. Reached %s pending"
|
|
" messages. The system's load is too high or an integration is"
|
|
" misbehaving; Last message was: %s"
|
|
),
|
|
self.description,
|
|
MAX_PENDING_MSG,
|
|
message,
|
|
)
|
|
self._cancel()
|
|
return
|
|
|
|
if self._release_ready_queue_size == 0:
|
|
# Try to coalesce more messages to reduce the number of writes
|
|
self._release_ready_queue_size = queue_size_after_add
|
|
self._loop.call_soon(self._release_ready_future_or_reschedule)
|
|
|
|
peak_checker_active = self._peak_checker_unsub is not None
|
|
|
|
if queue_size_after_add <= PENDING_MSG_PEAK:
|
|
if peak_checker_active:
|
|
self._cancel_peak_checker()
|
|
return
|
|
|
|
if not peak_checker_active:
|
|
self._peak_checker_unsub = async_call_later(
|
|
self._hass, PENDING_MSG_PEAK_TIME, self._check_write_peak
|
|
)
|
|
|
|
@callback
|
|
def _release_ready_future_or_reschedule(self) -> None:
|
|
"""Release the ready future or reschedule.
|
|
|
|
We will release the ready future if the queue did not grow since the
|
|
last time we tried to release the ready future.
|
|
|
|
If we reach PENDING_MSG_MAX_FORCE_READY, we will release the ready future
|
|
immediately so avoid the coalesced messages from growing too large.
|
|
"""
|
|
if not (ready_future := self._ready_future) or not (
|
|
queue_size := len(self._message_queue)
|
|
):
|
|
self._release_ready_queue_size = 0
|
|
return
|
|
# If we are below the max pending to force ready, and there are new messages
|
|
# in the queue since the last time we tried to release the ready future, we
|
|
# try again later so we can coalesce more messages.
|
|
if queue_size > self._release_ready_queue_size < PENDING_MSG_MAX_FORCE_READY:
|
|
self._release_ready_queue_size = queue_size
|
|
self._loop.call_soon(self._release_ready_future_or_reschedule)
|
|
return
|
|
self._release_ready_queue_size = 0
|
|
if not ready_future.done():
|
|
ready_future.set_result(queue_size)
|
|
|
|
@callback
|
|
def _check_write_peak(self, _utc_time: dt.datetime) -> None:
|
|
"""Check that we are no longer above the write peak."""
|
|
self._peak_checker_unsub = None
|
|
|
|
if len(self._message_queue) < PENDING_MSG_PEAK:
|
|
return
|
|
|
|
self._logger.error(
|
|
(
|
|
"%s: Client unable to keep up with pending messages. Stayed over %s for %s"
|
|
" seconds. The system's load is too high or an integration is"
|
|
" misbehaving; Last message was: %s"
|
|
),
|
|
self.description,
|
|
PENDING_MSG_PEAK,
|
|
PENDING_MSG_PEAK_TIME,
|
|
self._message_queue[-1],
|
|
)
|
|
self._cancel()
|
|
|
|
@callback
|
|
def _cancel(self) -> None:
|
|
"""Cancel the connection."""
|
|
self._closing = True
|
|
self._cancel_peak_checker()
|
|
if self._handle_task is not None:
|
|
self._handle_task.cancel()
|
|
if self._writer_task is not None:
|
|
self._writer_task.cancel()
|
|
|
|
@callback
|
|
def _async_handle_hass_stop(self, event: Event) -> None:
|
|
"""Cancel this connection."""
|
|
self._cancel()
|
|
|
|
async def async_handle(self) -> web.WebSocketResponse:
|
|
"""Handle a websocket response."""
|
|
request = self._request
|
|
wsock = self._wsock
|
|
logger = self._logger
|
|
debug = logger.debug
|
|
hass = self._hass
|
|
is_enabled_for = logger.isEnabledFor
|
|
logging_debug = logging.DEBUG
|
|
|
|
try:
|
|
async with asyncio.timeout(10):
|
|
await wsock.prepare(request)
|
|
except TimeoutError:
|
|
self._logger.warning("Timeout preparing request from %s", request.remote)
|
|
return wsock
|
|
|
|
debug("%s: Connected from %s", self.description, request.remote)
|
|
self._handle_task = asyncio.current_task()
|
|
|
|
unsub_stop = hass.bus.async_listen(
|
|
EVENT_HOMEASSISTANT_STOP, self._async_handle_hass_stop
|
|
)
|
|
|
|
writer = wsock._writer # noqa: SLF001
|
|
if TYPE_CHECKING:
|
|
assert writer is not None
|
|
|
|
send_bytes_text = partial(writer.send, binary=False)
|
|
auth = AuthPhase(
|
|
logger, hass, self._send_message, self._cancel, request, send_bytes_text
|
|
)
|
|
connection = None
|
|
disconnect_warn = None
|
|
|
|
try:
|
|
await send_bytes_text(AUTH_REQUIRED_MESSAGE)
|
|
|
|
# Auth Phase
|
|
try:
|
|
msg = await wsock.receive(10)
|
|
except TimeoutError as err:
|
|
disconnect_warn = "Did not receive auth message within 10 seconds"
|
|
raise Disconnect from err
|
|
|
|
if msg.type in (WSMsgType.CLOSE, WSMsgType.CLOSED, WSMsgType.CLOSING):
|
|
raise Disconnect
|
|
|
|
if msg.type != WSMsgType.TEXT:
|
|
disconnect_warn = "Received non-Text message."
|
|
raise Disconnect
|
|
|
|
try:
|
|
auth_msg_data = json_loads(msg.data)
|
|
except ValueError as err:
|
|
disconnect_warn = "Received invalid JSON."
|
|
raise Disconnect from err
|
|
|
|
if is_enabled_for(logging_debug):
|
|
debug("%s: Received %s", self.description, auth_msg_data)
|
|
connection = await auth.async_handle(auth_msg_data)
|
|
# As the webserver is now started before the start
|
|
# event we do not want to block for websocket responses
|
|
#
|
|
# We only start the writer queue after the auth phase is completed
|
|
# since there is no need to queue messages before the auth phase
|
|
self._connection = connection
|
|
self._writer_task = create_eager_task(self._writer(send_bytes_text))
|
|
hass.data[DATA_CONNECTIONS] = hass.data.get(DATA_CONNECTIONS, 0) + 1
|
|
async_dispatcher_send(hass, SIGNAL_WEBSOCKET_CONNECTED)
|
|
|
|
self._authenticated = True
|
|
#
|
|
#
|
|
# Our websocket implementation is backed by a deque
|
|
#
|
|
# As back-pressure builds, the queue will back up and use more memory
|
|
# until we disconnect the client when the queue size reaches
|
|
# MAX_PENDING_MSG. When we are generating a high volume of websocket messages,
|
|
# we hit a bottleneck in aiohttp where it will wait for
|
|
# the buffer to drain before sending the next message and messages
|
|
# start backing up in the queue.
|
|
#
|
|
# https://github.com/aio-libs/aiohttp/issues/1367 added drains
|
|
# to the websocket writer to handle malicious clients and network issues.
|
|
# The drain causes multiple problems for us since the buffer cannot be
|
|
# drained fast enough when we deliver a high volume or large messages:
|
|
#
|
|
# - We end up disconnecting the client. The client will then reconnect,
|
|
# and the cycle repeats itself, which results in a significant amount of
|
|
# CPU usage.
|
|
#
|
|
# - Messages latency increases because messages cannot be moved into
|
|
# the TCP buffer because it is blocked waiting for the drain to happen because
|
|
# of the low default limit of 16KiB. By increasing the limit, we instead
|
|
# rely on the underlying TCP buffer and stack to deliver the messages which
|
|
# can typically happen much faster.
|
|
#
|
|
# After the auth phase is completed, and we are not concerned about
|
|
# the user being a malicious client, we set the limit to force a drain
|
|
# to 1MiB. 1MiB is the maximum expected size of the serialized entity
|
|
# registry, which is the largest message we usually send.
|
|
#
|
|
# https://github.com/aio-libs/aiohttp/commit/b3c80ee3f7d5d8f0b8bc27afe52e4d46621eaf99
|
|
# added a way to set the limit, but there is no way to actually
|
|
# reach the code to set the limit, so we have to set it directly.
|
|
#
|
|
writer._limit = 2**20 # noqa: SLF001
|
|
async_handle_str = connection.async_handle
|
|
async_handle_binary = connection.async_handle_binary
|
|
|
|
# Command phase
|
|
while not wsock.closed:
|
|
msg = await wsock.receive()
|
|
|
|
if msg.type in (WSMsgType.CLOSE, WSMsgType.CLOSED, WSMsgType.CLOSING):
|
|
break
|
|
|
|
if msg.type is WSMsgType.BINARY:
|
|
if len(msg.data) < 1:
|
|
disconnect_warn = "Received invalid binary message."
|
|
break
|
|
handler = msg.data[0]
|
|
payload = msg.data[1:]
|
|
async_handle_binary(handler, payload)
|
|
continue
|
|
|
|
if msg.type is not WSMsgType.TEXT:
|
|
disconnect_warn = "Received non-Text message."
|
|
break
|
|
|
|
try:
|
|
command_msg_data = json_loads(msg.data)
|
|
except ValueError:
|
|
disconnect_warn = "Received invalid JSON."
|
|
break
|
|
|
|
if is_enabled_for(logging_debug):
|
|
debug("%s: Received %s", self.description, command_msg_data)
|
|
|
|
# command_msg_data is always deserialized from JSON as a list
|
|
if type(command_msg_data) is not list: # noqa: E721
|
|
async_handle_str(command_msg_data)
|
|
continue
|
|
|
|
for split_msg in command_msg_data:
|
|
async_handle_str(split_msg)
|
|
|
|
except asyncio.CancelledError:
|
|
debug("%s: Connection cancelled", self.description)
|
|
raise
|
|
|
|
except Disconnect as ex:
|
|
debug("%s: Connection closed by client: %s", self.description, ex)
|
|
|
|
except Exception:
|
|
self._logger.exception(
|
|
"%s: Unexpected error inside websocket API", self.description
|
|
)
|
|
|
|
finally:
|
|
unsub_stop()
|
|
|
|
self._cancel_peak_checker()
|
|
|
|
if connection is not None:
|
|
connection.async_handle_close()
|
|
|
|
self._closing = True
|
|
if self._ready_future and not self._ready_future.done():
|
|
self._ready_future.set_result(len(self._message_queue))
|
|
|
|
# If the writer gets canceled we still need to close the websocket
|
|
# so we have another finally block to make sure we close the websocket
|
|
# if the writer gets canceled.
|
|
try:
|
|
if self._writer_task:
|
|
await self._writer_task
|
|
finally:
|
|
try:
|
|
# Make sure all error messages are written before closing
|
|
await wsock.close()
|
|
finally:
|
|
if disconnect_warn is None:
|
|
debug("%s: Disconnected", self.description)
|
|
else:
|
|
self._logger.warning(
|
|
"%s: Disconnected: %s", self.description, disconnect_warn
|
|
)
|
|
|
|
if connection is not None:
|
|
hass.data[DATA_CONNECTIONS] -= 1
|
|
self._connection = None
|
|
|
|
async_dispatcher_send(hass, SIGNAL_WEBSOCKET_DISCONNECTED)
|
|
|
|
# Break reference cycles to make sure GC can happen sooner
|
|
self._wsock = None # type: ignore[assignment]
|
|
self._request = None # type: ignore[assignment]
|
|
self._hass = None # type: ignore[assignment]
|
|
self._logger = None # type: ignore[assignment]
|
|
self._message_queue = None # type: ignore[assignment]
|
|
self._handle_task = None
|
|
self._writer_task = None
|
|
self._ready_future = None
|
|
|
|
return wsock
|