framing.py 4.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135
  1. """
  2. :mod:`websockets.legacy.framing` reads and writes WebSocket frames.
  3. It deals with a single frame at a time. Anything that depends on the sequence
  4. of frames is implemented in :mod:`websockets.legacy.protocol`.
  5. See `section 5 of RFC 6455`_.
  6. .. _section 5 of RFC 6455: http://tools.ietf.org/html/rfc6455#section-5
  7. """
  8. import struct
  9. from typing import Any, Awaitable, Callable, Optional, Sequence
  10. from ..exceptions import PayloadTooBig, ProtocolError
  11. from ..frames import Frame as NewFrame, Opcode
  12. try:
  13. from ..speedups import apply_mask
  14. except ImportError: # pragma: no cover
  15. from ..utils import apply_mask
  16. class Frame(NewFrame):
  17. @classmethod
  18. async def read(
  19. cls,
  20. reader: Callable[[int], Awaitable[bytes]],
  21. *,
  22. mask: bool,
  23. max_size: Optional[int] = None,
  24. extensions: Optional[Sequence["extensions.Extension"]] = None,
  25. ) -> "Frame":
  26. """
  27. Read a WebSocket frame.
  28. :param reader: coroutine that reads exactly the requested number of
  29. bytes, unless the end of file is reached
  30. :param mask: whether the frame should be masked i.e. whether the read
  31. happens on the server side
  32. :param max_size: maximum payload size in bytes
  33. :param extensions: list of classes with a ``decode()`` method that
  34. transforms the frame and return a new frame; extensions are applied
  35. in reverse order
  36. :raises ~websockets.exceptions.PayloadTooBig: if the frame exceeds
  37. ``max_size``
  38. :raises ~websockets.exceptions.ProtocolError: if the frame
  39. contains incorrect values
  40. """
  41. # Read the header.
  42. data = await reader(2)
  43. head1, head2 = struct.unpack("!BB", data)
  44. # While not Pythonic, this is marginally faster than calling bool().
  45. fin = True if head1 & 0b10000000 else False
  46. rsv1 = True if head1 & 0b01000000 else False
  47. rsv2 = True if head1 & 0b00100000 else False
  48. rsv3 = True if head1 & 0b00010000 else False
  49. try:
  50. opcode = Opcode(head1 & 0b00001111)
  51. except ValueError as exc:
  52. raise ProtocolError("invalid opcode") from exc
  53. if (True if head2 & 0b10000000 else False) != mask:
  54. raise ProtocolError("incorrect masking")
  55. length = head2 & 0b01111111
  56. if length == 126:
  57. data = await reader(2)
  58. (length,) = struct.unpack("!H", data)
  59. elif length == 127:
  60. data = await reader(8)
  61. (length,) = struct.unpack("!Q", data)
  62. if max_size is not None and length > max_size:
  63. raise PayloadTooBig(f"over size limit ({length} > {max_size} bytes)")
  64. if mask:
  65. mask_bits = await reader(4)
  66. # Read the data.
  67. data = await reader(length)
  68. if mask:
  69. data = apply_mask(data, mask_bits)
  70. frame = cls(fin, opcode, data, rsv1, rsv2, rsv3)
  71. if extensions is None:
  72. extensions = []
  73. for extension in reversed(extensions):
  74. frame = cls(*extension.decode(frame, max_size=max_size))
  75. frame.check()
  76. return frame
  77. def write(
  78. self,
  79. write: Callable[[bytes], Any],
  80. *,
  81. mask: bool,
  82. extensions: Optional[Sequence["extensions.Extension"]] = None,
  83. ) -> None:
  84. """
  85. Write a WebSocket frame.
  86. :param frame: frame to write
  87. :param write: function that writes bytes
  88. :param mask: whether the frame should be masked i.e. whether the write
  89. happens on the client side
  90. :param extensions: list of classes with an ``encode()`` method that
  91. transform the frame and return a new frame; extensions are applied
  92. in order
  93. :raises ~websockets.exceptions.ProtocolError: if the frame
  94. contains incorrect values
  95. """
  96. # The frame is written in a single call to write in order to prevent
  97. # TCP fragmentation. See #68 for details. This also makes it safe to
  98. # send frames concurrently from multiple coroutines.
  99. write(self.serialize(mask=mask, extensions=extensions))
  100. # Backwards compatibility with previously documented public APIs
  101. from ..frames import parse_close # isort:skip # noqa
  102. from ..frames import prepare_ctrl as encode_data # isort:skip # noqa
  103. from ..frames import prepare_data # isort:skip # noqa
  104. from ..frames import serialize_close # isort:skip # noqa
  105. # at the bottom to allow circular import, because Extension depends on Frame
  106. from .. import extensions # isort:skip # noqa