You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
707 lines
23 KiB
707 lines
23 KiB
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, |
|
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 1006 |
|
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(1005, "") |
|
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 != 1006: |
|
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(1002, str(exc)) |
|
self.parser_exc = exc |
|
|
|
except EOFError as exc: |
|
self.fail(1006, str(exc)) |
|
self.parser_exc = exc |
|
|
|
except UnicodeDecodeError as exc: |
|
self.fail(1007, f"{exc.reason} at position {exc.start}") |
|
self.parser_exc = exc |
|
|
|
except PayloadTooBig as exc: |
|
self.fail(1009, 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(1011) |
|
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)
|
|
|