import asyncio
import base64
import binascii
import hashlib
import json
import sys
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
_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, 12):
# Optimization for Python 3.12, try to send the ping
# immediately to avoid having to schedule
# the task on the event loop.
ping_task = asyncio.Task(coro, loop=loop, eager_start=True)
else:
ping_task = loop.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()
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
)
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(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:
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
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)
