framing.py 6.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224
  1. from __future__ import annotations
  2. import struct
  3. from collections.abc import Awaitable, Sequence
  4. from typing import Any, Callable, NamedTuple
  5. from .. import extensions, frames
  6. from ..exceptions import PayloadTooBig, ProtocolError
  7. from ..typing import BytesLike, DataLike
  8. try:
  9. from ..speedups import apply_mask
  10. except ImportError:
  11. from ..utils import apply_mask
  12. class Frame(NamedTuple):
  13. fin: bool
  14. opcode: frames.Opcode
  15. data: BytesLike
  16. rsv1: bool = False
  17. rsv2: bool = False
  18. rsv3: bool = False
  19. @property
  20. def new_frame(self) -> frames.Frame:
  21. return frames.Frame(
  22. self.opcode,
  23. self.data,
  24. self.fin,
  25. self.rsv1,
  26. self.rsv2,
  27. self.rsv3,
  28. )
  29. def __str__(self) -> str:
  30. return str(self.new_frame)
  31. def check(self) -> None:
  32. return self.new_frame.check()
  33. @classmethod
  34. async def read(
  35. cls,
  36. reader: Callable[[int], Awaitable[bytes]],
  37. *,
  38. mask: bool,
  39. max_size: int | None = None,
  40. extensions: Sequence[extensions.Extension] | None = None,
  41. ) -> Frame:
  42. """
  43. Read a WebSocket frame.
  44. Args:
  45. reader: Coroutine that reads exactly the requested number of
  46. bytes, unless the end of file is reached.
  47. mask: Whether the frame should be masked i.e. whether the read
  48. happens on the server side.
  49. max_size: Maximum payload size in bytes.
  50. extensions: List of extensions, applied in reverse order.
  51. Raises:
  52. PayloadTooBig: If the frame exceeds ``max_size``.
  53. ProtocolError: If the frame contains incorrect values.
  54. """
  55. # Read the header.
  56. data = await reader(2)
  57. head1, head2 = struct.unpack("!BB", data)
  58. # While not Pythonic, this is marginally faster than calling bool().
  59. fin = True if head1 & 0b10000000 else False
  60. rsv1 = True if head1 & 0b01000000 else False
  61. rsv2 = True if head1 & 0b00100000 else False
  62. rsv3 = True if head1 & 0b00010000 else False
  63. try:
  64. opcode = frames.Opcode(head1 & 0b00001111)
  65. except ValueError as exc:
  66. raise ProtocolError("invalid opcode") from exc
  67. if (True if head2 & 0b10000000 else False) != mask:
  68. raise ProtocolError("incorrect masking")
  69. length = head2 & 0b01111111
  70. if length == 126:
  71. data = await reader(2)
  72. (length,) = struct.unpack("!H", data)
  73. elif length == 127:
  74. data = await reader(8)
  75. (length,) = struct.unpack("!Q", data)
  76. if max_size is not None and length > max_size:
  77. raise PayloadTooBig(length, max_size)
  78. if mask:
  79. mask_bits = await reader(4)
  80. # Read the data.
  81. data = await reader(length)
  82. if mask:
  83. data = apply_mask(data, mask_bits)
  84. new_frame = frames.Frame(opcode, data, fin, rsv1, rsv2, rsv3)
  85. if extensions is None:
  86. extensions = []
  87. for extension in reversed(extensions):
  88. new_frame = extension.decode(new_frame, max_size=max_size)
  89. new_frame.check()
  90. return cls(
  91. new_frame.fin,
  92. new_frame.opcode,
  93. new_frame.data,
  94. new_frame.rsv1,
  95. new_frame.rsv2,
  96. new_frame.rsv3,
  97. )
  98. def write(
  99. self,
  100. write: Callable[[bytes], Any],
  101. *,
  102. mask: bool,
  103. extensions: Sequence[extensions.Extension] | None = None,
  104. ) -> None:
  105. """
  106. Write a WebSocket frame.
  107. Args:
  108. frame: Frame to write.
  109. write: Function that writes bytes.
  110. mask: Whether the frame should be masked i.e. whether the write
  111. happens on the client side.
  112. extensions: List of extensions, applied in order.
  113. Raises:
  114. ProtocolError: If the frame contains incorrect values.
  115. """
  116. # The frame is written in a single call to write in order to prevent
  117. # TCP fragmentation. See #68 for details. This also makes it safe to
  118. # send frames concurrently from multiple coroutines.
  119. write(self.new_frame.serialize(mask=mask, extensions=extensions))
  120. def prepare_data(data: DataLike) -> tuple[int, BytesLike]:
  121. """
  122. Convert a string or byte-like object to an opcode and a bytes-like object.
  123. This function is designed for data frames.
  124. If ``data`` is a :class:`str`, return ``OP_TEXT`` and a :class:`bytes`
  125. object encoding ``data`` in UTF-8.
  126. If ``data`` is a bytes-like object, return ``OP_BINARY`` and a bytes-like
  127. object.
  128. Raises:
  129. TypeError: If ``data`` doesn't have a supported type.
  130. """
  131. if isinstance(data, str):
  132. return frames.Opcode.TEXT, data.encode()
  133. elif isinstance(data, BytesLike):
  134. return frames.Opcode.BINARY, data
  135. else:
  136. raise TypeError("data must be str or bytes-like")
  137. def prepare_ctrl(data: DataLike) -> bytes:
  138. """
  139. Convert a string or byte-like object to bytes.
  140. This function is designed for ping and pong frames.
  141. If ``data`` is a :class:`str`, return a :class:`bytes` object encoding
  142. ``data`` in UTF-8.
  143. If ``data`` is a bytes-like object, return a :class:`bytes` object.
  144. Raises:
  145. TypeError: If ``data`` doesn't have a supported type.
  146. """
  147. if isinstance(data, str):
  148. return data.encode()
  149. elif isinstance(data, BytesLike):
  150. return bytes(data)
  151. else:
  152. raise TypeError("data must be str or bytes-like")
  153. # Backwards compatibility with previously documented public APIs
  154. encode_data = prepare_ctrl
  155. # Backwards compatibility with previously documented public APIs
  156. from ..frames import Close # noqa: E402 F401, I001
  157. def parse_close(data: bytes) -> tuple[int, str]:
  158. """
  159. Parse the payload from a close frame.
  160. Returns:
  161. Close code and reason.
  162. Raises:
  163. ProtocolError: If data is ill-formed.
  164. UnicodeDecodeError: If the reason isn't valid UTF-8.
  165. """
  166. close = Close.parse(data)
  167. return close.code, close.reason
  168. def serialize_close(code: int, reason: str) -> bytes:
  169. """
  170. Serialize the payload for a close frame.
  171. """
  172. return Close(code, reason).serialize()