protocol.py

from __future__ import annotations

import enum
import logging
import uuid
from typing import Generator, List, Optional, Type, Union

from .exceptions import (
    ConnectionClosed,
    ConnectionClosedError,
    ConnectionClosedOK,
    InvalidState,
    PayloadTooBig,
    ProtocolError,
)
from .extensions import Extension
from .frames import (
    OK_CLOSE_CODES,
    OP_BINARY,
    OP_CLOSE,
    OP_CONT,
    OP_PING,
    OP_PONG,
    OP_TEXT,
    Close,
    CloseCode,
    Frame,
)
from .http11 import Request, Response
from .streams import StreamReader
from .typing import LoggerLike, Origin, Subprotocol


__all__ = [
    "Protocol",
    "Side",
    "State",
    "SEND_EOF",
]

Event = Union[Request, Response, Frame]
"""Events that :meth:`~Protocol.events_received` may return."""


class Side(enum.IntEnum):
    """A WebSocket connection is either a server or a client."""

    SERVER, CLIENT = range(2)


SERVER = Side.SERVER
CLIENT = Side.CLIENT


class State(enum.IntEnum):
    """A WebSocket connection is in one of these four states."""

    CONNECTING, OPEN, CLOSING, CLOSED = range(4)


CONNECTING = State.CONNECTING
OPEN = State.OPEN
CLOSING = State.CLOSING
CLOSED = State.CLOSED


SEND_EOF = b""
"""Sentinel signaling that the TCP connection must be half-closed."""


