Beta/venv/lib/python3.12/site-packages/aiohttp/web_ws.py

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


@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
    _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()

    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)

    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:
            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: ...

    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)