123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341 |
- """
- Parse and serialize WebSocket frames.
- """
- import enum
- import io
- import secrets
- import struct
- from typing import Callable, Generator, NamedTuple, Optional, Sequence, Tuple
- from .exceptions import PayloadTooBig, ProtocolError
- from .typing import Data
- try:
- from .speedups import apply_mask
- except ImportError: # pragma: no cover
- from .utils import apply_mask
- __all__ = [
- "Opcode",
- "OP_CONT",
- "OP_TEXT",
- "OP_BINARY",
- "OP_CLOSE",
- "OP_PING",
- "OP_PONG",
- "DATA_OPCODES",
- "CTRL_OPCODES",
- "Frame",
- "prepare_data",
- "prepare_ctrl",
- "parse_close",
- "serialize_close",
- ]
- class Opcode(enum.IntEnum):
- CONT, TEXT, BINARY = 0x00, 0x01, 0x02
- CLOSE, PING, PONG = 0x08, 0x09, 0x0A
- OP_CONT = Opcode.CONT
- OP_TEXT = Opcode.TEXT
- OP_BINARY = Opcode.BINARY
- OP_CLOSE = Opcode.CLOSE
- OP_PING = Opcode.PING
- OP_PONG = Opcode.PONG
- DATA_OPCODES = OP_CONT, OP_TEXT, OP_BINARY
- CTRL_OPCODES = OP_CLOSE, OP_PING, OP_PONG
- # Close code that are allowed in a close frame.
- # Using a set optimizes `code in EXTERNAL_CLOSE_CODES`.
- EXTERNAL_CLOSE_CODES = {
- 1000,
- 1001,
- 1002,
- 1003,
- 1007,
- 1008,
- 1009,
- 1010,
- 1011,
- 1012,
- 1013,
- 1014,
- }
- # Consider converting to a dataclass when dropping support for Python < 3.7.
- class Frame(NamedTuple):
- """
- WebSocket frame.
- :param bool fin: FIN bit
- :param bool rsv1: RSV1 bit
- :param bool rsv2: RSV2 bit
- :param bool rsv3: RSV3 bit
- :param int opcode: opcode
- :param bytes data: payload data
- Only these fields are needed. The MASK bit, payload length and masking-key
- are handled on the fly by :func:`parse_frame` and :meth:`serialize_frame`.
- """
- fin: bool
- opcode: Opcode
- data: bytes
- rsv1: bool = False
- rsv2: bool = False
- rsv3: bool = False
- @classmethod
- def parse(
- cls,
- read_exact: Callable[[int], Generator[None, None, bytes]],
- *,
- mask: bool,
- max_size: Optional[int] = None,
- extensions: Optional[Sequence["extensions.Extension"]] = None,
- ) -> Generator[None, None, "Frame"]:
- """
- Read a WebSocket frame.
- :param read_exact: generator-based coroutine that reads the requested
- number of bytes or raises an exception if there isn't enough data
- :param mask: whether the frame should be masked i.e. whether the read
- happens on the server side
- :param max_size: maximum payload size in bytes
- :param extensions: list of classes with a ``decode()`` method that
- transforms the frame and return a new frame; extensions are applied
- in reverse order
- :raises ~websockets.exceptions.PayloadTooBig: if the frame exceeds
- ``max_size``
- :raises ~websockets.exceptions.ProtocolError: if the frame
- contains incorrect values
- """
- # Read the header.
- data = yield from read_exact(2)
- head1, head2 = struct.unpack("!BB", data)
- # While not Pythonic, this is marginally faster than calling bool().
- fin = True if head1 & 0b10000000 else False
- rsv1 = True if head1 & 0b01000000 else False
- rsv2 = True if head1 & 0b00100000 else False
- rsv3 = True if head1 & 0b00010000 else False
- try:
- opcode = Opcode(head1 & 0b00001111)
- except ValueError as exc:
- raise ProtocolError("invalid opcode") from exc
- if (True if head2 & 0b10000000 else False) != mask:
- raise ProtocolError("incorrect masking")
- length = head2 & 0b01111111
- if length == 126:
- data = yield from read_exact(2)
- (length,) = struct.unpack("!H", data)
- elif length == 127:
- data = yield from read_exact(8)
- (length,) = struct.unpack("!Q", data)
- if max_size is not None and length > max_size:
- raise PayloadTooBig(f"over size limit ({length} > {max_size} bytes)")
- if mask:
- mask_bytes = yield from read_exact(4)
- # Read the data.
- data = yield from read_exact(length)
- if mask:
- data = apply_mask(data, mask_bytes)
- frame = cls(fin, opcode, data, rsv1, rsv2, rsv3)
- if extensions is None:
- extensions = []
- for extension in reversed(extensions):
- frame = extension.decode(frame, max_size=max_size)
- frame.check()
- return frame
- def serialize(
- self,
- *,
- mask: bool,
- extensions: Optional[Sequence["extensions.Extension"]] = None,
- ) -> bytes:
- """
- Write a WebSocket frame.
- :param frame: frame to write
- :param mask: whether the frame should be masked i.e. whether the write
- happens on the client side
- :param extensions: list of classes with an ``encode()`` method that
- transform the frame and return a new frame; extensions are applied
- in order
- :raises ~websockets.exceptions.ProtocolError: if the frame
- contains incorrect values
- """
- self.check()
- if extensions is None:
- extensions = []
- for extension in extensions:
- self = extension.encode(self)
- output = io.BytesIO()
- # Prepare the header.
- head1 = (
- (0b10000000 if self.fin else 0)
- | (0b01000000 if self.rsv1 else 0)
- | (0b00100000 if self.rsv2 else 0)
- | (0b00010000 if self.rsv3 else 0)
- | self.opcode
- )
- head2 = 0b10000000 if mask else 0
- length = len(self.data)
- if length < 126:
- output.write(struct.pack("!BB", head1, head2 | length))
- elif length < 65536:
- output.write(struct.pack("!BBH", head1, head2 | 126, length))
- else:
- output.write(struct.pack("!BBQ", head1, head2 | 127, length))
- if mask:
- mask_bytes = secrets.token_bytes(4)
- output.write(mask_bytes)
- # Prepare the data.
- if mask:
- data = apply_mask(self.data, mask_bytes)
- else:
- data = self.data
- output.write(data)
- return output.getvalue()
- def check(self) -> None:
- """
- Check that reserved bits and opcode have acceptable values.
- :raises ~websockets.exceptions.ProtocolError: if a reserved
- bit or the opcode is invalid
- """
- if self.rsv1 or self.rsv2 or self.rsv3:
- raise ProtocolError("reserved bits must be 0")
- if self.opcode in CTRL_OPCODES:
- if len(self.data) > 125:
- raise ProtocolError("control frame too long")
- if not self.fin:
- raise ProtocolError("fragmented control frame")
- def prepare_data(data: Data) -> Tuple[int, bytes]:
- """
- Convert a string or byte-like object to an opcode and a bytes-like object.
- This function is designed for data frames.
- If ``data`` is a :class:`str`, return ``OP_TEXT`` and a :class:`bytes`
- object encoding ``data`` in UTF-8.
- If ``data`` is a bytes-like object, return ``OP_BINARY`` and a bytes-like
- object.
- :raises TypeError: if ``data`` doesn't have a supported type
- """
- if isinstance(data, str):
- return OP_TEXT, data.encode("utf-8")
- elif isinstance(data, (bytes, bytearray, memoryview)):
- return OP_BINARY, data
- else:
- raise TypeError("data must be bytes-like or str")
- def prepare_ctrl(data: Data) -> bytes:
- """
- Convert a string or byte-like object to bytes.
- This function is designed for ping and pong frames.
- If ``data`` is a :class:`str`, return a :class:`bytes` object encoding
- ``data`` in UTF-8.
- If ``data`` is a bytes-like object, return a :class:`bytes` object.
- :raises TypeError: if ``data`` doesn't have a supported type
- """
- if isinstance(data, str):
- return data.encode("utf-8")
- elif isinstance(data, (bytes, bytearray, memoryview)):
- return bytes(data)
- else:
- raise TypeError("data must be bytes-like or str")
- def parse_close(data: bytes) -> Tuple[int, str]:
- """
- Parse the payload from a close frame.
- Return ``(code, reason)``.
- :raises ~websockets.exceptions.ProtocolError: if data is ill-formed
- :raises UnicodeDecodeError: if the reason isn't valid UTF-8
- """
- length = len(data)
- if length >= 2:
- (code,) = struct.unpack("!H", data[:2])
- check_close(code)
- reason = data[2:].decode("utf-8")
- return code, reason
- elif length == 0:
- return 1005, ""
- else:
- assert length == 1
- raise ProtocolError("close frame too short")
- def serialize_close(code: int, reason: str) -> bytes:
- """
- Serialize the payload for a close frame.
- This is the reverse of :func:`parse_close`.
- """
- check_close(code)
- return struct.pack("!H", code) + reason.encode("utf-8")
- def check_close(code: int) -> None:
- """
- Check that the close code has an acceptable value for a close frame.
- :raises ~websockets.exceptions.ProtocolError: if the close code
- is invalid
- """
- if not (code in EXTERNAL_CLOSE_CODES or 3000 <= code < 5000):
- raise ProtocolError("invalid status code")
- # at the bottom to allow circular import, because Extension depends on Frame
- from . import extensions # isort:skip # noqa
|