class Protocol:
    """
    Sans-I/O implementation of a WebSocket connection.

    Args:
        side: :attr:`~Side.CLIENT` or :attr:`~Side.SERVER`.
        state: initial state of the WebSocket connection.
        max_size: maximum size of incoming messages in bytes;
            :obj:`None` disables the limit.
        logger: logger for this connection; depending on ``side``,
            defaults to ``logging.getLogger("websockets.client")``
            or ``logging.getLogger("websockets.server")``;
            see the :doc:`logging guide <../../topics/logging>` for details.

    """

    def __init__(
        self,
        side: Side,
        *,
        state: State = OPEN,
        max_size: Optional[int] = 2**20,
        logger: Optional[LoggerLike] = None,
    ) -> None:
        # Unique identifier. For logs.
        self.id: uuid.UUID = uuid.uuid4()
        """Unique identifier of the connection. Useful in logs."""

        # Logger or LoggerAdapter for this connection.
        if logger is None:
            logger = logging.getLogger(f"websockets.{side.name.lower()}")
        self.logger: LoggerLike = logger
        """Logger for this connection."""

        # Track if DEBUG is enabled. Shortcut logging calls if it isn't.
        self.debug = logger.isEnabledFor(logging.DEBUG)

        # Connection side. CLIENT or SERVER.
        self.side = side

        # Connection state. Initially OPEN because subclasses handle CONNECTING.
        self.state = state

        # Maximum size of incoming messages in bytes.
        self.max_size = max_size

        # Current size of incoming message in bytes. Only set while reading a
        # fragmented message i.e. a data frames with the FIN bit not set.
        self.cur_size: Optional[int] = None

        # True while sending a fragmented message i.e. a data frames with the
        # FIN bit not set.
        self.expect_continuation_frame = False

        # WebSocket protocol parameters.
        self.origin: Optional[Origin] = None
        self.extensions: List[Extension] = []
        self.subprotocol: Optional[Subprotocol] = None

        # Close code and reason, set when a close frame is sent or received.
        self.close_rcvd: Optional[Close] = None
        self.close_sent: Optional[Close] = None
        self.close_rcvd_then_sent: Optional[bool] = None

        # Track if an exception happened during the handshake.
        self.handshake_exc: Optional[Exception] = None
        """
        Exception to raise if the opening handshake failed.

        :obj:`None` if the opening handshake succeeded.

        """

        # Track if send_eof() was called.
        self.eof_sent = False

        # Parser state.
        self.reader = StreamReader()
        self.events: List[Event] = []
        self.writes: List[bytes] = []
        self.parser = self.parse()
        next(self.parser)  # start coroutine
        self.parser_exc: Optional[Exception] = None

    @property
    def state(self) -> State:
        """
        WebSocket connection state.

        Defined in 4.1, 4.2, 7.1.3, and 7.1.4 of :rfc:`6455`.

        """
        return self._state

    @state.setter
    def state(self, state: State) -> None:
        if self.debug:
            self.logger.debug("= connection is %s", state.name)
        self._state = state

    @property
    def close_code(self) -> Optional[int]:
        """
        `WebSocket close code`_.

        .. _WebSocket close code:
            https://www.rfc-editor.org/rfc/rfc6455.html#section-7.1.5

        :obj:`None` if the connection isn't closed yet.

        """
        if self.state is not CLOSED:
            return None
        elif self.close_rcvd is None:
            return CloseCode.ABNORMAL_CLOSURE
        else:
            return self.close_rcvd.code

    @property
    def close_reason(self) -> Optional[str]:
        """
        `WebSocket close reason`_.

        .. _WebSocket close reason:
            https://www.rfc-editor.org/rfc/rfc6455.html#section-7.1.6

        :obj:`None` if the connection isn't closed yet.

        """
        if self.state is not CLOSED:
            return None
        elif self.close_rcvd is None:
            return ""
        else:
            return self.close_rcvd.reason

    @property
    def close_exc(self) -> ConnectionClosed:
        """
        Exception to raise when trying to interact with a closed connection.

        Don't raise this exception while the connection :attr:`state`
        is :attr:`~websockets.protocol.State.CLOSING`; wait until
        it's :attr:`~websockets.protocol.State.CLOSED`.

        Indeed, the exception includes the close code and reason, which are
        known only once the connection is closed.

        Raises:
            AssertionError: if the connection isn't closed yet.

        """
        assert self.state is CLOSED, "connection isn't closed yet"
        exc_type: Type[ConnectionClosed]
        if (
            self.close_rcvd is not None
            and self.close_sent is not None
            and self.close_rcvd.code in OK_CLOSE_CODES
            and self.close_sent.code in OK_CLOSE_CODES
        ):
            exc_type = ConnectionClosedOK
        else:
            exc_type = ConnectionClosedError
        exc: ConnectionClosed = exc_type(
            self.close_rcvd,
            self.close_sent,
            self.close_rcvd_then_sent,
        )
        # Chain to the exception raised in the parser, if any.
        exc.__cause__ = self.parser_exc
        return exc

    # Public methods for receiving data.

    def receive_data(self, data: bytes) -> None:
        """
        Receive data from the network.

        After calling this method:

        - You must call :meth:`data_to_send` and send this data to the network.
        - You should call :meth:`events_received` and process resulting events.

        Raises:
            EOFError: if :meth:`receive_eof` was called earlier.

        """
        self.reader.feed_data(data)
        next(self.parser)

    def receive_eof(self) -> None:
        """
        Receive the end of the data stream from the network.

        After calling this method:

        - You must call :meth:`data_to_send` and send this data to the network;
          it will return ``[b""]``, signaling the end of the stream, or ``[]``.
        - You aren't expected to call :meth:`events_received`; it won't return
          any new events.

        Raises:
            EOFError: if :meth:`receive_eof` was called earlier.

        """
        self.reader.feed_eof()
        next(self.parser)

    # Public methods for sending events.

    def send_continuation(self, data: bytes, fin: bool) -> None:
        """
        Send a `Continuation frame`_.

        .. _Continuation frame:
            https://datatracker.ietf.org/doc/html/rfc6455#section-5.6

        Parameters:
            data: payload containing the same kind of data
                as the initial frame.
            fin: FIN bit; set it to :obj:`True` if this is the last frame
                of a fragmented message and to :obj:`False` otherwise.

        Raises:
            ProtocolError: if a fragmented message isn't in progress.

        """
        if not self.expect_continuation_frame:
            raise ProtocolError("unexpected continuation frame")
        self.expect_continuation_frame = not fin
        self.send_frame(Frame(OP_CONT, data, fin))

    def send_text(self, data: bytes, fin: bool = True) -> None:
        """
        Send a `Text frame`_.

        .. _Text frame:
            https://datatracker.ietf.org/doc/html/rfc6455#section-5.6

        Parameters:
            data: payload containing text encoded with UTF-8.
            fin: FIN bit; set it to :obj:`False` if this is the first frame of
                a fragmented message.

        Raises:
            ProtocolError: if a fragmented message is in progress.

        """
        if self.expect_continuation_frame:
            raise ProtocolError("expected a continuation frame")
        self.expect_continuation_frame = not fin
        self.send_frame(Frame(OP_TEXT, data, fin))

    def send_binary(self, data: bytes, fin: bool = True) -> None:
        """
        Send a `Binary frame`_.

        .. _Binary frame:
            https://datatracker.ietf.org/doc/html/rfc6455#section-5.6

        Parameters:
            data: payload containing arbitrary binary data.
            fin: FIN bit; set it to :obj:`False` if this is the first frame of
                a fragmented message.

        Raises:
            ProtocolError: if a fragmented message is in progress.

        """
        if self.expect_continuation_frame:
            raise ProtocolError("expected a continuation frame")
        self.expect_continuation_frame = not fin
        self.send_frame(Frame(OP_BINARY, data, fin))

    def send_close(self, code: Optional[int] = None, reason: str = "") -> None:
        """
        Send a `Close frame`_.

        .. _Close frame:
            https://datatracker.ietf.org/doc/html/rfc6455#section-5.5.1

        Parameters:
            code: close code.
            reason: close reason.

        Raises:
            ProtocolError: if a fragmented message is being sent, if the code
                isn't valid, or if a reason is provided without a code

        """
        if self.expect_continuation_frame:
            raise ProtocolError("expected a continuation frame")
        if code is None:
            if reason != "":
                raise ProtocolError("cannot send a reason without a code")
            close = Close(CloseCode.NO_STATUS_RCVD, "")
            data = b""
        else:
            close = Close(code, reason)
            data = close.serialize()
        # send_frame() guarantees that self.state is OPEN at this point.
        # 7.1.3. The WebSocket Closing Handshake is Started
        self.send_frame(Frame(OP_CLOSE, data))
        self.close_sent = close
        self.state = CLOSING

    def send_ping(self, data: bytes) -> None:
        """
        Send a `Ping frame`_.

        .. _Ping frame:
            https://datatracker.ietf.org/doc/html/rfc6455#section-5.5.2

        Parameters:
            data: payload containing arbitrary binary data.

        """
        self.send_frame(Frame(OP_PING, data))

    def send_pong(self, data: bytes) -> None:
        """
        Send a `Pong frame`_.

        .. _Pong frame:
            https://datatracker.ietf.org/doc/html/rfc6455#section-5.5.3

        Parameters:
            data: payload containing arbitrary binary data.

        """
        self.send_frame(Frame(OP_PONG, data))

    def fail(self, code: int, reason: str = "") -> None:
        """
        `Fail the WebSocket connection`_.

        .. _Fail the WebSocket connection:
            https://datatracker.ietf.org/doc/html/rfc6455#section-7.1.7

        Parameters:
            code: close code
            reason: close reason

        Raises:
            ProtocolError: if the code isn't valid.
        """
        # 7.1.7. Fail the WebSocket Connection

        # Send a close frame when the state is OPEN (a close frame was already
        # sent if it's CLOSING), except when failing the connection because
        # of an error reading from or writing to the network.
        if self.state is OPEN:
            if code != CloseCode.ABNORMAL_CLOSURE:
                close = Close(code, reason)
                data = close.serialize()
                self.send_frame(Frame(OP_CLOSE, data))
                self.close_sent = close
                self.state = CLOSING

        # When failing the connection, a server closes the TCP connection
        # without waiting for the client to complete the handshake, while a
        # client waits for the server to close the TCP connection, possibly
        # after sending a close frame that the client will ignore.
        if self.side is SERVER and not self.eof_sent:
            self.send_eof()

        # 7.1.7. Fail the WebSocket Connection "An endpoint MUST NOT continue
        # to attempt to process data(including a responding Close frame) from
        # the remote endpoint after being instructed to _Fail the WebSocket
        # Connection_."
        self.parser = self.discard()
        next(self.parser)  # start coroutine

    # Public method for getting incoming events after receiving data.

    def events_received(self) -> List[Event]:
        """
        Fetch events generated from data received from the network.

        Call this method immediately after any of the ``receive_*()`` methods.

        Process resulting events, likely by passing them to the application.

        Returns:
            List[Event]: Events read from the connection.
        """
        events, self.events = self.events, []
        return events

    # Public method for getting outgoing data after receiving data or sending events.

    def data_to_send(self) -> List[bytes]:
        """
        Obtain data to send to the network.

        Call this method immediately after any of the ``receive_*()``,
        ``send_*()``, or :meth:`fail` methods.

        Write resulting data to the connection.

        The empty bytestring :data:`~websockets.protocol.SEND_EOF` signals
        the end of the data stream. When you receive it, half-close the TCP
        connection.

        Returns:
            List[bytes]: Data to write to the connection.

        """
        writes, self.writes = self.writes, []
        return writes

    def close_expected(self) -> bool:
        """
        Tell if the TCP connection is expected to close soon.

        Call this method immediately after any of the ``receive_*()``,
        ``send_close()``, or :meth:`fail` methods.

        If it returns :obj:`True`, schedule closing the TCP connection after a
        short timeout if the other side hasn't already closed it.

        Returns:
            bool: Whether the TCP connection is expected to close soon.

        """
        # We expect a TCP close if and only if we sent a close frame:
        # * Normal closure: once we send a close frame, we expect a TCP close:
        #   server waits for client to complete the TCP closing handshake;
        #   client waits for server to initiate the TCP closing handshake.
        # * Abnormal closure: we always send a close frame and the same logic
        #   applies, except on EOFError where we don't send a close frame
        #   because we already received the TCP close, so we don't expect it.
        # We already got a TCP Close if and only if the state is CLOSED.
        return self.state is CLOSING or self.handshake_exc is not None

    # Private methods for receiving data.

    def parse(self) -> Generator[None, None, None]:
        """
        Parse incoming data into frames.

        :meth:`receive_data` and :meth:`receive_eof` run this generator
        coroutine until it needs more data or reaches EOF.

        :meth:`parse` never raises an exception. Instead, it sets the
        :attr:`parser_exc` and yields control.

        """
        try:
            while True:
                if (yield from self.reader.at_eof()):
                    if self.debug:
                        self.logger.debug("< EOF")
                    # If the WebSocket connection is closed cleanly, with a
                    # closing handhshake, recv_frame() substitutes parse()
                    # with discard(). This branch is reached only when the
                    # connection isn't closed cleanly.
                    raise EOFError("unexpected end of stream")

                if self.max_size is None:
                    max_size = None
                elif self.cur_size is None:
                    max_size = self.max_size
                else:
                    max_size = self.max_size - self.cur_size

                # During a normal closure, execution ends here on the next
                # iteration of the loop after receiving a close frame. At
                # this point, recv_frame() replaced parse() by discard().
                frame = yield from Frame.parse(
                    self.reader.read_exact,
                    mask=self.side is SERVER,
                    max_size=max_size,
                    extensions=self.extensions,
                )

                if self.debug:
                    self.logger.debug("< %s", frame)

                self.recv_frame(frame)

        except ProtocolError as exc:
            self.fail(CloseCode.PROTOCOL_ERROR, str(exc))
            self.parser_exc = exc

        except EOFError as exc:
            self.fail(CloseCode.ABNORMAL_CLOSURE, str(exc))
            self.parser_exc = exc

        except UnicodeDecodeError as exc:
            self.fail(CloseCode.INVALID_DATA, f"{exc.reason} at position {exc.start}")
            self.parser_exc = exc

        except PayloadTooBig as exc:
            self.fail(CloseCode.MESSAGE_TOO_BIG, str(exc))
            self.parser_exc = exc

        except Exception as exc:
            self.logger.error("parser failed", exc_info=True)
            # Don't include exception details, which may be security-sensitive.
            self.fail(CloseCode.INTERNAL_ERROR)
            self.parser_exc = exc

        # During an abnormal closure, execution ends here after catching an
        # exception. At this point, fail() replaced parse() by discard().
        yield
        raise AssertionError("parse() shouldn't step after error")

    def discard(self) -> Generator[None, None, None]:
        """
        Discard incoming data.

        This coroutine replaces :meth:`parse`:

        - after receiving a close frame, during a normal closure (1.4);
        - after sending a close frame, during an abnormal closure (7.1.7).

        """
        # The server close the TCP connection in the same circumstances where
        # discard() replaces parse(). The client closes the connection later,
        # after the server closes the connection or a timeout elapses.
        # (The latter case cannot be handled in this Sans-I/O layer.)
        assert (self.side is SERVER) == (self.eof_sent)
        while not (yield from self.reader.at_eof()):
            self.reader.discard()
        if self.debug:
            self.logger.debug("< EOF")
        # A server closes the TCP connection immediately, while a client
        # waits for the server to close the TCP connection.
        if self.side is CLIENT:
            self.send_eof()
        self.state = CLOSED
        # If discard() completes normally, execution ends here.
        yield
        # Once the reader reaches EOF, its feed_data/eof() methods raise an
        # error, so our receive_data/eof() methods don't step the generator.
        raise AssertionError("discard() shouldn't step after EOF")

    def recv_frame(self, frame: Frame) -> None:
        """
        Process an incoming frame.

        """
        if frame.opcode is OP_TEXT or frame.opcode is OP_BINARY:
            if self.cur_size is not None:
                raise ProtocolError("expected a continuation frame")
            if frame.fin:
                self.cur_size = None
            else:
                self.cur_size = len(frame.data)

        elif frame.opcode is OP_CONT:
            if self.cur_size is None:
                raise ProtocolError("unexpected continuation frame")
            if frame.fin:
                self.cur_size = None
            else:
                self.cur_size += len(frame.data)

        elif frame.opcode is OP_PING:
            # 5.5.2. Ping: "Upon receipt of a Ping frame, an endpoint MUST
            # send a Pong frame in response"
            pong_frame = Frame(OP_PONG, frame.data)
            self.send_frame(pong_frame)

        elif frame.opcode is OP_PONG:
            # 5.5.3 Pong: "A response to an unsolicited Pong frame is not
            # expected."
            pass

        elif frame.opcode is OP_CLOSE:
            # 7.1.5.  The WebSocket Connection Close Code
            # 7.1.6.  The WebSocket Connection Close Reason
            self.close_rcvd = Close.parse(frame.data)
            if self.state is CLOSING:
                assert self.close_sent is not None
                self.close_rcvd_then_sent = False

            if self.cur_size is not None:
                raise ProtocolError("incomplete fragmented message")

            # 5.5.1 Close: "If an endpoint receives a Close frame and did
            # not previously send a Close frame, the endpoint MUST send a
            # Close frame in response. (When sending a Close frame in
            # response, the endpoint typically echos the status code it
            # received.)"

            if self.state is OPEN:
                # Echo the original data instead of re-serializing it with
                # Close.serialize() because that fails when the close frame
                # is empty and Close.parse() synthesizes a 1005 close code.
                # The rest is identical to send_close().
                self.send_frame(Frame(OP_CLOSE, frame.data))
                self.close_sent = self.close_rcvd
                self.close_rcvd_then_sent = True
                self.state = CLOSING

            # 7.1.2. Start the WebSocket Closing Handshake: "Once an
            # endpoint has both sent and received a Close control frame,
            # that endpoint SHOULD _Close the WebSocket Connection_"

            # A server closes the TCP connection immediately, while a client
            # waits for the server to close the TCP connection.
            if self.side is SERVER:
                self.send_eof()

            # 1.4. Closing Handshake: "after receiving a control frame
            # indicating the connection should be closed, a peer discards
            # any further data received."
            self.parser = self.discard()
            next(self.parser)  # start coroutine

        else:
            # This can't happen because Frame.parse() validates opcodes.
            raise AssertionError(f"unexpected opcode: {frame.opcode:02x}")

        self.events.append(frame)

    # Private methods for sending events.

    def send_frame(self, frame: Frame) -> None:
        if self.state is not OPEN:
            raise InvalidState(
                f"cannot write to a WebSocket in the {self.state.name} state"
            )

        if self.debug:
            self.logger.debug("> %s", frame)
        self.writes.append(
            frame.serialize(mask=self.side is CLIENT, extensions=self.extensions)
        )

    def send_eof(self) -> None:
        assert not self.eof_sent
        self.eof_sent = True
        if self.debug:
            self.logger.debug("> EOF")
        self.writes.append(SEND_EOF)