connection.py 14 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409
  1. import enum
  2. import logging
  3. from typing import Generator, List, Optional, Union
  4. from .exceptions import InvalidState, PayloadTooBig, ProtocolError
  5. from .extensions.base import Extension
  6. from .frames import (
  7. OP_BINARY,
  8. OP_CLOSE,
  9. OP_CONT,
  10. OP_PING,
  11. OP_PONG,
  12. OP_TEXT,
  13. Frame,
  14. parse_close,
  15. serialize_close,
  16. )
  17. from .http11 import Request, Response
  18. from .streams import StreamReader
  19. from .typing import Origin, Subprotocol
  20. __all__ = [
  21. "Connection",
  22. "Side",
  23. "State",
  24. "SEND_EOF",
  25. ]
  26. logger = logging.getLogger(__name__)
  27. Event = Union[Request, Response, Frame]
  28. # A WebSocket connection is either a server or a client.
  29. class Side(enum.IntEnum):
  30. SERVER, CLIENT = range(2)
  31. SERVER = Side.SERVER
  32. CLIENT = Side.CLIENT
  33. # A WebSocket connection goes through the following four states, in order:
  34. class State(enum.IntEnum):
  35. CONNECTING, OPEN, CLOSING, CLOSED = range(4)
  36. CONNECTING = State.CONNECTING
  37. OPEN = State.OPEN
  38. CLOSING = State.CLOSING
  39. CLOSED = State.CLOSED
  40. # Sentinel to signal that the connection should be closed.
  41. SEND_EOF = b""
  42. class Connection:
  43. def __init__(
  44. self,
  45. side: Side,
  46. state: State = OPEN,
  47. max_size: Optional[int] = 2 ** 20,
  48. ) -> None:
  49. # Connection side. CLIENT or SERVER.
  50. self.side = side
  51. # Connnection state. CONNECTING and CLOSED states are handled in subclasses.
  52. logger.debug("%s - initial state: %s", self.side, state.name)
  53. self.state = state
  54. # Maximum size of incoming messages in bytes.
  55. self.max_size = max_size
  56. # Current size of incoming message in bytes. Only set while reading a
  57. # fragmented message i.e. a data frames with the FIN bit not set.
  58. self.cur_size: Optional[int] = None
  59. # True while sending a fragmented message i.e. a data frames with the
  60. # FIN bit not set.
  61. self.expect_continuation_frame = False
  62. # WebSocket protocol parameters.
  63. self.origin: Optional[Origin] = None
  64. self.extensions: List[Extension] = []
  65. self.subprotocol: Optional[Subprotocol] = None
  66. # Connection state isn't enough to tell if a close frame was received:
  67. # when this side closes the connection, state is CLOSING as soon as a
  68. # close frame is sent, before a close frame is received.
  69. self.close_frame_received = False
  70. # Close code and reason. Set when receiving a close frame or when the
  71. # TCP connection drops.
  72. self.close_code: int
  73. self.close_reason: str
  74. # Track if send_eof() was called.
  75. self.eof_sent = False
  76. # Parser state.
  77. self.reader = StreamReader()
  78. self.events: List[Event] = []
  79. self.writes: List[bytes] = []
  80. self.parser = self.parse()
  81. next(self.parser) # start coroutine
  82. self.parser_exc: Optional[Exception] = None
  83. def set_state(self, state: State) -> None:
  84. logger.debug(
  85. "%s - state change: %s > %s", self.side, self.state.name, state.name
  86. )
  87. self.state = state
  88. # Public APIs for receiving data.
  89. def receive_data(self, data: bytes) -> None:
  90. """
  91. Receive data from the connection.
  92. After calling this method:
  93. - You must call :meth:`data_to_send` and send this data.
  94. - You should call :meth:`events_received` and process these events.
  95. """
  96. self.reader.feed_data(data)
  97. self.step_parser()
  98. def receive_eof(self) -> None:
  99. """
  100. Receive the end of the data stream from the connection.
  101. After calling this method:
  102. - You must call :meth:`data_to_send` and send this data.
  103. - You shouldn't call :meth:`events_received` as it won't
  104. return any new events.
  105. """
  106. self.reader.feed_eof()
  107. self.step_parser()
  108. # Public APIs for sending events.
  109. def send_continuation(self, data: bytes, fin: bool) -> None:
  110. """
  111. Send a continuation frame.
  112. """
  113. if not self.expect_continuation_frame:
  114. raise ProtocolError("unexpected continuation frame")
  115. self.expect_continuation_frame = not fin
  116. self.send_frame(Frame(fin, OP_CONT, data))
  117. def send_text(self, data: bytes, fin: bool = True) -> None:
  118. """
  119. Send a text frame.
  120. """
  121. if self.expect_continuation_frame:
  122. raise ProtocolError("expected a continuation frame")
  123. self.expect_continuation_frame = not fin
  124. self.send_frame(Frame(fin, OP_TEXT, data))
  125. def send_binary(self, data: bytes, fin: bool = True) -> None:
  126. """
  127. Send a binary frame.
  128. """
  129. if self.expect_continuation_frame:
  130. raise ProtocolError("expected a continuation frame")
  131. self.expect_continuation_frame = not fin
  132. self.send_frame(Frame(fin, OP_BINARY, data))
  133. def send_close(self, code: Optional[int] = None, reason: str = "") -> None:
  134. """
  135. Send a connection close frame.
  136. """
  137. if self.expect_continuation_frame:
  138. raise ProtocolError("expected a continuation frame")
  139. if code is None:
  140. if reason != "":
  141. raise ValueError("cannot send a reason without a code")
  142. data = b""
  143. else:
  144. data = serialize_close(code, reason)
  145. self.send_frame(Frame(True, OP_CLOSE, data))
  146. # send_frame() guarantees that self.state is OPEN at this point.
  147. # 7.1.3. The WebSocket Closing Handshake is Started
  148. self.set_state(CLOSING)
  149. if self.side is SERVER:
  150. self.send_eof()
  151. def send_ping(self, data: bytes) -> None:
  152. """
  153. Send a ping frame.
  154. """
  155. self.send_frame(Frame(True, OP_PING, data))
  156. def send_pong(self, data: bytes) -> None:
  157. """
  158. Send a pong frame.
  159. """
  160. self.send_frame(Frame(True, OP_PONG, data))
  161. # Public API for getting incoming events after receiving data.
  162. def events_received(self) -> List[Event]:
  163. """
  164. Return events read from the connection.
  165. Call this method immediately after calling any of the ``receive_*()``
  166. methods and process the events.
  167. """
  168. events, self.events = self.events, []
  169. return events
  170. # Public API for getting outgoing data after receiving data or sending events.
  171. def data_to_send(self) -> List[bytes]:
  172. """
  173. Return data to write to the connection.
  174. Call this method immediately after calling any of the ``receive_*()``
  175. or ``send_*()`` methods and write the data to the connection.
  176. The empty bytestring signals the end of the data stream.
  177. """
  178. writes, self.writes = self.writes, []
  179. return writes
  180. # Private APIs for receiving data.
  181. def fail_connection(self, code: int = 1006, reason: str = "") -> None:
  182. # Send a close frame when the state is OPEN (a close frame was already
  183. # sent if it's CLOSING), except when failing the connection because of
  184. # an error reading from or writing to the network.
  185. if code != 1006 and self.state is OPEN:
  186. self.send_frame(Frame(True, OP_CLOSE, serialize_close(code, reason)))
  187. self.set_state(CLOSING)
  188. if not self.eof_sent:
  189. self.send_eof()
  190. def step_parser(self) -> None:
  191. # Run parser until more data is needed or EOF
  192. try:
  193. next(self.parser)
  194. except StopIteration:
  195. # This happens if receive_data() or receive_eof() is called after
  196. # the parser raised an exception. (It cannot happen after reaching
  197. # EOF because receive_data() or receive_eof() would fail earlier.)
  198. assert self.parser_exc is not None
  199. raise RuntimeError(
  200. "cannot receive data or EOF after an error"
  201. ) from self.parser_exc
  202. except ProtocolError as exc:
  203. self.fail_connection(1002, str(exc))
  204. self.parser_exc = exc
  205. raise
  206. except EOFError as exc:
  207. self.fail_connection(1006, str(exc))
  208. self.parser_exc = exc
  209. raise
  210. except UnicodeDecodeError as exc:
  211. self.fail_connection(1007, f"{exc.reason} at position {exc.start}")
  212. self.parser_exc = exc
  213. raise
  214. except PayloadTooBig as exc:
  215. self.fail_connection(1009, str(exc))
  216. self.parser_exc = exc
  217. raise
  218. except Exception as exc:
  219. logger.error("unexpected exception in parser", exc_info=True)
  220. # Don't include exception details, which may be security-sensitive.
  221. self.fail_connection(1011)
  222. self.parser_exc = exc
  223. raise
  224. def parse(self) -> Generator[None, None, None]:
  225. while True:
  226. eof = yield from self.reader.at_eof()
  227. if eof:
  228. if self.close_frame_received:
  229. if not self.eof_sent:
  230. self.send_eof()
  231. yield
  232. # Once the reader reaches EOF, its feed_data/eof() methods
  233. # raise an error, so our receive_data/eof() methods never
  234. # call step_parser(), so the generator shouldn't resume
  235. # executing until it's garbage collected.
  236. raise AssertionError(
  237. "parser shouldn't step after EOF"
  238. ) # pragma: no cover
  239. else:
  240. raise EOFError("unexpected end of stream")
  241. if self.max_size is None:
  242. max_size = None
  243. elif self.cur_size is None:
  244. max_size = self.max_size
  245. else:
  246. max_size = self.max_size - self.cur_size
  247. frame = yield from Frame.parse(
  248. self.reader.read_exact,
  249. mask=self.side is SERVER,
  250. max_size=max_size,
  251. extensions=self.extensions,
  252. )
  253. if frame.opcode is OP_TEXT or frame.opcode is OP_BINARY:
  254. # 5.5.1 Close: "The application MUST NOT send any more data
  255. # frames after sending a Close frame."
  256. if self.close_frame_received:
  257. raise ProtocolError("data frame after close frame")
  258. if self.cur_size is not None:
  259. raise ProtocolError("expected a continuation frame")
  260. if frame.fin:
  261. self.cur_size = None
  262. else:
  263. self.cur_size = len(frame.data)
  264. elif frame.opcode is OP_CONT:
  265. # 5.5.1 Close: "The application MUST NOT send any more data
  266. # frames after sending a Close frame."
  267. if self.close_frame_received:
  268. raise ProtocolError("data frame after close frame")
  269. if self.cur_size is None:
  270. raise ProtocolError("unexpected continuation frame")
  271. if frame.fin:
  272. self.cur_size = None
  273. else:
  274. self.cur_size += len(frame.data)
  275. elif frame.opcode is OP_PING:
  276. # 5.5.2. Ping: "Upon receipt of a Ping frame, an endpoint MUST
  277. # send a Pong frame in response, unless it already received a
  278. # Close frame."
  279. if not self.close_frame_received:
  280. pong_frame = Frame(True, OP_PONG, frame.data)
  281. self.send_frame(pong_frame)
  282. elif frame.opcode is OP_PONG:
  283. # 5.5.3 Pong: "A response to an unsolicited Pong frame is not
  284. # expected."
  285. pass
  286. elif frame.opcode is OP_CLOSE:
  287. self.close_frame_received = True
  288. # 7.1.5. The WebSocket Connection Close Code
  289. # 7.1.6. The WebSocket Connection Close Reason
  290. self.close_code, self.close_reason = parse_close(frame.data)
  291. if self.cur_size is not None:
  292. raise ProtocolError("incomplete fragmented message")
  293. # 5.5.1 Close: "If an endpoint receives a Close frame and did
  294. # not previously send a Close frame, the endpoint MUST send a
  295. # Close frame in response. (When sending a Close frame in
  296. # response, the endpoint typically echos the status code it
  297. # received.)"
  298. if self.state is OPEN:
  299. # Echo the original data instead of re-serializing it with
  300. # serialize_close() because that fails when the close frame
  301. # is empty and parse_close() synthetizes a 1005 close code.
  302. # The rest is identical to send_close().
  303. self.send_frame(Frame(True, OP_CLOSE, frame.data))
  304. self.set_state(CLOSING)
  305. if self.side is SERVER:
  306. self.send_eof()
  307. else: # pragma: no cover
  308. # This can't happen because Frame.parse() validates opcodes.
  309. raise AssertionError(f"unexpected opcode: {frame.opcode:02x}")
  310. self.events.append(frame)
  311. # Private APIs for sending events.
  312. def send_frame(self, frame: Frame) -> None:
  313. # Defensive assertion for protocol compliance.
  314. if self.state is not OPEN:
  315. raise InvalidState(
  316. f"cannot write to a WebSocket in the {self.state.name} state"
  317. )
  318. logger.debug("%s > %r", self.side, frame)
  319. self.writes.append(
  320. frame.serialize(mask=self.side is CLIENT, extensions=self.extensions)
  321. )
  322. def send_eof(self) -> None:
  323. assert not self.eof_sent
  324. self.eof_sent = True
  325. logger.debug("%s > EOF", self.side)
  326. self.writes.append(SEND_EOF)