python3-websockets 3.0-1 / usr / lib / python3 / dist-packages / websockets / protocol.py

"""
The :mod:`websockets.protocol` module handles WebSocket control and data
frames as specified in `sections 4 to 8 of RFC 6455`_.

.. _sections 4 to 8 of RFC 6455: http://tools.ietf.org/html/rfc6455#section-4

"""

__all__ = ['WebSocketCommonProtocol']

import asyncio
import asyncio.queues
import codecs
import collections
import logging
import random
import struct

from .compatibility import asyncio_ensure_future
from .exceptions import (ConnectionClosed, InvalidState, PayloadTooBig,
                         WebSocketProtocolError)
from .framing import *
from .handshake import *


logger = logging.getLogger(__name__)


# A WebSocket connection goes through the following four states, in order:

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

# In order to ensure consistency, the code always checks the current value of
# WebSocketCommonProtocol.state before assigning a new value and never yields
# between the check and the assignment.


class WebSocketCommonProtocol(asyncio.StreamReaderProtocol):
    """
    This class implements common parts of the WebSocket protocol.

    It assumes that the WebSocket connection is established. The handshake is
    managed in subclasses such as
    :class:`~websockets.server.WebSocketServerProtocol` and
    :class:`~websockets.client.WebSocketClientProtocol`.

    It runs a task that stores incoming data frames in a queue and deals with
    control frames automatically. It sends outgoing data frames and performs
    the closing handshake.

    The ``host``, ``port`` and ``secure`` parameters are simply stored as
    attributes for handlers that need them.

    The ``timeout`` parameter defines the maximum wait time in seconds for
    completing the closing handshake and, only on the client side, for
    terminating the TCP connection. :meth:`close()` will complete in at most
    this time on the server side and twice this time on the client side.

    The ``max_size`` parameter enforces the maximum size for incoming messages
    in bytes. The default value is 1MB. ``None`` disables the limit. If a
    message larger than the maximum size is received, :meth:`recv()` will
    raise :exc:`~websockets.exceptions.ConnectionClosed` and the connection
    will be closed with status code 1009.

    Once the handshake is complete, request and response HTTP headers are
    available:

    * as a MIME :class:`~email.message.Message` in the :attr:`request_headers`
      and :attr:`response_headers` attributes
    * as an iterable of (name, value) pairs in the :attr:`raw_request_headers`
      and :attr:`raw_response_headers` attributes

    If a subprotocol was negotiated, it's available in the :attr:`subprotocol`
    attribute.

    Once the connection is closed, the status code is available in the
    :attr:`close_code` attribute and the reason in :attr:`close_reason`.

    """
    # There are only two differences between the client-side and the server-
    # side behavior: masking the payload and closing the underlying TCP
    # connection. This class implements the server-side behavior by default.
    # To get the client-side behavior, set is_client = True.

    is_client = False
    state = OPEN

    def __init__(self, *,
                 host=None, port=None, secure=None,
                 timeout=10, max_size=2 ** 20, loop=None,
                 legacy_recv=False):
        self.host = host
        self.port = port
        self.secure = secure

        self.timeout = timeout
        self.max_size = max_size
        # Store a reference to loop to avoid relying on self._loop, a private
        # attribute of StreamReaderProtocol, inherited from FlowControlMixin.
        self.loop = loop

        self.legacy_recv = legacy_recv

        stream_reader = asyncio.StreamReader(loop=loop)
        super().__init__(stream_reader, self.client_connected, loop)

        self.reader = None
        self.writer = None

        self.request_headers = None
        self.raw_request_headers = None
        self.response_headers = None
        self.raw_response_headers = None

        self.subprotocol = None

        # Code and reason must be set when the closing handshake completes.
        self.close_code = None
        self.close_reason = ''

        # Futures tracking steps in the connection's lifecycle.
        # Set to True when the opening handshake has completed properly.
        self.opening_handshake = asyncio.Future(loop=loop)
        # Set to True when the closing handshake has completed properly and to
        # False when the connection terminates abnormally.
        self.closing_handshake = asyncio.Future(loop=loop)
        # Set to None when the connection state becomes CLOSED.
        self.connection_closed = asyncio.Future(loop=loop)

        # Queue of received messages.
        self.messages = asyncio.queues.Queue(loop=loop)

        # Mapping of ping IDs to waiters, in chronological order.
        self.pings = collections.OrderedDict()

        # Task managing the connection, initalized in self.client_connected.
        self.worker = None

        # In a subclass implementing the opening handshake, the state will be
        # CONNECTING at this point.
        if self.state == OPEN:
            self.opening_handshake.set_result(True)

    # Public API

    @property
    def local_address(self):
        """
        Local address of the connection.

        This is a ``(host, port)`` tuple or ``None`` if the connection hasn't
        been established yet.

        """
        if self.writer is None:
            return None
        return self.writer.get_extra_info('sockname')

    @property
    def remote_address(self):
        """
        Remote address of the connection.

        This is a ``(host, port)`` tuple or ``None`` if the connection hasn't
        been established yet.

        """
        if self.writer is None:
            return None
        return self.writer.get_extra_info('peername')

    @property
    def open(self):
        """
        This property is ``True`` when the connection is usable.

        It may be used to detect disconnections but this is discouraged per
        the EAFP_ principle. When ``open`` is ``False``, using the connection
        raises a :exc:`~websockets.exceptions.ConnectionClosed` exception.

        .. _EAFP: https://docs.python.org/3/glossary.html#term-eafp

        """
        return self.state == OPEN

    @property
    def state_name(self):
        """
        Current connection state, as a string.

        Possible states are defined in the WebSocket specification:
        CONNECTING, OPEN, CLOSING, or CLOSED.

        To check if the connection is open, use :attr:`open` instead.

        """
        return ['CONNECTING', 'OPEN', 'CLOSING', 'CLOSED'][self.state]

    @asyncio.coroutine
    def close(self, code=1000, reason=''):
        """
        This coroutine performs the closing handshake.

        It waits for the other end to complete the handshake. It doesn't do
        anything once the connection is closed. Thus it's idemptotent.

        It's safe to wrap this coroutine in :func:`~asyncio.ensure_future`
        since errors during connection termination aren't particularly useful.

        ``code`` must be an :class:`int` and ``reason`` a :class:`str`.

        """
        if self.state == OPEN:
            # 7.1.2. Start the WebSocket Closing Handshake
            # 7.1.3. The WebSocket Closing Handshake is Started
            frame_data = serialize_close(code, reason)
            yield from self.write_frame(OP_CLOSE, frame_data)

        # If the connection doesn't terminate within the timeout, break out of
        # the worker loop.
        try:
            yield from asyncio.wait_for(
                self.worker, self.timeout, loop=self.loop)
        except asyncio.TimeoutError:
            self.worker.cancel()

        # The worker should terminate quickly once it has been cancelled.
        yield from self.worker

    @asyncio.coroutine
    def recv(self):
        """
        This coroutine receives the next message.

        It returns a :class:`str` for a text frame and :class:`bytes` for a
        binary frame.

        When the end of the message stream is reached, :meth:`recv` raises
        :exc:`~websockets.exceptions.ConnectionClosed`. This can happen after
        a normal connection closure, a protocol error or a network failure.

        .. versionchanged:: 3.0

            :meth:`recv` used to return ``None`` instead. Refer to the
            changelog for details.

        """
        # Don't yield from self.ensure_open() here because messages could be
        # available in the queue even if the connection is closed.

        # Return any available message
        try:
            return self.messages.get_nowait()
        except asyncio.queues.QueueEmpty:
            pass

        # Don't yield from self.ensure_open() here because messages could be
        # received before the closing frame even if the connection is closing.

        # Wait for a message until the connection is closed
        next_message = asyncio_ensure_future(
            self.messages.get(), loop=self.loop)
        try:
            done, pending = yield from asyncio.wait(
                [next_message, self.worker],
                loop=self.loop, return_when=asyncio.FIRST_COMPLETED)
        except asyncio.CancelledError:
            # Handle the Task.cancel()
            next_message.cancel()
            raise

        # Now there's no need to yield from self.ensure_open(). Either a
        # message was received or the connection was closed.

        if next_message in done:
            return next_message.result()
        else:
            next_message.cancel()
            if not self.legacy_recv:
                raise ConnectionClosed(self.close_code, self.close_reason)

    @asyncio.coroutine
    def send(self, data):
        """
        This coroutine sends a message.

        It sends :class:`str` as a text frame and :class:`bytes` as a binary
        frame. It raises a :exc:`TypeError` for other inputs.

        """
        yield from self.ensure_open()

        if isinstance(data, str):
            opcode = 1
            data = data.encode('utf-8')
        elif isinstance(data, bytes):
            opcode = 2
        else:
            raise TypeError("data must be bytes or str")

        yield from self.write_frame(opcode, data)

    @asyncio.coroutine
    def ping(self, data=None):
        """
        This coroutine sends a ping.

        It returns a :class:`~asyncio.Future` which will be completed when the
        corresponding pong is received and which you may ignore if you don't
        want to wait.

        A ping may serve as a keepalive or as a check that the remote endpoint
        received all messages up to this point, with ``yield from ws.ping()``.

        By default, the ping contains four random bytes. The content may be
        overridden with the optional ``data`` argument which must be of type
        :class:`str` (which will be encoded to UTF-8) or :class:`bytes`.

        """
        yield from self.ensure_open()

        if data is not None:
            data = self.encode_data(data)

        # Protect against duplicates if a payload is explicitly set.
        if data in self.pings:
            raise ValueError("Already waiting for a pong with the same data")

        # Generate a unique random payload otherwise.
        while data is None or data in self.pings:
            data = struct.pack('!I', random.getrandbits(32))

        self.pings[data] = asyncio.Future(loop=self.loop)
        yield from self.write_frame(OP_PING, data)
        return self.pings[data]

    @asyncio.coroutine
    def pong(self, data=b''):
        """
        This coroutine sends a pong.

        An unsolicited pong may serve as a unidirectional heartbeat.

        The content may be overridden with the optional ``data`` argument
        which must be of type :class:`str` (which will be encoded to UTF-8) or
        :class:`bytes`.

        """
        yield from self.ensure_open()

        data = self.encode_data(data)

        yield from self.write_frame(OP_PONG, data)

    # Private methods - no guarantees.

    def encode_data(self, data):
        # Expect str or bytes, return bytes.
        if isinstance(data, str):
            return data.encode('utf-8')
        elif isinstance(data, bytes):
            return data
        else:
            raise TypeError("data must be bytes or str")

    @asyncio.coroutine
    def ensure_open(self):
        # Raise a suitable exception if the connection isn't open.
        # Handle cases from the most common to the least common.

        if self.state == OPEN:
            return

        if self.state == CLOSED:
            raise ConnectionClosed(self.close_code, self.close_reason)

        # If the closing handshake is in progress, let it complete to get the
        # proper close status and code. As an safety measure, the timeout is
        # longer than the worst case (2 * self.timeout) but not unlimited.
        if self.state == CLOSING:
            yield from asyncio.wait_for(
                self.worker, 3 * self.timeout, loop=self.loop)
            raise ConnectionClosed(self.close_code, self.close_reason)

        # Control may only reach this point in buggy third-party subclasses.
        assert self.state == CONNECTING
        raise InvalidState("WebSocket connection isn't established yet.")

    @asyncio.coroutine
    def run(self):
        # This coroutine guarantees that the connection is closed at exit.
        yield from self.opening_handshake
        while not self.closing_handshake.done():
            try:
                msg = yield from self.read_message()
                if msg is None:
                    break
                self.messages.put_nowait(msg)
            except asyncio.CancelledError:
                break
            except WebSocketProtocolError:
                yield from self.fail_connection(1002)
            except asyncio.IncompleteReadError:
                yield from self.fail_connection(1006)
            except UnicodeDecodeError:
                yield from self.fail_connection(1007)
            except PayloadTooBig:
                yield from self.fail_connection(1009)
            except Exception:
                yield from self.fail_connection(1011)
                raise
        yield from self.close_connection()

    @asyncio.coroutine
    def read_message(self):
        # Reassemble fragmented messages.
        frame = yield from self.read_data_frame(max_size=self.max_size)
        if frame is None:
            return
        if frame.opcode == OP_TEXT:
            text = True
        elif frame.opcode == OP_BINARY:
            text = False
        else:   # frame.opcode == OP_CONT
            raise WebSocketProtocolError("Unexpected opcode")

        # Shortcut for the common case - no fragmentation
        if frame.fin:
            return frame.data.decode('utf-8') if text else frame.data

        # 5.4. Fragmentation
        chunks = []
        max_size = self.max_size
        if text:
            decoder = codecs.getincrementaldecoder('utf-8')(errors='strict')
            if max_size is None:
                def append(frame):
                    nonlocal chunks
                    chunks.append(decoder.decode(frame.data, frame.fin))
            else:
                def append(frame):
                    nonlocal chunks, max_size
                    chunks.append(decoder.decode(frame.data, frame.fin))
                    max_size -= len(frame.data)
        else:
            if max_size is None:
                def append(frame):
                    nonlocal chunks
                    chunks.append(frame.data)
            else:
                def append(frame):
                    nonlocal chunks, max_size
                    chunks.append(frame.data)
                    max_size -= len(frame.data)
        append(frame)

        while not frame.fin:
            frame = yield from self.read_data_frame(max_size=max_size)
            if frame is None:
                raise WebSocketProtocolError("Incomplete fragmented message")
            if frame.opcode != OP_CONT:
                raise WebSocketProtocolError("Unexpected opcode")
            append(frame)

        return ('' if text else b'').join(chunks)

    @asyncio.coroutine
    def read_data_frame(self, max_size):
        # Deal with control frames automatically and return next data frame.
        # 6.2. Receiving Data
        while True:
            frame = yield from self.read_frame(max_size)

            # 5.5. Control Frames
            if frame.opcode == OP_CLOSE:
                # Make sure the close frame is valid before echoing it.
                code, reason = parse_close(frame.data)
                if self.state == OPEN:
                    # 7.1.3. The WebSocket Closing Handshake is Started
                    yield from self.write_frame(OP_CLOSE, frame.data)
                self.close_code, self.close_reason = code, reason
                self.closing_handshake.set_result(True)
                return

            elif frame.opcode == OP_PING:
                # Answer pings.
                yield from self.pong(frame.data)

            elif frame.opcode == OP_PONG:
                # Do not acknowledge pings on unsolicited pongs.
                if frame.data in self.pings:
                    # Acknowledge all pings up to the one matching this pong.
                    ping_id = None
                    while ping_id != frame.data:
                        ping_id, waiter = self.pings.popitem(0)
                        if not waiter.cancelled():
                            waiter.set_result(None)

            # 5.6. Data Frames
            else:
                return frame

    @asyncio.coroutine
    def read_frame(self, max_size):
        is_masked = not self.is_client
        frame = yield from read_frame(
            self.reader.readexactly, is_masked, max_size=max_size)
        side = 'client' if self.is_client else 'server'
        logger.debug("%s << %s", side, frame)
        return frame

    @asyncio.coroutine
    def write_frame(self, opcode, data=b''):
        # Defensive assertion for protocol compliance.
        if self.state != OPEN:                              # pragma: no cover
            raise InvalidState("Cannot write to a WebSocket "
                               "in the {} state".format(self.state_name))

        # Make sure no other frame will be sent after a close frame. Do this
        # before yielding control to avoid sending more than one close frame.
        if opcode == OP_CLOSE:
            self.state = CLOSING
        frame = Frame(True, opcode, data)
        side = 'client' if self.is_client else 'server'
        logger.debug("%s >> %s", side, frame)
        is_masked = self.is_client
        write_frame(frame, self.writer.write, is_masked)

        # Backport of the combined logic of:
        # https://github.com/python/asyncio/pull/280
        # https://github.com/python/asyncio/pull/291
        # Remove when dropping support for Python < 3.6.
        transport = self.writer._transport
        if transport is not None:                           # pragma: no cover
            # PR 291 added the is_closing method to transports shortly after
            # PR 280 fixed the bug we're trying to work around in this block.
            if not hasattr(transport, 'is_closing'):
                # This emulates what is_closing would return if it existed.
                try:
                    is_closing = transport._closing
                except AttributeError:
                    is_closing = transport._closed
                if is_closing:
                    yield

        try:
            # Handle flow control automatically.
            yield from self.writer.drain()
        except ConnectionError:
            # Terminate the connection if the socket died.
            yield from self.fail_connection(1006)
            # And raise an exception, since the frame couldn't be sent.
            raise ConnectionClosed(self.close_code, self.close_reason)

    @asyncio.coroutine
    def close_connection(self, force=False):
        # 7.1.1. Close the WebSocket Connection
        if self.state == CLOSED:
            return

        # Defensive assertion for protocol compliance.
        if self.state != CLOSING and not force:             # pragma: no cover
            raise InvalidState("Cannot close a WebSocket connection "
                               "in the {} state".format(self.state_name))

        if self.is_client and not force:
            try:
                yield from asyncio.wait_for(
                    self.connection_closed, self.timeout, loop=self.loop)
            except (asyncio.CancelledError, asyncio.TimeoutError):
                pass

            if self.state == CLOSED:
                return

        # Attempt to terminate the TCP connection properly.
        # If the socket is already closed, this may crash.
        try:
            if self.writer.can_write_eof():
                self.writer.write_eof()
        except Exception:                                   # pragma: no cover
            pass

        self.writer.close()

        try:
            yield from asyncio.wait_for(
                self.connection_closed, self.timeout, loop=self.loop)
        except (asyncio.CancelledError, asyncio.TimeoutError):
            pass

    @asyncio.coroutine
    def fail_connection(self, code=1011, reason=''):
        # 7.1.7. Fail the WebSocket Connection
        logger.info("Failing the WebSocket connection: %d %s", code, reason)
        if self.state == OPEN:
            if code == 1006:
                # Don't send a close frame if the connection is broken. Set
                # the state to CLOSING to allow close_connection to proceed.
                self.state = CLOSING
            else:
                frame_data = serialize_close(code, reason)
                yield from self.write_frame(OP_CLOSE, frame_data)
        if not self.closing_handshake.done():
            self.close_code, self.close_reason = code, reason
            self.closing_handshake.set_result(False)
        yield from self.close_connection()

    # asyncio StreamReaderProtocol methods

    def client_connected(self, reader, writer):
        self.reader = reader
        self.writer = writer
        # Start the task that handles incoming messages.
        self.worker = asyncio_ensure_future(self.run(), loop=self.loop)

    def eof_received(self):
        super().eof_received()
        # Since Python 3.5, StreamReaderProtocol.eof_received() returns True
        # to leave the transport open (http://bugs.python.org/issue24539).
        # This is inappropriate for websockets for at least three reasons.
        # 1. The use case is to read data until EOF with self.reader.read(-1).
        #    Since websockets is a TLV protocol, this never happens.
        # 2. It doesn't work on SSL connections. A falsy value must be
        #    returned to have the same behavior on SSL and plain connections.
        # 3. The websockets protocol has its own closing handshake. Endpoints
        #    close the TCP connection after sending a Close frame.
        # As a consequence we revert to the previous, more useful behavior.
        return

    def connection_lost(self, exc):
        # 7.1.4. The WebSocket Connection is Closed
        self.state = CLOSED
        if not self.closing_handshake.done():
            self.close_code, self.close_reason = 1006, ''
            self.closing_handshake.set_result(False)
        if not self.connection_closed.done():
            self.connection_closed.set_result(None)
        # Close the transport in case close_connection() wasn't executed.
        if self.writer is not None:
            self.writer.close()
        super().connection_lost(exc)