frames.py 9.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341
  1. """
  2. Parse and serialize WebSocket frames.
  3. """
  4. import enum
  5. import io
  6. import secrets
  7. import struct
  8. from typing import Callable, Generator, NamedTuple, Optional, Sequence, Tuple
  9. from .exceptions import PayloadTooBig, ProtocolError
  10. from .typing import Data
  11. try:
  12. from .speedups import apply_mask
  13. except ImportError: # pragma: no cover
  14. from .utils import apply_mask
  15. __all__ = [
  16. "Opcode",
  17. "OP_CONT",
  18. "OP_TEXT",
  19. "OP_BINARY",
  20. "OP_CLOSE",
  21. "OP_PING",
  22. "OP_PONG",
  23. "DATA_OPCODES",
  24. "CTRL_OPCODES",
  25. "Frame",
  26. "prepare_data",
  27. "prepare_ctrl",
  28. "parse_close",
  29. "serialize_close",
  30. ]
  31. class Opcode(enum.IntEnum):
  32. CONT, TEXT, BINARY = 0x00, 0x01, 0x02
  33. CLOSE, PING, PONG = 0x08, 0x09, 0x0A
  34. OP_CONT = Opcode.CONT
  35. OP_TEXT = Opcode.TEXT
  36. OP_BINARY = Opcode.BINARY
  37. OP_CLOSE = Opcode.CLOSE
  38. OP_PING = Opcode.PING
  39. OP_PONG = Opcode.PONG
  40. DATA_OPCODES = OP_CONT, OP_TEXT, OP_BINARY
  41. CTRL_OPCODES = OP_CLOSE, OP_PING, OP_PONG
  42. # Close code that are allowed in a close frame.
  43. # Using a set optimizes `code in EXTERNAL_CLOSE_CODES`.
  44. EXTERNAL_CLOSE_CODES = {
  45. 1000,
  46. 1001,
  47. 1002,
  48. 1003,
  49. 1007,
  50. 1008,
  51. 1009,
  52. 1010,
  53. 1011,
  54. 1012,
  55. 1013,
  56. 1014,
  57. }
  58. # Consider converting to a dataclass when dropping support for Python < 3.7.
  59. class Frame(NamedTuple):
  60. """
  61. WebSocket frame.
  62. :param bool fin: FIN bit
  63. :param bool rsv1: RSV1 bit
  64. :param bool rsv2: RSV2 bit
  65. :param bool rsv3: RSV3 bit
  66. :param int opcode: opcode
  67. :param bytes data: payload data
  68. Only these fields are needed. The MASK bit, payload length and masking-key
  69. are handled on the fly by :func:`parse_frame` and :meth:`serialize_frame`.
  70. """
  71. fin: bool
  72. opcode: Opcode
  73. data: bytes
  74. rsv1: bool = False
  75. rsv2: bool = False
  76. rsv3: bool = False
  77. @classmethod
  78. def parse(
  79. cls,
  80. read_exact: Callable[[int], Generator[None, None, bytes]],
  81. *,
  82. mask: bool,
  83. max_size: Optional[int] = None,
  84. extensions: Optional[Sequence["extensions.Extension"]] = None,
  85. ) -> Generator[None, None, "Frame"]:
  86. """
  87. Read a WebSocket frame.
  88. :param read_exact: generator-based coroutine that reads the requested
  89. number of bytes or raises an exception if there isn't enough data
  90. :param mask: whether the frame should be masked i.e. whether the read
  91. happens on the server side
  92. :param max_size: maximum payload size in bytes
  93. :param extensions: list of classes with a ``decode()`` method that
  94. transforms the frame and return a new frame; extensions are applied
  95. in reverse order
  96. :raises ~websockets.exceptions.PayloadTooBig: if the frame exceeds
  97. ``max_size``
  98. :raises ~websockets.exceptions.ProtocolError: if the frame
  99. contains incorrect values
  100. """
  101. # Read the header.
  102. data = yield from read_exact(2)
  103. head1, head2 = struct.unpack("!BB", data)
  104. # While not Pythonic, this is marginally faster than calling bool().
  105. fin = True if head1 & 0b10000000 else False
  106. rsv1 = True if head1 & 0b01000000 else False
  107. rsv2 = True if head1 & 0b00100000 else False
  108. rsv3 = True if head1 & 0b00010000 else False
  109. try:
  110. opcode = Opcode(head1 & 0b00001111)
  111. except ValueError as exc:
  112. raise ProtocolError("invalid opcode") from exc
  113. if (True if head2 & 0b10000000 else False) != mask:
  114. raise ProtocolError("incorrect masking")
  115. length = head2 & 0b01111111
  116. if length == 126:
  117. data = yield from read_exact(2)
  118. (length,) = struct.unpack("!H", data)
  119. elif length == 127:
  120. data = yield from read_exact(8)
  121. (length,) = struct.unpack("!Q", data)
  122. if max_size is not None and length > max_size:
  123. raise PayloadTooBig(f"over size limit ({length} > {max_size} bytes)")
  124. if mask:
  125. mask_bytes = yield from read_exact(4)
  126. # Read the data.
  127. data = yield from read_exact(length)
  128. if mask:
  129. data = apply_mask(data, mask_bytes)
  130. frame = cls(fin, opcode, data, rsv1, rsv2, rsv3)
  131. if extensions is None:
  132. extensions = []
  133. for extension in reversed(extensions):
  134. frame = extension.decode(frame, max_size=max_size)
  135. frame.check()
  136. return frame
  137. def serialize(
  138. self,
  139. *,
  140. mask: bool,
  141. extensions: Optional[Sequence["extensions.Extension"]] = None,
  142. ) -> bytes:
  143. """
  144. Write a WebSocket frame.
  145. :param frame: frame to write
  146. :param mask: whether the frame should be masked i.e. whether the write
  147. happens on the client side
  148. :param extensions: list of classes with an ``encode()`` method that
  149. transform the frame and return a new frame; extensions are applied
  150. in order
  151. :raises ~websockets.exceptions.ProtocolError: if the frame
  152. contains incorrect values
  153. """
  154. self.check()
  155. if extensions is None:
  156. extensions = []
  157. for extension in extensions:
  158. self = extension.encode(self)
  159. output = io.BytesIO()
  160. # Prepare the header.
  161. head1 = (
  162. (0b10000000 if self.fin else 0)
  163. | (0b01000000 if self.rsv1 else 0)
  164. | (0b00100000 if self.rsv2 else 0)
  165. | (0b00010000 if self.rsv3 else 0)
  166. | self.opcode
  167. )
  168. head2 = 0b10000000 if mask else 0
  169. length = len(self.data)
  170. if length < 126:
  171. output.write(struct.pack("!BB", head1, head2 | length))
  172. elif length < 65536:
  173. output.write(struct.pack("!BBH", head1, head2 | 126, length))
  174. else:
  175. output.write(struct.pack("!BBQ", head1, head2 | 127, length))
  176. if mask:
  177. mask_bytes = secrets.token_bytes(4)
  178. output.write(mask_bytes)
  179. # Prepare the data.
  180. if mask:
  181. data = apply_mask(self.data, mask_bytes)
  182. else:
  183. data = self.data
  184. output.write(data)
  185. return output.getvalue()
  186. def check(self) -> None:
  187. """
  188. Check that reserved bits and opcode have acceptable values.
  189. :raises ~websockets.exceptions.ProtocolError: if a reserved
  190. bit or the opcode is invalid
  191. """
  192. if self.rsv1 or self.rsv2 or self.rsv3:
  193. raise ProtocolError("reserved bits must be 0")
  194. if self.opcode in CTRL_OPCODES:
  195. if len(self.data) > 125:
  196. raise ProtocolError("control frame too long")
  197. if not self.fin:
  198. raise ProtocolError("fragmented control frame")
  199. def prepare_data(data: Data) -> Tuple[int, bytes]:
  200. """
  201. Convert a string or byte-like object to an opcode and a bytes-like object.
  202. This function is designed for data frames.
  203. If ``data`` is a :class:`str`, return ``OP_TEXT`` and a :class:`bytes`
  204. object encoding ``data`` in UTF-8.
  205. If ``data`` is a bytes-like object, return ``OP_BINARY`` and a bytes-like
  206. object.
  207. :raises TypeError: if ``data`` doesn't have a supported type
  208. """
  209. if isinstance(data, str):
  210. return OP_TEXT, data.encode("utf-8")
  211. elif isinstance(data, (bytes, bytearray, memoryview)):
  212. return OP_BINARY, data
  213. else:
  214. raise TypeError("data must be bytes-like or str")
  215. def prepare_ctrl(data: Data) -> bytes:
  216. """
  217. Convert a string or byte-like object to bytes.
  218. This function is designed for ping and pong frames.
  219. If ``data`` is a :class:`str`, return a :class:`bytes` object encoding
  220. ``data`` in UTF-8.
  221. If ``data`` is a bytes-like object, return a :class:`bytes` object.
  222. :raises TypeError: if ``data`` doesn't have a supported type
  223. """
  224. if isinstance(data, str):
  225. return data.encode("utf-8")
  226. elif isinstance(data, (bytes, bytearray, memoryview)):
  227. return bytes(data)
  228. else:
  229. raise TypeError("data must be bytes-like or str")
  230. def parse_close(data: bytes) -> Tuple[int, str]:
  231. """
  232. Parse the payload from a close frame.
  233. Return ``(code, reason)``.
  234. :raises ~websockets.exceptions.ProtocolError: if data is ill-formed
  235. :raises UnicodeDecodeError: if the reason isn't valid UTF-8
  236. """
  237. length = len(data)
  238. if length >= 2:
  239. (code,) = struct.unpack("!H", data[:2])
  240. check_close(code)
  241. reason = data[2:].decode("utf-8")
  242. return code, reason
  243. elif length == 0:
  244. return 1005, ""
  245. else:
  246. assert length == 1
  247. raise ProtocolError("close frame too short")
  248. def serialize_close(code: int, reason: str) -> bytes:
  249. """
  250. Serialize the payload for a close frame.
  251. This is the reverse of :func:`parse_close`.
  252. """
  253. check_close(code)
  254. return struct.pack("!H", code) + reason.encode("utf-8")
  255. def check_close(code: int) -> None:
  256. """
  257. Check that the close code has an acceptable value for a close frame.
  258. :raises ~websockets.exceptions.ProtocolError: if the close code
  259. is invalid
  260. """
  261. if not (code in EXTERNAL_CLOSE_CODES or 3000 <= code < 5000):
  262. raise ProtocolError("invalid status code")
  263. # at the bottom to allow circular import, because Extension depends on Frame
  264. from . import extensions # isort:skip # noqa