import asyncio import base64 import binascii import hashlib import json import sys from asyncio.base_events import BaseEventLoop from collections.abc import Callable, Iterable from typing import Any, Final, Generic, Literal, cast, overload import attr from multidict import CIMultiDict from . import hdrs from ._websocket.reader import WebSocketDataQueue from .abc import AbstractStreamWriter from .client_exceptions import WSMessageTypeError from .helpers import ( DEFAULT_CHUNK_SIZE, calculate_timeout_when, set_exception, set_result, ) from .http import ( WS_CLOSED_MESSAGE, WS_CLOSING_MESSAGE, WS_KEY, WebSocketError, WebSocketReader, WebSocketWriter, WSCloseCode, WSMessage, WSMessageDecodeText, WSMessageNoDecodeText, WSMsgType as WSMsgType, ws_ext_gen, ws_ext_parse, ) from .http_websocket import _INTERNAL_RECEIVE_TYPES from .log import ws_logger from .streams import EofStream from .typedefs import JSONBytesEncoder, JSONDecoder, JSONEncoder from .web_exceptions import HTTPBadRequest, HTTPException from .web_request import BaseRequest from .web_response import StreamResponse if sys.version_info >= (3, 13): from typing import TypeVar else: from typing_extensions import TypeVar if sys.version_info >= (3, 12): from collections.abc import Buffer else: from typing import Union Buffer = Union[bytes, bytearray, "memoryview[int]", "memoryview[bytes]"] if sys.version_info >= (3, 11): import asyncio as async_timeout from typing import Self else: import async_timeout from typing_extensions import Self __all__ = ( "WebSocketResponse", "WebSocketReady", "WSMsgType", ) THRESHOLD_CONNLOST_ACCESS: Final[int] = 5 # TypeVar for whether text messages are decoded to str (True) or kept as bytes (False) _DecodeText = TypeVar("_DecodeText", bound=bool, covariant=True, default=Literal[True]) @attr.s(auto_attribs=True, frozen=True, slots=True) class WebSocketReady: ok: bool protocol: str | None def __bool__(self) -> bool: return self.ok class WebSocketResponse(StreamResponse, Generic[_DecodeText]): _length_check: bool = False _ws_protocol: str | None = None _writer: WebSocketWriter | None = None _reader: WebSocketDataQueue | None = None _parser: WebSocketReader | None = None _closed: bool = False _closing: bool = False _conn_lost: int = 0 _close_code: int | None = None _loop: asyncio.AbstractEventLoop | None = None _waiting: bool = False _close_wait: asyncio.Future[None] | None = None _exception: BaseException | None = None _heartbeat_when: float = 0.0 _heartbeat_cb: asyncio.TimerHandle | None = None _pong_response_cb: asyncio.TimerHandle | None = None _ping_task: asyncio.Task[None] | None = None _need_heartbeat_reset: bool = False _heartbeat_reset_handle: asyncio.Handle | None = None def __init__( self, *, timeout: float = 10.0, receive_timeout: float | None = None, autoclose: bool = True, autoping: bool = True, heartbeat: float | None = None, protocols: Iterable[str] = (), compress: bool = True, max_msg_size: int = 4 * 1024 * 1024, writer_limit: int = DEFAULT_CHUNK_SIZE, decode_text: bool = True, ) -> None: super().__init__(status=101) self._protocols = protocols self._timeout = timeout self._receive_timeout = receive_timeout self._autoclose = autoclose self._autoping = autoping self._heartbeat = heartbeat if heartbeat is not None: self._pong_heartbeat = heartbeat / 2.0 self._compress: bool | int = compress self._max_msg_size = max_msg_size self._writer_limit = writer_limit self._decode_text = decode_text self._need_heartbeat_reset = False self._heartbeat_reset_handle = None def _cancel_heartbeat(self) -> None: self._cancel_pong_response_cb() if self._heartbeat_reset_handle is not None: self._heartbeat_reset_handle.cancel() self._heartbeat_reset_handle = None self._need_heartbeat_reset = False if self._heartbeat_cb is not None: self._heartbeat_cb.cancel() self._heartbeat_cb = None if self._ping_task is not None: self._ping_task.cancel() self._ping_task = None def _cancel_pong_response_cb(self) -> None: if self._pong_response_cb is not None: self._pong_response_cb.cancel() self._pong_response_cb = None def _on_data_received(self) -> None: if self._heartbeat is None or self._need_heartbeat_reset: return loop = self._loop assert loop is not None # Coalesce multiple chunks received in the same loop tick into a single # heartbeat reset. Resetting immediately per chunk increases timer churn. self._need_heartbeat_reset = True self._heartbeat_reset_handle = loop.call_soon(self._flush_heartbeat_reset) def _flush_heartbeat_reset(self) -> None: self._heartbeat_reset_handle = None if not self._need_heartbeat_reset: return self._reset_heartbeat() self._need_heartbeat_reset = False def _reset_heartbeat(self) -> None: if self._heartbeat is None: return self._cancel_pong_response_cb() req = self._req timeout_ceil_threshold = ( req._protocol._timeout_ceil_threshold if req is not None else 5 ) loop = self._loop assert loop is not None now = loop.time() when = calculate_timeout_when(now, self._heartbeat, timeout_ceil_threshold) self._heartbeat_when = when if self._heartbeat_cb is None: # We do not cancel the previous heartbeat_cb here because # it generates a significant amount of TimerHandle churn # which causes asyncio to rebuild the heap frequently. # Instead _send_heartbeat() will reschedule the next # heartbeat if it fires too early. self._heartbeat_cb = loop.call_at(when, self._send_heartbeat) def _send_heartbeat(self) -> None: self._heartbeat_cb = None # If heartbeat reset is pending (data is being received), skip sending # the ping and let the reset callback handle rescheduling the heartbeat. if self._need_heartbeat_reset: return loop = self._loop assert loop is not None and self._writer is not None now = loop.time() if now < self._heartbeat_when: # Heartbeat fired too early, reschedule self._heartbeat_cb = loop.call_at( self._heartbeat_when, self._send_heartbeat ) return req = self._req timeout_ceil_threshold = ( req._protocol._timeout_ceil_threshold if req is not None else 5 ) when = calculate_timeout_when(now, self._pong_heartbeat, timeout_ceil_threshold) self._cancel_pong_response_cb() self._pong_response_cb = loop.call_at(when, self._pong_not_received) coro = self._writer.send_frame(b"", WSMsgType.PING) if sys.version_info >= (3, 14): # Try to send the ping immediately to avoid having to schedule # the task on the event loop. if isinstance(loop, BaseEventLoop): ping_task = asyncio.create_task(coro, eager_start=True) else: ping_task = asyncio.Task(coro, loop=loop, eager_start=True) elif sys.version_info >= (3, 12): ping_task = asyncio.Task(coro, loop=loop, eager_start=True) else: ping_task = asyncio.create_task(coro) if not ping_task.done(): self._ping_task = ping_task ping_task.add_done_callback(self._ping_task_done) else: self._ping_task_done(ping_task) def _ping_task_done(self, task: "asyncio.Task[None]") -> None: """Callback for when the ping task completes.""" if not task.cancelled() and (exc := task.exception()): self._handle_ping_pong_exception(exc) self._ping_task = None def _pong_not_received(self) -> None: if self._req is not None and self._req.transport is not None: self._handle_ping_pong_exception( asyncio.TimeoutError( f"No PONG received after {self._pong_heartbeat} seconds" ) ) def _handle_ping_pong_exception(self, exc: BaseException) -> None: """Handle exceptions raised during ping/pong processing.""" if self._closed: return self._set_closed() # close() is never reached after this; release the parser here. self._parser = None self._set_code_close_transport(WSCloseCode.ABNORMAL_CLOSURE) self._exception = exc if self._waiting and not self._closing and self._reader is not None: self._reader.feed_data(WSMessage(WSMsgType.ERROR, exc, None), 0) def _set_closed(self) -> None: """Set the connection to closed. Cancel any heartbeat timers and set the closed flag. """ self._closed = True self._cancel_heartbeat() async def prepare(self, request: BaseRequest) -> AbstractStreamWriter: # make pre-check to don't hide it by do_handshake() exceptions if self._payload_writer is not None: return self._payload_writer protocol, writer = self._pre_start(request) payload_writer = await super().prepare(request) assert payload_writer is not None self._post_start(request, protocol, writer) await payload_writer.drain() return payload_writer def _handshake( self, request: BaseRequest ) -> tuple["CIMultiDict[str]", str | None, int, bool]: headers = request.headers if "websocket" != headers.get(hdrs.UPGRADE, "").lower().strip(): raise HTTPBadRequest( text=( f"No WebSocket UPGRADE hdr: {headers.get(hdrs.UPGRADE)}\n Can " '"Upgrade" only to "WebSocket".' ) ) if not request._message.upgrade: raise HTTPBadRequest( text=f"No CONNECTION upgrade hdr: {headers.get(hdrs.CONNECTION)}" ) # find common sub-protocol between client and server protocol: str | None = None if hdrs.SEC_WEBSOCKET_PROTOCOL in headers: req_protocols = [ str(proto.strip()) for proto in headers[hdrs.SEC_WEBSOCKET_PROTOCOL].split(",") ] for proto in req_protocols: if proto in self._protocols: protocol = proto break else: # No overlap found: Return no protocol as per spec ws_logger.warning( "%s: Client protocols %r don’t overlap server-known ones %r", request.remote, req_protocols, self._protocols, ) # check supported version version = headers.get(hdrs.SEC_WEBSOCKET_VERSION, "") if version not in ("13", "8", "7"): raise HTTPBadRequest(text=f"Unsupported version: {version}") # check client handshake for validity key = headers.get(hdrs.SEC_WEBSOCKET_KEY) try: if not key or len(base64.b64decode(key)) != 16: raise HTTPBadRequest(text=f"Handshake error: {key!r}") except binascii.Error: raise HTTPBadRequest(text=f"Handshake error: {key!r}") from None accept_val = base64.b64encode( hashlib.sha1(key.encode() + WS_KEY).digest() ).decode() response_headers = CIMultiDict( { hdrs.UPGRADE: "websocket", hdrs.CONNECTION: "upgrade", hdrs.SEC_WEBSOCKET_ACCEPT: accept_val, } ) notakeover = False compress = 0 if self._compress: extensions = headers.get(hdrs.SEC_WEBSOCKET_EXTENSIONS) # Server side always get return with no exception. # If something happened, just drop compress extension compress, notakeover = ws_ext_parse(extensions, isserver=True) if compress: enabledext = ws_ext_gen( compress=compress, isserver=True, server_notakeover=notakeover ) response_headers[hdrs.SEC_WEBSOCKET_EXTENSIONS] = enabledext if protocol: response_headers[hdrs.SEC_WEBSOCKET_PROTOCOL] = protocol return ( response_headers, protocol, compress, notakeover, ) def _pre_start(self, request: BaseRequest) -> tuple[str | None, WebSocketWriter]: self._loop = request._loop headers, protocol, compress, notakeover = self._handshake(request) self.set_status(101) self.headers.update(headers) self.force_close() self._compress = compress transport = request._protocol.transport if transport is None: raise ConnectionResetError("Connection lost") writer = WebSocketWriter( request._protocol, transport, compress=compress, notakeover=notakeover, limit=self._writer_limit, ) return protocol, writer def _post_start( self, request: BaseRequest, protocol: str | None, writer: WebSocketWriter ) -> None: self._ws_protocol = protocol self._writer = writer self._reset_heartbeat() loop = self._loop assert loop is not None self._reader = WebSocketDataQueue( request._protocol, DEFAULT_CHUNK_SIZE, loop=loop ) # Owns the parser so a stalled reader parked on the queue by weakref # stays alive while this response can still be drained. self._parser = WebSocketReader( self._reader, self._max_msg_size, compress=bool(self._compress), decode_text=self._decode_text, ) cb = None if self._heartbeat is None else self._on_data_received request.protocol.set_parser(self._parser, data_received_cb=cb) # disable HTTP keepalive for WebSocket request.protocol.keep_alive(False) def can_prepare(self, request: BaseRequest) -> WebSocketReady: if self._writer is not None: raise RuntimeError("Already started") try: _, protocol, _, _ = self._handshake(request) except HTTPException: return WebSocketReady(False, None) else: return WebSocketReady(True, protocol) @property def prepared(self) -> bool: return self._writer is not None @property def closed(self) -> bool: return self._closed @property def close_code(self) -> int | None: return self._close_code @property def ws_protocol(self) -> str | None: return self._ws_protocol @property def compress(self) -> int | bool: return self._compress def get_extra_info(self, name: str, default: Any = None) -> Any: """Get optional transport information. If no value associated with ``name`` is found, ``default`` is returned. """ writer = self._writer if writer is None: return default transport = writer.transport if transport is None: return default return transport.get_extra_info(name, default) def exception(self) -> BaseException | None: return self._exception async def ping(self, message: bytes = b"") -> None: if self._writer is None: raise RuntimeError("Call .prepare() first") await self._writer.send_frame(message, WSMsgType.PING) async def pong(self, message: bytes = b"") -> None: # unsolicited pong if self._writer is None: raise RuntimeError("Call .prepare() first") await self._writer.send_frame(message, WSMsgType.PONG) async def send_frame( self, message: bytes, opcode: WSMsgType, compress: int | None = None ) -> None: """Send a frame over the websocket.""" if self._writer is None: raise RuntimeError("Call .prepare() first") await self._writer.send_frame(message, opcode, compress) async def send_str(self, data: str, compress: int | None = None) -> None: if self._writer is None: raise RuntimeError("Call .prepare() first") if not isinstance(data, str): raise TypeError("data argument must be str (%r)" % type(data)) await self._writer.send_frame( data.encode("utf-8"), WSMsgType.TEXT, compress=compress ) async def send_bytes(self, data: bytes, compress: int | None = None) -> None: if self._writer is None: raise RuntimeError("Call .prepare() first") if not isinstance(data, (bytes, bytearray, memoryview)): raise TypeError("data argument must be byte-ish (%r)" % type(data)) await self._writer.send_frame(data, WSMsgType.BINARY, compress=compress) async def send_json( self, data: Any, compress: int | None = None, *, dumps: JSONEncoder = json.dumps, ) -> None: await self.send_str(dumps(data), compress=compress) async def send_json_bytes( self, data: Any, compress: int | None = None, *, dumps: JSONBytesEncoder, ) -> None: """Send JSON data using a bytes-returning encoder as a binary frame. Use this when your JSON encoder (like orjson) returns bytes instead of str, avoiding the encode/decode overhead. """ await self.send_bytes(dumps(data), compress=compress) async def write_eof(self) -> None: # type: ignore[override] if self._eof_sent: return if self._payload_writer is None: raise RuntimeError("Response has not been started") await self.close() self._eof_sent = True async def close( self, *, code: int = WSCloseCode.OK, message: bytes = b"", drain: bool = True ) -> bool: """Close websocket connection.""" if self._writer is None: raise RuntimeError("Call .prepare() first") if self._closed: return False self._set_closed() try: try: await self._writer.close(code, message) writer = self._payload_writer assert writer is not None if drain: await writer.drain() except (asyncio.CancelledError, asyncio.TimeoutError): self._set_code_close_transport(WSCloseCode.ABNORMAL_CLOSURE) raise except Exception as exc: self._exception = exc self._set_code_close_transport(WSCloseCode.ABNORMAL_CLOSURE) return True reader = self._reader assert reader is not None # we need to break `receive()` cycle before we can call # `reader.read()` as `close()` may be called from different task if self._waiting: assert self._loop is not None assert self._close_wait is None self._close_wait = self._loop.create_future() reader.feed_data(WS_CLOSING_MESSAGE, 0) await self._close_wait if self._closing: self._close_transport() return True try: async with async_timeout.timeout(self._timeout): while True: msg = await reader.read() if msg.type is WSMsgType.CLOSE: self._set_code_close_transport(msg.data) return True except asyncio.CancelledError: self._set_code_close_transport(WSCloseCode.ABNORMAL_CLOSURE) raise except Exception as exc: self._exception = exc self._set_code_close_transport(WSCloseCode.ABNORMAL_CLOSURE) return True finally: # Once closed the response can no longer be drained; release the # parser and the stash it retains. self._parser = None def _set_closing(self, code: WSCloseCode) -> None: """Set the close code and mark the connection as closing.""" self._closing = True self._close_code = code self._cancel_heartbeat() def _set_code_close_transport(self, code: WSCloseCode) -> None: """Set the close code and close the transport.""" self._close_code = code self._close_transport() def _close_transport(self) -> None: """Close the transport.""" if self._req is not None and self._req.transport is not None: self._req.transport.close() @overload async def receive( self: "WebSocketResponse[Literal[True]]", timeout: float | None = None ) -> WSMessageDecodeText: ... @overload async def receive( self: "WebSocketResponse[Literal[False]]", timeout: float | None = None ) -> WSMessageNoDecodeText: ... @overload async def receive( self: "WebSocketResponse[_DecodeText]", timeout: float | None = None ) -> WSMessageDecodeText | WSMessageNoDecodeText: ... async def receive( self, timeout: float | None = None ) -> WSMessageDecodeText | WSMessageNoDecodeText: if self._reader is None: raise RuntimeError("Call .prepare() first") receive_timeout = timeout or self._receive_timeout while True: if self._waiting: raise RuntimeError("Concurrent call to receive() is not allowed") if self._closed: self._conn_lost += 1 if self._conn_lost >= THRESHOLD_CONNLOST_ACCESS: raise RuntimeError("WebSocket connection is closed.") return WS_CLOSED_MESSAGE elif self._closing: return WS_CLOSING_MESSAGE try: self._waiting = True try: if receive_timeout: # Entering the context manager and creating # Timeout() object can take almost 50% of the # run time in this loop so we avoid it if # there is no read timeout. async with async_timeout.timeout(receive_timeout): msg = await self._reader.read() else: msg = await self._reader.read() finally: self._waiting = False if self._close_wait: set_result(self._close_wait, None) except asyncio.TimeoutError: raise except EofStream: self._close_code = WSCloseCode.OK await self.close() return WSMessage(WSMsgType.CLOSED, None, None) except WebSocketError as exc: self._close_code = exc.code await self.close(code=exc.code) return WSMessage(WSMsgType.ERROR, exc, None) except Exception as exc: self._exception = exc self._set_closing(WSCloseCode.ABNORMAL_CLOSURE) await self.close() return WSMessage(WSMsgType.ERROR, exc, None) if msg.type not in _INTERNAL_RECEIVE_TYPES: # If its not a close/closing/ping/pong message # we can return it immediately return msg if msg.type is WSMsgType.CLOSE: self._set_closing(msg.data) # Could be closed while awaiting reader. if not self._closed and self._autoclose: # The client is likely going to close the # connection out from under us so we do not # want to drain any pending writes as it will # likely result writing to a broken pipe. await self.close(drain=False) elif msg.type is WSMsgType.CLOSING: self._set_closing(WSCloseCode.OK) elif msg.type is WSMsgType.PING and self._autoping: await self.pong(msg.data) continue elif msg.type is WSMsgType.PONG and self._autoping: continue return msg @overload async def receive_str( self: "WebSocketResponse[Literal[True]]", *, timeout: float | None = None ) -> str: ... @overload async def receive_str( self: "WebSocketResponse[Literal[False]]", *, timeout: float | None = None ) -> bytes: ... @overload async def receive_str( self: "WebSocketResponse[_DecodeText]", *, timeout: float | None = None ) -> str | bytes: ... async def receive_str(self, *, timeout: float | None = None) -> str | bytes: """Receive TEXT message. Returns str when decode_text=True (default), bytes when decode_text=False. """ msg = await self.receive(timeout) if msg.type is not WSMsgType.TEXT: raise WSMessageTypeError( f"Received message {msg.type}:{msg.data!r} is not WSMsgType.TEXT" ) return cast(str, msg.data) async def receive_bytes(self, *, timeout: float | None = None) -> bytes: msg = await self.receive(timeout) if msg.type is not WSMsgType.BINARY: raise WSMessageTypeError( f"Received message {msg.type}:{msg.data!r} is not WSMsgType.BINARY" ) return cast(bytes, msg.data) @overload async def receive_json( self: "WebSocketResponse[Literal[True]]", *, loads: JSONDecoder = ..., timeout: float | None = None, ) -> Any: ... @overload async def receive_json( self: "WebSocketResponse[Literal[False]]", *, loads: Callable[[bytes], Any] = ..., timeout: float | None = None, ) -> Any: ... @overload async def receive_json( self: "WebSocketResponse[_DecodeText]", *, loads: JSONDecoder | Callable[[bytes], Any] = ..., timeout: float | None = None, ) -> Any: ... async def receive_json( self, *, loads: JSONDecoder | Callable[[bytes], Any] = json.loads, timeout: float | None = None, ) -> Any: data = await self.receive_str(timeout=timeout) return loads(data) # type: ignore[arg-type] async def write(self, data: Buffer) -> None: raise RuntimeError("Cannot call .write() for websocket") def __aiter__(self) -> Self: return self @overload async def __anext__( self: "WebSocketResponse[Literal[True]]", ) -> WSMessageDecodeText: ... @overload async def __anext__( self: "WebSocketResponse[Literal[False]]", ) -> WSMessageNoDecodeText: ... @overload async def __anext__( self: "WebSocketResponse[_DecodeText]", ) -> WSMessageDecodeText | WSMessageNoDecodeText: ... async def __anext__(self) -> WSMessageDecodeText | WSMessageNoDecodeText: msg = await self.receive() if msg.type in (WSMsgType.CLOSE, WSMsgType.CLOSING, WSMsgType.CLOSED): raise StopAsyncIteration return msg def _cancel(self, exc: BaseException) -> None: # web_protocol calls this from connection_lost # or when the server is shutting down. self._closing = True self._cancel_heartbeat() if self._reader is not None: set_exception(self._reader, exc)