permessage_deflate.py 23 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635
  1. """
  2. :mod:`websockets.extensions.permessage_deflate` implements the Compression
  3. Extensions for WebSocket as specified in :rfc:`7692`.
  4. """
  5. import zlib
  6. from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
  7. from ..exceptions import (
  8. DuplicateParameter,
  9. InvalidParameterName,
  10. InvalidParameterValue,
  11. NegotiationError,
  12. PayloadTooBig,
  13. )
  14. from ..frames import CTRL_OPCODES, OP_CONT, Frame
  15. from ..typing import ExtensionName, ExtensionParameter
  16. from .base import ClientExtensionFactory, Extension, ServerExtensionFactory
  17. __all__ = [
  18. "PerMessageDeflate",
  19. "ClientPerMessageDeflateFactory",
  20. "enable_client_permessage_deflate",
  21. "ServerPerMessageDeflateFactory",
  22. "enable_server_permessage_deflate",
  23. ]
  24. _EMPTY_UNCOMPRESSED_BLOCK = b"\x00\x00\xff\xff"
  25. _MAX_WINDOW_BITS_VALUES = [str(bits) for bits in range(8, 16)]
  26. class PerMessageDeflate(Extension):
  27. """
  28. Per-Message Deflate extension.
  29. """
  30. name = ExtensionName("permessage-deflate")
  31. def __init__(
  32. self,
  33. remote_no_context_takeover: bool,
  34. local_no_context_takeover: bool,
  35. remote_max_window_bits: int,
  36. local_max_window_bits: int,
  37. compress_settings: Optional[Dict[Any, Any]] = None,
  38. ) -> None:
  39. """
  40. Configure the Per-Message Deflate extension.
  41. """
  42. if compress_settings is None:
  43. compress_settings = {}
  44. assert remote_no_context_takeover in [False, True]
  45. assert local_no_context_takeover in [False, True]
  46. assert 8 <= remote_max_window_bits <= 15
  47. assert 8 <= local_max_window_bits <= 15
  48. assert "wbits" not in compress_settings
  49. self.remote_no_context_takeover = remote_no_context_takeover
  50. self.local_no_context_takeover = local_no_context_takeover
  51. self.remote_max_window_bits = remote_max_window_bits
  52. self.local_max_window_bits = local_max_window_bits
  53. self.compress_settings = compress_settings
  54. if not self.remote_no_context_takeover:
  55. self.decoder = zlib.decompressobj(wbits=-self.remote_max_window_bits)
  56. if not self.local_no_context_takeover:
  57. self.encoder = zlib.compressobj(
  58. wbits=-self.local_max_window_bits, **self.compress_settings
  59. )
  60. # To handle continuation frames properly, we must keep track of
  61. # whether that initial frame was encoded.
  62. self.decode_cont_data = False
  63. # There's no need for self.encode_cont_data because we always encode
  64. # outgoing frames, so it would always be True.
  65. def __repr__(self) -> str:
  66. return (
  67. f"PerMessageDeflate("
  68. f"remote_no_context_takeover={self.remote_no_context_takeover}, "
  69. f"local_no_context_takeover={self.local_no_context_takeover}, "
  70. f"remote_max_window_bits={self.remote_max_window_bits}, "
  71. f"local_max_window_bits={self.local_max_window_bits})"
  72. )
  73. def decode(self, frame: Frame, *, max_size: Optional[int] = None) -> Frame:
  74. """
  75. Decode an incoming frame.
  76. """
  77. # Skip control frames.
  78. if frame.opcode in CTRL_OPCODES:
  79. return frame
  80. # Handle continuation data frames:
  81. # - skip if the message isn't encoded
  82. # - reset "decode continuation data" flag if it's a final frame
  83. if frame.opcode == OP_CONT:
  84. if not self.decode_cont_data:
  85. return frame
  86. if frame.fin:
  87. self.decode_cont_data = False
  88. # Handle text and binary data frames:
  89. # - skip if the message isn't encoded
  90. # - unset the rsv1 flag on the first frame of a compressed message
  91. # - set "decode continuation data" flag if it's a non-final frame
  92. else:
  93. if not frame.rsv1:
  94. return frame
  95. frame = frame._replace(rsv1=False)
  96. if not frame.fin:
  97. self.decode_cont_data = True
  98. # Re-initialize per-message decoder.
  99. if self.remote_no_context_takeover:
  100. self.decoder = zlib.decompressobj(wbits=-self.remote_max_window_bits)
  101. # Uncompress data. Protect against zip bombs by preventing zlib from
  102. # decompressing more than max_length bytes (except when the limit is
  103. # disabled with max_size = None).
  104. data = frame.data
  105. if frame.fin:
  106. data += _EMPTY_UNCOMPRESSED_BLOCK
  107. max_length = 0 if max_size is None else max_size
  108. data = self.decoder.decompress(data, max_length)
  109. if self.decoder.unconsumed_tail:
  110. raise PayloadTooBig(f"over size limit (? > {max_size} bytes)")
  111. # Allow garbage collection of the decoder if it won't be reused.
  112. if frame.fin and self.remote_no_context_takeover:
  113. del self.decoder
  114. return frame._replace(data=data)
  115. def encode(self, frame: Frame) -> Frame:
  116. """
  117. Encode an outgoing frame.
  118. """
  119. # Skip control frames.
  120. if frame.opcode in CTRL_OPCODES:
  121. return frame
  122. # Since we always encode messages, there's no "encode continuation
  123. # data" flag similar to "decode continuation data" at this time.
  124. if frame.opcode != OP_CONT:
  125. # Set the rsv1 flag on the first frame of a compressed message.
  126. frame = frame._replace(rsv1=True)
  127. # Re-initialize per-message decoder.
  128. if self.local_no_context_takeover:
  129. self.encoder = zlib.compressobj(
  130. wbits=-self.local_max_window_bits, **self.compress_settings
  131. )
  132. # Compress data.
  133. data = self.encoder.compress(frame.data) + self.encoder.flush(zlib.Z_SYNC_FLUSH)
  134. if frame.fin and data.endswith(_EMPTY_UNCOMPRESSED_BLOCK):
  135. data = data[:-4]
  136. # Allow garbage collection of the encoder if it won't be reused.
  137. if frame.fin and self.local_no_context_takeover:
  138. del self.encoder
  139. return frame._replace(data=data)
  140. def _build_parameters(
  141. server_no_context_takeover: bool,
  142. client_no_context_takeover: bool,
  143. server_max_window_bits: Optional[int],
  144. client_max_window_bits: Optional[Union[int, bool]],
  145. ) -> List[ExtensionParameter]:
  146. """
  147. Build a list of ``(name, value)`` pairs for some compression parameters.
  148. """
  149. params: List[ExtensionParameter] = []
  150. if server_no_context_takeover:
  151. params.append(("server_no_context_takeover", None))
  152. if client_no_context_takeover:
  153. params.append(("client_no_context_takeover", None))
  154. if server_max_window_bits:
  155. params.append(("server_max_window_bits", str(server_max_window_bits)))
  156. if client_max_window_bits is True: # only in handshake requests
  157. params.append(("client_max_window_bits", None))
  158. elif client_max_window_bits:
  159. params.append(("client_max_window_bits", str(client_max_window_bits)))
  160. return params
  161. def _extract_parameters(
  162. params: Sequence[ExtensionParameter], *, is_server: bool
  163. ) -> Tuple[bool, bool, Optional[int], Optional[Union[int, bool]]]:
  164. """
  165. Extract compression parameters from a list of ``(name, value)`` pairs.
  166. If ``is_server`` is ``True``, ``client_max_window_bits`` may be provided
  167. without a value. This is only allow in handshake requests.
  168. """
  169. server_no_context_takeover: bool = False
  170. client_no_context_takeover: bool = False
  171. server_max_window_bits: Optional[int] = None
  172. client_max_window_bits: Optional[Union[int, bool]] = None
  173. for name, value in params:
  174. if name == "server_no_context_takeover":
  175. if server_no_context_takeover:
  176. raise DuplicateParameter(name)
  177. if value is None:
  178. server_no_context_takeover = True
  179. else:
  180. raise InvalidParameterValue(name, value)
  181. elif name == "client_no_context_takeover":
  182. if client_no_context_takeover:
  183. raise DuplicateParameter(name)
  184. if value is None:
  185. client_no_context_takeover = True
  186. else:
  187. raise InvalidParameterValue(name, value)
  188. elif name == "server_max_window_bits":
  189. if server_max_window_bits is not None:
  190. raise DuplicateParameter(name)
  191. if value in _MAX_WINDOW_BITS_VALUES:
  192. server_max_window_bits = int(value)
  193. else:
  194. raise InvalidParameterValue(name, value)
  195. elif name == "client_max_window_bits":
  196. if client_max_window_bits is not None:
  197. raise DuplicateParameter(name)
  198. if is_server and value is None: # only in handshake requests
  199. client_max_window_bits = True
  200. elif value in _MAX_WINDOW_BITS_VALUES:
  201. client_max_window_bits = int(value)
  202. else:
  203. raise InvalidParameterValue(name, value)
  204. else:
  205. raise InvalidParameterName(name)
  206. return (
  207. server_no_context_takeover,
  208. client_no_context_takeover,
  209. server_max_window_bits,
  210. client_max_window_bits,
  211. )
  212. class ClientPerMessageDeflateFactory(ClientExtensionFactory):
  213. """
  214. Client-side extension factory for the Per-Message Deflate extension.
  215. Parameters behave as described in `section 7.1 of RFC 7692`_. Set them to
  216. ``True`` to include them in the negotiation offer without a value or to an
  217. integer value to include them with this value.
  218. .. _section 7.1 of RFC 7692: https://tools.ietf.org/html/rfc7692#section-7.1
  219. :param server_no_context_takeover: defaults to ``False``
  220. :param client_no_context_takeover: defaults to ``False``
  221. :param server_max_window_bits: optional, defaults to ``None``
  222. :param client_max_window_bits: optional, defaults to ``None``
  223. :param compress_settings: optional, keyword arguments for
  224. :func:`zlib.compressobj`, excluding ``wbits``
  225. """
  226. name = ExtensionName("permessage-deflate")
  227. def __init__(
  228. self,
  229. server_no_context_takeover: bool = False,
  230. client_no_context_takeover: bool = False,
  231. server_max_window_bits: Optional[int] = None,
  232. client_max_window_bits: Optional[Union[int, bool]] = None,
  233. compress_settings: Optional[Dict[str, Any]] = None,
  234. ) -> None:
  235. """
  236. Configure the Per-Message Deflate extension factory.
  237. """
  238. if not (server_max_window_bits is None or 8 <= server_max_window_bits <= 15):
  239. raise ValueError("server_max_window_bits must be between 8 and 15")
  240. if not (
  241. client_max_window_bits is None
  242. or client_max_window_bits is True
  243. or 8 <= client_max_window_bits <= 15
  244. ):
  245. raise ValueError("client_max_window_bits must be between 8 and 15")
  246. if compress_settings is not None and "wbits" in compress_settings:
  247. raise ValueError(
  248. "compress_settings must not include wbits, "
  249. "set client_max_window_bits instead"
  250. )
  251. self.server_no_context_takeover = server_no_context_takeover
  252. self.client_no_context_takeover = client_no_context_takeover
  253. self.server_max_window_bits = server_max_window_bits
  254. self.client_max_window_bits = client_max_window_bits
  255. self.compress_settings = compress_settings
  256. def get_request_params(self) -> List[ExtensionParameter]:
  257. """
  258. Build request parameters.
  259. """
  260. return _build_parameters(
  261. self.server_no_context_takeover,
  262. self.client_no_context_takeover,
  263. self.server_max_window_bits,
  264. self.client_max_window_bits,
  265. )
  266. def process_response_params(
  267. self,
  268. params: Sequence[ExtensionParameter],
  269. accepted_extensions: Sequence["Extension"],
  270. ) -> PerMessageDeflate:
  271. """
  272. Process response parameters.
  273. Return an extension instance.
  274. """
  275. if any(other.name == self.name for other in accepted_extensions):
  276. raise NegotiationError(f"received duplicate {self.name}")
  277. # Request parameters are available in instance variables.
  278. # Load response parameters in local variables.
  279. (
  280. server_no_context_takeover,
  281. client_no_context_takeover,
  282. server_max_window_bits,
  283. client_max_window_bits,
  284. ) = _extract_parameters(params, is_server=False)
  285. # After comparing the request and the response, the final
  286. # configuration must be available in the local variables.
  287. # server_no_context_takeover
  288. #
  289. # Req. Resp. Result
  290. # ------ ------ --------------------------------------------------
  291. # False False False
  292. # False True True
  293. # True False Error!
  294. # True True True
  295. if self.server_no_context_takeover:
  296. if not server_no_context_takeover:
  297. raise NegotiationError("expected server_no_context_takeover")
  298. # client_no_context_takeover
  299. #
  300. # Req. Resp. Result
  301. # ------ ------ --------------------------------------------------
  302. # False False False
  303. # False True True
  304. # True False True - must change value
  305. # True True True
  306. if self.client_no_context_takeover:
  307. if not client_no_context_takeover:
  308. client_no_context_takeover = True
  309. # server_max_window_bits
  310. # Req. Resp. Result
  311. # ------ ------ --------------------------------------------------
  312. # None None None
  313. # None 8≤M≤15 M
  314. # 8≤N≤15 None Error!
  315. # 8≤N≤15 8≤M≤N M
  316. # 8≤N≤15 N<M≤15 Error!
  317. if self.server_max_window_bits is None:
  318. pass
  319. else:
  320. if server_max_window_bits is None:
  321. raise NegotiationError("expected server_max_window_bits")
  322. elif server_max_window_bits > self.server_max_window_bits:
  323. raise NegotiationError("unsupported server_max_window_bits")
  324. # client_max_window_bits
  325. # Req. Resp. Result
  326. # ------ ------ --------------------------------------------------
  327. # None None None
  328. # None 8≤M≤15 Error!
  329. # True None None
  330. # True 8≤M≤15 M
  331. # 8≤N≤15 None N - must change value
  332. # 8≤N≤15 8≤M≤N M
  333. # 8≤N≤15 N<M≤15 Error!
  334. if self.client_max_window_bits is None:
  335. if client_max_window_bits is not None:
  336. raise NegotiationError("unexpected client_max_window_bits")
  337. elif self.client_max_window_bits is True:
  338. pass
  339. else:
  340. if client_max_window_bits is None:
  341. client_max_window_bits = self.client_max_window_bits
  342. elif client_max_window_bits > self.client_max_window_bits:
  343. raise NegotiationError("unsupported client_max_window_bits")
  344. return PerMessageDeflate(
  345. server_no_context_takeover, # remote_no_context_takeover
  346. client_no_context_takeover, # local_no_context_takeover
  347. server_max_window_bits or 15, # remote_max_window_bits
  348. client_max_window_bits or 15, # local_max_window_bits
  349. self.compress_settings,
  350. )
  351. def enable_client_permessage_deflate(
  352. extensions: Optional[Sequence[ClientExtensionFactory]],
  353. ) -> Sequence[ClientExtensionFactory]:
  354. """
  355. Enable Per-Message Deflate with default settings in client extensions.
  356. If the extension is already present, perhaps with non-default settings,
  357. the configuration isn't changed.
  358. """
  359. if extensions is None:
  360. extensions = []
  361. if not any(
  362. extension_factory.name == ClientPerMessageDeflateFactory.name
  363. for extension_factory in extensions
  364. ):
  365. extensions = list(extensions) + [
  366. ClientPerMessageDeflateFactory(client_max_window_bits=True)
  367. ]
  368. return extensions
  369. class ServerPerMessageDeflateFactory(ServerExtensionFactory):
  370. """
  371. Server-side extension factory for the Per-Message Deflate extension.
  372. Parameters behave as described in `section 7.1 of RFC 7692`_. Set them to
  373. ``True`` to include them in the negotiation offer without a value or to an
  374. integer value to include them with this value.
  375. .. _section 7.1 of RFC 7692: https://tools.ietf.org/html/rfc7692#section-7.1
  376. :param server_no_context_takeover: defaults to ``False``
  377. :param client_no_context_takeover: defaults to ``False``
  378. :param server_max_window_bits: optional, defaults to ``None``
  379. :param client_max_window_bits: optional, defaults to ``None``
  380. :param compress_settings: optional, keyword arguments for
  381. :func:`zlib.compressobj`, excluding ``wbits``
  382. """
  383. name = ExtensionName("permessage-deflate")
  384. def __init__(
  385. self,
  386. server_no_context_takeover: bool = False,
  387. client_no_context_takeover: bool = False,
  388. server_max_window_bits: Optional[int] = None,
  389. client_max_window_bits: Optional[int] = None,
  390. compress_settings: Optional[Dict[str, Any]] = None,
  391. ) -> None:
  392. """
  393. Configure the Per-Message Deflate extension factory.
  394. """
  395. if not (server_max_window_bits is None or 8 <= server_max_window_bits <= 15):
  396. raise ValueError("server_max_window_bits must be between 8 and 15")
  397. if not (client_max_window_bits is None or 8 <= client_max_window_bits <= 15):
  398. raise ValueError("client_max_window_bits must be between 8 and 15")
  399. if compress_settings is not None and "wbits" in compress_settings:
  400. raise ValueError(
  401. "compress_settings must not include wbits, "
  402. "set server_max_window_bits instead"
  403. )
  404. self.server_no_context_takeover = server_no_context_takeover
  405. self.client_no_context_takeover = client_no_context_takeover
  406. self.server_max_window_bits = server_max_window_bits
  407. self.client_max_window_bits = client_max_window_bits
  408. self.compress_settings = compress_settings
  409. def process_request_params(
  410. self,
  411. params: Sequence[ExtensionParameter],
  412. accepted_extensions: Sequence["Extension"],
  413. ) -> Tuple[List[ExtensionParameter], PerMessageDeflate]:
  414. """
  415. Process request parameters.
  416. Return response params and an extension instance.
  417. """
  418. if any(other.name == self.name for other in accepted_extensions):
  419. raise NegotiationError(f"skipped duplicate {self.name}")
  420. # Load request parameters in local variables.
  421. (
  422. server_no_context_takeover,
  423. client_no_context_takeover,
  424. server_max_window_bits,
  425. client_max_window_bits,
  426. ) = _extract_parameters(params, is_server=True)
  427. # Configuration parameters are available in instance variables.
  428. # After comparing the request and the configuration, the response must
  429. # be available in the local variables.
  430. # server_no_context_takeover
  431. #
  432. # Config Req. Resp.
  433. # ------ ------ --------------------------------------------------
  434. # False False False
  435. # False True True
  436. # True False True - must change value to True
  437. # True True True
  438. if self.server_no_context_takeover:
  439. if not server_no_context_takeover:
  440. server_no_context_takeover = True
  441. # client_no_context_takeover
  442. #
  443. # Config Req. Resp.
  444. # ------ ------ --------------------------------------------------
  445. # False False False
  446. # False True True (or False)
  447. # True False True - must change value to True
  448. # True True True (or False)
  449. if self.client_no_context_takeover:
  450. if not client_no_context_takeover:
  451. client_no_context_takeover = True
  452. # server_max_window_bits
  453. # Config Req. Resp.
  454. # ------ ------ --------------------------------------------------
  455. # None None None
  456. # None 8≤M≤15 M
  457. # 8≤N≤15 None N - must change value
  458. # 8≤N≤15 8≤M≤N M
  459. # 8≤N≤15 N<M≤15 N - must change value
  460. if self.server_max_window_bits is None:
  461. pass
  462. else:
  463. if server_max_window_bits is None:
  464. server_max_window_bits = self.server_max_window_bits
  465. elif server_max_window_bits > self.server_max_window_bits:
  466. server_max_window_bits = self.server_max_window_bits
  467. # client_max_window_bits
  468. # Config Req. Resp.
  469. # ------ ------ --------------------------------------------------
  470. # None None None
  471. # None True None - must change value
  472. # None 8≤M≤15 M (or None)
  473. # 8≤N≤15 None Error!
  474. # 8≤N≤15 True N - must change value
  475. # 8≤N≤15 8≤M≤N M (or None)
  476. # 8≤N≤15 N<M≤15 N
  477. if self.client_max_window_bits is None:
  478. if client_max_window_bits is True:
  479. client_max_window_bits = self.client_max_window_bits
  480. else:
  481. if client_max_window_bits is None:
  482. raise NegotiationError("required client_max_window_bits")
  483. elif client_max_window_bits is True:
  484. client_max_window_bits = self.client_max_window_bits
  485. elif self.client_max_window_bits < client_max_window_bits:
  486. client_max_window_bits = self.client_max_window_bits
  487. return (
  488. _build_parameters(
  489. server_no_context_takeover,
  490. client_no_context_takeover,
  491. server_max_window_bits,
  492. client_max_window_bits,
  493. ),
  494. PerMessageDeflate(
  495. client_no_context_takeover, # remote_no_context_takeover
  496. server_no_context_takeover, # local_no_context_takeover
  497. client_max_window_bits or 15, # remote_max_window_bits
  498. server_max_window_bits or 15, # local_max_window_bits
  499. self.compress_settings,
  500. ),
  501. )
  502. def enable_server_permessage_deflate(
  503. extensions: Optional[Sequence[ServerExtensionFactory]],
  504. ) -> Sequence[ServerExtensionFactory]:
  505. """
  506. Enable Per-Message Deflate with default settings in server extensions.
  507. If the extension is already present, perhaps with non-default settings,
  508. the configuration isn't changed.
  509. """
  510. if extensions is None:
  511. extensions = []
  512. if not any(
  513. ext_factory.name == ServerPerMessageDeflateFactory.name
  514. for ext_factory in extensions
  515. ):
  516. extensions = list(extensions) + [ServerPerMessageDeflateFactory()]
  517. return extensions