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])
[docs]
@attr.s(auto_attribs=True, frozen=True, slots=True)
class WebSocketReady:
ok: bool
protocol: str | None
def __bool__(self) -> bool:
return self.ok
[docs]
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()
[docs]
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)
[docs]
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
[docs]
def exception(self) -> BaseException | None:
return self._exception
[docs]
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)
[docs]
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)
[docs]
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)
[docs]
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
)
[docs]
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)
[docs]
async def send_json(
self,
data: Any,
compress: int | None = None,
*,
dumps: JSONEncoder = json.dumps,
) -> None:
await self.send_str(dumps(data), compress=compress)
[docs]
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
[docs]
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: ...
[docs]
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: ...
[docs]
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)
[docs]
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: ...
[docs]
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)
