frames.py 13 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431
  1. from __future__ import annotations
  2. import dataclasses
  3. import enum
  4. import io
  5. import os
  6. import secrets
  7. import struct
  8. from collections.abc import Generator, Sequence
  9. from typing import Callable
  10. from .exceptions import PayloadTooBig, ProtocolError
  11. from .typing import BytesLike
  12. try:
  13. from .speedups import apply_mask
  14. except ImportError:
  15. from .utils import apply_mask
  16. __all__ = [
  17. "Opcode",
  18. "OP_CONT",
  19. "OP_TEXT",
  20. "OP_BINARY",
  21. "OP_CLOSE",
  22. "OP_PING",
  23. "OP_PONG",
  24. "DATA_OPCODES",
  25. "CTRL_OPCODES",
  26. "CloseCode",
  27. "Frame",
  28. "Close",
  29. ]
  30. class Opcode(enum.IntEnum):
  31. """Opcode values for WebSocket frames."""
  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. class CloseCode(enum.IntEnum):
  43. """Close code values for WebSocket close frames."""
  44. NORMAL_CLOSURE = 1000
  45. GOING_AWAY = 1001
  46. PROTOCOL_ERROR = 1002
  47. UNSUPPORTED_DATA = 1003
  48. # 1004 is reserved
  49. NO_STATUS_RCVD = 1005
  50. ABNORMAL_CLOSURE = 1006
  51. INVALID_DATA = 1007
  52. POLICY_VIOLATION = 1008
  53. MESSAGE_TOO_BIG = 1009
  54. MANDATORY_EXTENSION = 1010
  55. INTERNAL_ERROR = 1011
  56. SERVICE_RESTART = 1012
  57. TRY_AGAIN_LATER = 1013
  58. BAD_GATEWAY = 1014
  59. TLS_HANDSHAKE = 1015
  60. # See https://www.iana.org/assignments/websocket/websocket.xhtml
  61. CLOSE_CODE_EXPLANATIONS: dict[int, str] = {
  62. CloseCode.NORMAL_CLOSURE: "OK",
  63. CloseCode.GOING_AWAY: "going away",
  64. CloseCode.PROTOCOL_ERROR: "protocol error",
  65. CloseCode.UNSUPPORTED_DATA: "unsupported data",
  66. CloseCode.NO_STATUS_RCVD: "no status received [internal]",
  67. CloseCode.ABNORMAL_CLOSURE: "abnormal closure [internal]",
  68. CloseCode.INVALID_DATA: "invalid frame payload data",
  69. CloseCode.POLICY_VIOLATION: "policy violation",
  70. CloseCode.MESSAGE_TOO_BIG: "message too big",
  71. CloseCode.MANDATORY_EXTENSION: "mandatory extension",
  72. CloseCode.INTERNAL_ERROR: "internal error",
  73. CloseCode.SERVICE_RESTART: "service restart",
  74. CloseCode.TRY_AGAIN_LATER: "try again later",
  75. CloseCode.BAD_GATEWAY: "bad gateway",
  76. CloseCode.TLS_HANDSHAKE: "TLS handshake failure [internal]",
  77. }
  78. # Close code that are allowed in a close frame.
  79. # Using a set optimizes `code in EXTERNAL_CLOSE_CODES`.
  80. EXTERNAL_CLOSE_CODES = {
  81. CloseCode.NORMAL_CLOSURE,
  82. CloseCode.GOING_AWAY,
  83. CloseCode.PROTOCOL_ERROR,
  84. CloseCode.UNSUPPORTED_DATA,
  85. CloseCode.INVALID_DATA,
  86. CloseCode.POLICY_VIOLATION,
  87. CloseCode.MESSAGE_TOO_BIG,
  88. CloseCode.MANDATORY_EXTENSION,
  89. CloseCode.INTERNAL_ERROR,
  90. CloseCode.SERVICE_RESTART,
  91. CloseCode.TRY_AGAIN_LATER,
  92. CloseCode.BAD_GATEWAY,
  93. }
  94. OK_CLOSE_CODES = {
  95. CloseCode.NORMAL_CLOSURE,
  96. CloseCode.GOING_AWAY,
  97. CloseCode.NO_STATUS_RCVD,
  98. }
  99. @dataclasses.dataclass
  100. class Frame:
  101. """
  102. WebSocket frame.
  103. Attributes:
  104. opcode: Opcode.
  105. data: Payload data.
  106. fin: FIN bit.
  107. rsv1: RSV1 bit.
  108. rsv2: RSV2 bit.
  109. rsv3: RSV3 bit.
  110. Only these fields are needed. The MASK bit, payload length and masking-key
  111. are handled on the fly when parsing and serializing frames.
  112. """
  113. opcode: Opcode
  114. data: BytesLike
  115. fin: bool = True
  116. rsv1: bool = False
  117. rsv2: bool = False
  118. rsv3: bool = False
  119. # Configure if you want to see more in logs. Should be a multiple of 3.
  120. MAX_LOG_SIZE = int(os.environ.get("WEBSOCKETS_MAX_LOG_SIZE", "75"))
  121. def __str__(self) -> str:
  122. """
  123. Return a human-readable representation of a frame.
  124. """
  125. coding = None
  126. length = f"{len(self.data)} byte{'' if len(self.data) == 1 else 's'}"
  127. non_final = "" if self.fin else "continued"
  128. if self.opcode is OP_TEXT:
  129. # Decoding only the beginning and the end is needlessly hard.
  130. # Decode the entire payload then elide later if necessary.
  131. data = repr(bytes(self.data).decode())
  132. elif self.opcode is OP_BINARY:
  133. # We'll show at most the first 16 bytes and the last 8 bytes.
  134. # Encode just what we need, plus two dummy bytes to elide later.
  135. binary = self.data
  136. if len(binary) > self.MAX_LOG_SIZE // 3:
  137. cut = (self.MAX_LOG_SIZE // 3 - 1) // 3 # by default cut = 8
  138. binary = b"".join([binary[: 2 * cut], b"\x00\x00", binary[-cut:]])
  139. data = " ".join(f"{byte:02x}" for byte in binary)
  140. elif self.opcode is OP_CLOSE:
  141. data = str(Close.parse(self.data))
  142. elif self.data:
  143. # We don't know if a Continuation frame contains text or binary.
  144. # Ping and Pong frames could contain UTF-8.
  145. # Attempt to decode as UTF-8 and display it as text; fallback to
  146. # binary. If self.data is a memoryview, it has no decode() method,
  147. # which raises AttributeError.
  148. try:
  149. data = repr(bytes(self.data).decode())
  150. coding = "text"
  151. except (UnicodeDecodeError, AttributeError):
  152. binary = self.data
  153. if len(binary) > self.MAX_LOG_SIZE // 3:
  154. cut = (self.MAX_LOG_SIZE // 3 - 1) // 3 # by default cut = 8
  155. binary = b"".join([binary[: 2 * cut], b"\x00\x00", binary[-cut:]])
  156. data = " ".join(f"{byte:02x}" for byte in binary)
  157. coding = "binary"
  158. else:
  159. data = "''"
  160. if len(data) > self.MAX_LOG_SIZE:
  161. cut = self.MAX_LOG_SIZE // 3 - 1 # by default cut = 24
  162. data = data[: 2 * cut] + "..." + data[-cut:]
  163. metadata = ", ".join(filter(None, [coding, length, non_final]))
  164. return f"{self.opcode.name} {data} [{metadata}]"
  165. @classmethod
  166. def parse(
  167. cls,
  168. read_exact: Callable[[int], Generator[None, None, bytes | bytearray]],
  169. *,
  170. mask: bool,
  171. max_size: int | None = None,
  172. extensions: Sequence[extensions.Extension] | None = None,
  173. ) -> Generator[None, None, Frame]:
  174. """
  175. Parse a WebSocket frame.
  176. This is a generator-based coroutine.
  177. Args:
  178. read_exact: Generator-based coroutine that reads the requested
  179. bytes or raises an exception if there isn't enough data.
  180. mask: Whether the frame should be masked i.e. whether the read
  181. happens on the server side.
  182. max_size: Maximum payload size in bytes.
  183. extensions: List of extensions, applied in reverse order.
  184. Raises:
  185. EOFError: If the connection is closed without a full WebSocket frame.
  186. PayloadTooBig: If the frame's payload size exceeds ``max_size``.
  187. ProtocolError: If the frame contains incorrect values.
  188. """
  189. # Read the header.
  190. data = yield from read_exact(2)
  191. head1, head2 = struct.unpack("!BB", data)
  192. # While not Pythonic, this is marginally faster than calling bool().
  193. fin = True if head1 & 0b10000000 else False
  194. rsv1 = True if head1 & 0b01000000 else False
  195. rsv2 = True if head1 & 0b00100000 else False
  196. rsv3 = True if head1 & 0b00010000 else False
  197. try:
  198. opcode = Opcode(head1 & 0b00001111)
  199. except ValueError as exc:
  200. raise ProtocolError("invalid opcode") from exc
  201. if (True if head2 & 0b10000000 else False) != mask:
  202. raise ProtocolError("incorrect masking")
  203. length = head2 & 0b01111111
  204. if length == 126:
  205. data = yield from read_exact(2)
  206. (length,) = struct.unpack("!H", data)
  207. elif length == 127:
  208. data = yield from read_exact(8)
  209. (length,) = struct.unpack("!Q", data)
  210. if max_size is not None and length > max_size:
  211. raise PayloadTooBig(length, max_size)
  212. if mask:
  213. mask_bytes = yield from read_exact(4)
  214. # Read the data.
  215. data = yield from read_exact(length)
  216. if mask:
  217. data = apply_mask(data, mask_bytes)
  218. frame = cls(opcode, data, fin, rsv1, rsv2, rsv3)
  219. if extensions is None:
  220. extensions = []
  221. for extension in reversed(extensions):
  222. frame = extension.decode(frame, max_size=max_size)
  223. frame.check()
  224. return frame
  225. def serialize(
  226. self,
  227. *,
  228. mask: bool,
  229. extensions: Sequence[extensions.Extension] | None = None,
  230. ) -> bytes:
  231. """
  232. Serialize a WebSocket frame.
  233. Args:
  234. mask: Whether the frame should be masked i.e. whether the write
  235. happens on the client side.
  236. extensions: List of extensions, applied in order.
  237. Raises:
  238. ProtocolError: If the frame contains incorrect values.
  239. """
  240. self.check()
  241. if extensions is None:
  242. extensions = []
  243. for extension in extensions:
  244. self = extension.encode(self)
  245. output = io.BytesIO()
  246. # Prepare the header.
  247. head1 = (
  248. (0b10000000 if self.fin else 0)
  249. | (0b01000000 if self.rsv1 else 0)
  250. | (0b00100000 if self.rsv2 else 0)
  251. | (0b00010000 if self.rsv3 else 0)
  252. | self.opcode
  253. )
  254. head2 = 0b10000000 if mask else 0
  255. length = len(self.data)
  256. if length < 126:
  257. output.write(struct.pack("!BB", head1, head2 | length))
  258. elif length < 65536:
  259. output.write(struct.pack("!BBH", head1, head2 | 126, length))
  260. else:
  261. output.write(struct.pack("!BBQ", head1, head2 | 127, length))
  262. if mask:
  263. mask_bytes = secrets.token_bytes(4)
  264. output.write(mask_bytes)
  265. # Prepare the data.
  266. data: BytesLike
  267. if mask:
  268. data = apply_mask(self.data, mask_bytes)
  269. else:
  270. data = self.data
  271. output.write(data)
  272. return output.getvalue()
  273. def check(self) -> None:
  274. """
  275. Check that reserved bits and opcode have acceptable values.
  276. Raises:
  277. ProtocolError: If a reserved bit or the opcode is invalid.
  278. """
  279. if self.rsv1 or self.rsv2 or self.rsv3:
  280. raise ProtocolError("reserved bits must be 0")
  281. if self.opcode in CTRL_OPCODES:
  282. if len(self.data) > 125:
  283. raise ProtocolError("control frame too long")
  284. if not self.fin:
  285. raise ProtocolError("fragmented control frame")
  286. @dataclasses.dataclass
  287. class Close:
  288. """
  289. Code and reason for WebSocket close frames.
  290. Attributes:
  291. code: Close code.
  292. reason: Close reason.
  293. """
  294. code: CloseCode | int
  295. reason: str
  296. def __str__(self) -> str:
  297. """
  298. Return a human-readable representation of a close code and reason.
  299. """
  300. if 3000 <= self.code < 4000:
  301. explanation = "registered"
  302. elif 4000 <= self.code < 5000:
  303. explanation = "private use"
  304. else:
  305. explanation = CLOSE_CODE_EXPLANATIONS.get(self.code, "unknown")
  306. result = f"{self.code} ({explanation})"
  307. if self.reason:
  308. result = f"{result} {self.reason}"
  309. return result
  310. @classmethod
  311. def parse(cls, data: BytesLike) -> Close:
  312. """
  313. Parse the payload of a close frame.
  314. Args:
  315. data: Payload of the close frame.
  316. Raises:
  317. ProtocolError: If data is ill-formed.
  318. UnicodeDecodeError: If the reason isn't valid UTF-8.
  319. """
  320. if isinstance(data, memoryview):
  321. raise AssertionError("only compressed outgoing frames use memoryview")
  322. if len(data) >= 2:
  323. (code,) = struct.unpack("!H", data[:2])
  324. reason = data[2:].decode()
  325. close = cls(code, reason)
  326. close.check()
  327. return close
  328. elif len(data) == 0:
  329. return cls(CloseCode.NO_STATUS_RCVD, "")
  330. else:
  331. raise ProtocolError("close frame too short")
  332. def serialize(self) -> bytes:
  333. """
  334. Serialize the payload of a close frame.
  335. """
  336. self.check()
  337. return struct.pack("!H", self.code) + self.reason.encode()
  338. def check(self) -> None:
  339. """
  340. Check that the close code has a valid value for a close frame.
  341. Raises:
  342. ProtocolError: If the close code is invalid.
  343. """
  344. if not (self.code in EXTERNAL_CLOSE_CODES or 3000 <= self.code < 5000):
  345. raise ProtocolError("invalid status code")
  346. # At the bottom to break import cycles created by type annotations.
  347. from . import extensions # noqa: E402