protocol.py 27 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768
  1. from __future__ import annotations
  2. import enum
  3. import logging
  4. import uuid
  5. from collections.abc import Generator
  6. from .exceptions import (
  7. ConnectionClosed,
  8. ConnectionClosedError,
  9. ConnectionClosedOK,
  10. InvalidState,
  11. PayloadTooBig,
  12. ProtocolError,
  13. )
  14. from .extensions import Extension
  15. from .frames import (
  16. OK_CLOSE_CODES,
  17. OP_BINARY,
  18. OP_CLOSE,
  19. OP_CONT,
  20. OP_PING,
  21. OP_PONG,
  22. OP_TEXT,
  23. Close,
  24. CloseCode,
  25. Frame,
  26. )
  27. from .http11 import Request, Response
  28. from .streams import StreamReader
  29. from .typing import BytesLike, LoggerLike, Origin, Subprotocol
  30. __all__ = [
  31. "Protocol",
  32. "Side",
  33. "State",
  34. "SEND_EOF",
  35. ]
  36. Event = Request | Response | Frame
  37. """Events that :meth:`~Protocol.events_received` may return."""
  38. class Side(enum.IntEnum):
  39. """A WebSocket connection is either a server or a client."""
  40. SERVER, CLIENT = range(2)
  41. SERVER = Side.SERVER
  42. CLIENT = Side.CLIENT
  43. class State(enum.IntEnum):
  44. """A WebSocket connection is in one of these four states."""
  45. CONNECTING, OPEN, CLOSING, CLOSED = range(4)
  46. CONNECTING = State.CONNECTING
  47. OPEN = State.OPEN
  48. CLOSING = State.CLOSING
  49. CLOSED = State.CLOSED
  50. SEND_EOF = b""
  51. """Sentinel signaling that the TCP connection must be half-closed."""
  52. class Protocol:
  53. """
  54. Sans-I/O implementation of a WebSocket connection.
  55. Args:
  56. side: :attr:`~Side.CLIENT` or :attr:`~Side.SERVER`.
  57. state: Initial state of the WebSocket connection.
  58. max_size: Maximum size of incoming messages in bytes.
  59. :obj:`None` disables the limit. You may pass a ``(max_message_size,
  60. max_fragment_size)`` tuple to set different limits for messages and
  61. fragments when you expect long messages sent in short fragments.
  62. logger: Logger for this connection; depending on ``side``,
  63. defaults to ``logging.getLogger("websockets.client")``
  64. or ``logging.getLogger("websockets.server")``;
  65. see the :doc:`logging guide <../../topics/logging>` for details.
  66. """
  67. def __init__(
  68. self,
  69. side: Side,
  70. *,
  71. state: State = OPEN,
  72. max_size: tuple[int | None, int | None] | int | None = 2**20,
  73. logger: LoggerLike | None = None,
  74. ) -> None:
  75. # Unique identifier. For logs.
  76. self.id: uuid.UUID = uuid.uuid4()
  77. """Unique identifier of the connection. Useful in logs."""
  78. # Logger or LoggerAdapter for this connection.
  79. if logger is None:
  80. logger = logging.getLogger(f"websockets.{side.name.lower()}")
  81. self.logger: LoggerLike = logger
  82. """Logger for this connection."""
  83. # Track if DEBUG is enabled. Shortcut logging calls if it isn't.
  84. self.debug = logger.isEnabledFor(logging.DEBUG)
  85. # Connection side. CLIENT or SERVER.
  86. self.side = side
  87. # Connection state. Initially OPEN because subclasses handle CONNECTING.
  88. self.state = state
  89. # Maximum size of incoming messages in bytes.
  90. if isinstance(max_size, int) or max_size is None:
  91. self.max_message_size, self.max_fragment_size = max_size, None
  92. else:
  93. self.max_message_size, self.max_fragment_size = max_size
  94. # Current size of incoming message in bytes. Only set while reading a
  95. # fragmented message i.e. a data frames with the FIN bit not set.
  96. self.current_size: int | None = None
  97. # True while sending a fragmented message i.e. a data frames with the
  98. # FIN bit not set.
  99. self.expect_continuation_frame = False
  100. # WebSocket protocol parameters.
  101. self.origin: Origin | None = None
  102. self.extensions: list[Extension] = []
  103. self.subprotocol: Subprotocol | None = None
  104. # Close code and reason, set when a close frame is sent or received.
  105. self.close_rcvd: Close | None = None
  106. self.close_sent: Close | None = None
  107. self.close_rcvd_then_sent: bool | None = None
  108. # Track if an exception happened during the handshake.
  109. self.handshake_exc: Exception | None = None
  110. """
  111. Exception to raise if the opening handshake failed.
  112. :obj:`None` if the opening handshake succeeded.
  113. """
  114. # Track if send_eof() was called.
  115. self.eof_sent = False
  116. # Parser state.
  117. self.reader = StreamReader()
  118. self.events: list[Event] = []
  119. self.writes: list[bytes] = []
  120. self.parser = self.parse()
  121. next(self.parser) # start coroutine
  122. self.parser_exc: Exception | None = None
  123. @property
  124. def state(self) -> State:
  125. """
  126. State of the WebSocket connection.
  127. Defined in 4.1_, 4.2_, 7.1.3_, and 7.1.4_ of :rfc:`6455`.
  128. .. _4.1: https://datatracker.ietf.org/doc/html/rfc6455#section-4.1
  129. .. _4.2: https://datatracker.ietf.org/doc/html/rfc6455#section-4.2
  130. .. _7.1.3: https://datatracker.ietf.org/doc/html/rfc6455#section-7.1.3
  131. .. _7.1.4: https://datatracker.ietf.org/doc/html/rfc6455#section-7.1.4
  132. """
  133. return self._state
  134. @state.setter
  135. def state(self, state: State) -> None:
  136. if self.debug:
  137. self.logger.debug("= connection is %s", state.name)
  138. self._state = state
  139. @property
  140. def close_code(self) -> int | None:
  141. """
  142. WebSocket close code received from the remote endpoint.
  143. Defined in 7.1.5_ of :rfc:`6455`.
  144. .. _7.1.5: https://datatracker.ietf.org/doc/html/rfc6455#section-7.1.5
  145. :obj:`None` if the connection isn't closed yet.
  146. """
  147. if self.state is not CLOSED:
  148. return None
  149. elif self.close_rcvd is None:
  150. return CloseCode.ABNORMAL_CLOSURE
  151. else:
  152. return self.close_rcvd.code
  153. @property
  154. def close_reason(self) -> str | None:
  155. """
  156. WebSocket close reason received from the remote endpoint.
  157. Defined in 7.1.6_ of :rfc:`6455`.
  158. .. _7.1.6: https://datatracker.ietf.org/doc/html/rfc6455#section-7.1.6
  159. :obj:`None` if the connection isn't closed yet.
  160. """
  161. if self.state is not CLOSED:
  162. return None
  163. elif self.close_rcvd is None:
  164. return ""
  165. else:
  166. return self.close_rcvd.reason
  167. @property
  168. def close_exc(self) -> ConnectionClosed:
  169. """
  170. Exception to raise when trying to interact with a closed connection.
  171. Don't raise this exception while the connection :attr:`state`
  172. is :attr:`~websockets.protocol.State.CLOSING`; wait until
  173. it's :attr:`~websockets.protocol.State.CLOSED`.
  174. Indeed, the exception includes the close code and reason, which are
  175. known only once the connection is closed.
  176. Raises:
  177. AssertionError: If the connection isn't closed yet.
  178. """
  179. assert self.state is CLOSED, "connection isn't closed yet"
  180. exc_type: type[ConnectionClosed]
  181. if (
  182. self.close_rcvd is not None
  183. and self.close_sent is not None
  184. and self.close_rcvd.code in OK_CLOSE_CODES
  185. and self.close_sent.code in OK_CLOSE_CODES
  186. ):
  187. exc_type = ConnectionClosedOK
  188. else:
  189. exc_type = ConnectionClosedError
  190. exc: ConnectionClosed = exc_type(
  191. self.close_rcvd,
  192. self.close_sent,
  193. self.close_rcvd_then_sent,
  194. )
  195. # Chain to the exception raised in the parser, if any.
  196. exc.__cause__ = self.parser_exc
  197. return exc
  198. # Public methods for receiving data.
  199. def receive_data(self, data: bytes | bytearray) -> None:
  200. """
  201. Receive data from the network.
  202. After calling this method:
  203. - You must call :meth:`data_to_send` and send this data to the network.
  204. - You should call :meth:`events_received` and process resulting events.
  205. Raises:
  206. EOFError: If :meth:`receive_eof` was called earlier.
  207. """
  208. self.reader.feed_data(data)
  209. next(self.parser)
  210. def receive_eof(self) -> None:
  211. """
  212. Receive the end of the data stream from the network.
  213. After calling this method:
  214. - You must call :meth:`data_to_send` and send this data to the network;
  215. it will return ``[b""]``, signaling the end of the stream, or ``[]``.
  216. - You aren't expected to call :meth:`events_received`; it won't return
  217. any new events.
  218. :meth:`receive_eof` is idempotent.
  219. """
  220. if self.reader.eof:
  221. return
  222. self.reader.feed_eof()
  223. next(self.parser)
  224. # Public methods for sending events.
  225. def send_continuation(self, data: BytesLike, fin: bool) -> None:
  226. """
  227. Send a `Continuation frame`_.
  228. .. _Continuation frame:
  229. https://datatracker.ietf.org/doc/html/rfc6455#section-5.6
  230. Parameters:
  231. data: payload containing the same kind of data
  232. as the initial frame.
  233. fin: FIN bit; set it to :obj:`True` if this is the last frame
  234. of a fragmented message and to :obj:`False` otherwise.
  235. Raises:
  236. ProtocolError: If a fragmented message isn't in progress.
  237. """
  238. if not self.expect_continuation_frame:
  239. raise ProtocolError("unexpected continuation frame")
  240. if self._state is not OPEN:
  241. raise InvalidState(f"connection is {self.state.name.lower()}")
  242. self.expect_continuation_frame = not fin
  243. self.send_frame(Frame(OP_CONT, data, fin))
  244. def send_text(self, data: BytesLike, fin: bool = True) -> None:
  245. """
  246. Send a `Text frame`_.
  247. .. _Text frame:
  248. https://datatracker.ietf.org/doc/html/rfc6455#section-5.6
  249. Parameters:
  250. data: payload containing text encoded with UTF-8.
  251. fin: FIN bit; set it to :obj:`False` if this is the first frame of
  252. a fragmented message.
  253. Raises:
  254. ProtocolError: If a fragmented message is in progress.
  255. """
  256. if self.expect_continuation_frame:
  257. raise ProtocolError("expected a continuation frame")
  258. if self._state is not OPEN:
  259. raise InvalidState(f"connection is {self.state.name.lower()}")
  260. self.expect_continuation_frame = not fin
  261. self.send_frame(Frame(OP_TEXT, data, fin))
  262. def send_binary(self, data: BytesLike, fin: bool = True) -> None:
  263. """
  264. Send a `Binary frame`_.
  265. .. _Binary frame:
  266. https://datatracker.ietf.org/doc/html/rfc6455#section-5.6
  267. Parameters:
  268. data: payload containing arbitrary binary data.
  269. fin: FIN bit; set it to :obj:`False` if this is the first frame of
  270. a fragmented message.
  271. Raises:
  272. ProtocolError: If a fragmented message is in progress.
  273. """
  274. if self.expect_continuation_frame:
  275. raise ProtocolError("expected a continuation frame")
  276. if self._state is not OPEN:
  277. raise InvalidState(f"connection is {self.state.name.lower()}")
  278. self.expect_continuation_frame = not fin
  279. self.send_frame(Frame(OP_BINARY, data, fin))
  280. def send_close(self, code: CloseCode | int | None = None, reason: str = "") -> None:
  281. """
  282. Send a `Close frame`_.
  283. .. _Close frame:
  284. https://datatracker.ietf.org/doc/html/rfc6455#section-5.5.1
  285. Parameters:
  286. code: close code.
  287. reason: close reason.
  288. Raises:
  289. ProtocolError: If the code isn't valid or if a reason is provided
  290. without a code.
  291. """
  292. # While RFC 6455 doesn't rule out sending more than one close Frame,
  293. # websockets is conservative in what it sends and doesn't allow that.
  294. if self._state is not OPEN:
  295. raise InvalidState(f"connection is {self.state.name.lower()}")
  296. if code is None:
  297. if reason != "":
  298. raise ProtocolError("cannot send a reason without a code")
  299. close = Close(CloseCode.NO_STATUS_RCVD, "")
  300. data = b""
  301. else:
  302. close = Close(code, reason)
  303. data = close.serialize()
  304. # 7.1.3. The WebSocket Closing Handshake is Started
  305. self.send_frame(Frame(OP_CLOSE, data))
  306. # Since the state is OPEN, no close frame was received yet.
  307. # As a consequence, self.close_rcvd_then_sent remains None.
  308. assert self.close_rcvd is None
  309. self.close_sent = close
  310. self.state = CLOSING
  311. def send_ping(self, data: BytesLike) -> None:
  312. """
  313. Send a `Ping frame`_.
  314. .. _Ping frame:
  315. https://datatracker.ietf.org/doc/html/rfc6455#section-5.5.2
  316. Parameters:
  317. data: payload containing arbitrary binary data.
  318. """
  319. # RFC 6455 allows control frames after starting the closing handshake.
  320. if self._state is not OPEN and self._state is not CLOSING:
  321. raise InvalidState(f"connection is {self.state.name.lower()}")
  322. self.send_frame(Frame(OP_PING, data))
  323. def send_pong(self, data: BytesLike) -> None:
  324. """
  325. Send a `Pong frame`_.
  326. .. _Pong frame:
  327. https://datatracker.ietf.org/doc/html/rfc6455#section-5.5.3
  328. Parameters:
  329. data: payload containing arbitrary binary data.
  330. """
  331. # RFC 6455 allows control frames after starting the closing handshake.
  332. if self._state is not OPEN and self._state is not CLOSING:
  333. raise InvalidState(f"connection is {self.state.name.lower()}")
  334. self.send_frame(Frame(OP_PONG, data))
  335. def fail(self, code: CloseCode | int, reason: str = "") -> None:
  336. """
  337. `Fail the WebSocket connection`_.
  338. .. _Fail the WebSocket connection:
  339. https://datatracker.ietf.org/doc/html/rfc6455#section-7.1.7
  340. Parameters:
  341. code: close code
  342. reason: close reason
  343. Raises:
  344. ProtocolError: If the code isn't valid.
  345. """
  346. # 7.1.7. Fail the WebSocket Connection
  347. # Send a close frame when the state is OPEN (a close frame was already
  348. # sent if it's CLOSING), except when failing the connection because
  349. # of an error reading from or writing to the network.
  350. if self.state is OPEN:
  351. if code != CloseCode.ABNORMAL_CLOSURE:
  352. close = Close(code, reason)
  353. data = close.serialize()
  354. self.send_frame(Frame(OP_CLOSE, data))
  355. self.close_sent = close
  356. # If recv_messages() raised an exception upon receiving a close
  357. # frame but before echoing it, then close_rcvd is not None even
  358. # though the state is OPEN. This happens when the connection is
  359. # closed while receiving a fragmented message.
  360. if self.close_rcvd is not None:
  361. self.close_rcvd_then_sent = True
  362. self.state = CLOSING
  363. # When failing the connection, a server closes the TCP connection
  364. # without waiting for the client to complete the handshake, while a
  365. # client waits for the server to close the TCP connection, possibly
  366. # after sending a close frame that the client will ignore.
  367. if self.side is SERVER and not self.eof_sent:
  368. self.send_eof()
  369. # 7.1.7. Fail the WebSocket Connection "An endpoint MUST NOT continue
  370. # to attempt to process data(including a responding Close frame) from
  371. # the remote endpoint after being instructed to _Fail the WebSocket
  372. # Connection_."
  373. self.parser = self.discard()
  374. next(self.parser) # start coroutine
  375. # Public method for getting incoming events after receiving data.
  376. def events_received(self) -> list[Event]:
  377. """
  378. Fetch events generated from data received from the network.
  379. Call this method immediately after any of the ``receive_*()`` methods.
  380. Process resulting events, likely by passing them to the application.
  381. Returns:
  382. Events read from the connection.
  383. """
  384. events, self.events = self.events, []
  385. return events
  386. # Public method for getting outgoing data after receiving data or sending events.
  387. def data_to_send(self) -> list[bytes]:
  388. """
  389. Obtain data to send to the network.
  390. Call this method immediately after any of the ``receive_*()``,
  391. ``send_*()``, or :meth:`fail` methods.
  392. Write resulting data to the connection.
  393. The empty bytestring :data:`~websockets.protocol.SEND_EOF` signals
  394. the end of the data stream. When you receive it, half-close the TCP
  395. connection.
  396. Returns:
  397. Data to write to the connection.
  398. """
  399. writes, self.writes = self.writes, []
  400. return writes
  401. def close_expected(self) -> bool:
  402. """
  403. Tell if the TCP connection is expected to close soon.
  404. Call this method immediately after any of the ``receive_*()``,
  405. ``send_close()``, or :meth:`fail` methods.
  406. If it returns :obj:`True`, schedule closing the TCP connection after a
  407. short timeout if the other side hasn't already closed it.
  408. Returns:
  409. Whether the TCP connection is expected to close soon.
  410. """
  411. # During the opening handshake, when our state is CONNECTING, we expect
  412. # a TCP close if and only if the hansdake fails. When it does, we start
  413. # the TCP closing handshake by sending EOF with send_eof().
  414. # Once the opening handshake completes successfully, we expect a TCP
  415. # close if and only if we sent a close frame, meaning that our state
  416. # progressed to CLOSING:
  417. # * Normal closure: once we send a close frame, we expect a TCP close:
  418. # server waits for client to complete the TCP closing handshake;
  419. # client waits for server to initiate the TCP closing handshake.
  420. # * Abnormal closure: we always send a close frame and the same logic
  421. # applies, except on EOFError where we don't send a close frame
  422. # because we already received the TCP close, so we don't expect it.
  423. # If our state is CLOSED, we already received a TCP close so we don't
  424. # expect it anymore.
  425. # Micro-optimization: put the most common case first
  426. if self.state is OPEN:
  427. return False
  428. if self.state is CLOSING:
  429. return True
  430. if self.state is CLOSED:
  431. return False
  432. assert self.state is CONNECTING
  433. return self.eof_sent
  434. # Private methods for receiving data.
  435. def parse(self) -> Generator[None]:
  436. """
  437. Parse incoming data into frames.
  438. :meth:`receive_data` and :meth:`receive_eof` run this generator
  439. coroutine until it needs more data or reaches EOF.
  440. :meth:`parse` never raises an exception. Instead, it sets the
  441. :attr:`parser_exc` and yields control.
  442. """
  443. try:
  444. while True:
  445. if (yield from self.reader.at_eof()):
  446. if self.debug:
  447. self.logger.debug("< EOF")
  448. # If the WebSocket connection is closed cleanly, with a
  449. # closing handhshake, recv_frame() substitutes parse()
  450. # with discard(). This branch is reached only when the
  451. # connection isn't closed cleanly.
  452. raise EOFError("unexpected end of stream")
  453. max_size = None
  454. if self.max_message_size is not None:
  455. if self.current_size is None:
  456. max_size = self.max_message_size
  457. else:
  458. max_size = self.max_message_size - self.current_size
  459. if self.max_fragment_size is not None:
  460. if max_size is None:
  461. max_size = self.max_fragment_size
  462. else:
  463. max_size = min(max_size, self.max_fragment_size)
  464. # During a normal closure, execution ends here on the next
  465. # iteration of the loop after receiving a close frame. At
  466. # this point, recv_frame() replaced parse() by discard().
  467. frame = yield from Frame.parse(
  468. self.reader.read_exact,
  469. mask=self.side is SERVER,
  470. max_size=max_size,
  471. extensions=self.extensions,
  472. )
  473. if self.debug:
  474. self.logger.debug("< %s", frame)
  475. self.recv_frame(frame)
  476. except ProtocolError as exc:
  477. self.fail(CloseCode.PROTOCOL_ERROR, str(exc))
  478. self.parser_exc = exc
  479. except EOFError as exc:
  480. self.fail(CloseCode.ABNORMAL_CLOSURE, str(exc))
  481. self.parser_exc = exc
  482. except UnicodeDecodeError as exc:
  483. self.fail(CloseCode.INVALID_DATA, f"{exc.reason} at position {exc.start}")
  484. self.parser_exc = exc
  485. except PayloadTooBig as exc:
  486. exc.set_current_size(self.current_size)
  487. self.fail(CloseCode.MESSAGE_TOO_BIG, str(exc))
  488. self.parser_exc = exc
  489. except Exception as exc:
  490. self.logger.error("parser failed", exc_info=True)
  491. # Don't include exception details, which may be security-sensitive.
  492. self.fail(CloseCode.INTERNAL_ERROR)
  493. self.parser_exc = exc
  494. # During an abnormal closure, execution ends here after catching an
  495. # exception. At this point, fail() replaced parse() by discard().
  496. yield
  497. raise AssertionError("parse() shouldn't step after error")
  498. def discard(self) -> Generator[None]:
  499. """
  500. Discard incoming data.
  501. This coroutine replaces :meth:`parse`:
  502. - after receiving a close frame, during a normal closure (1.4);
  503. - after sending a close frame, during an abnormal closure (7.1.7).
  504. """
  505. # After the opening handshake completes, the server closes the TCP
  506. # connection in the same circumstances where discard() replaces parse().
  507. # The client closes it when it receives EOF from the server or times
  508. # out. (The latter case cannot be handled in this Sans-I/O layer.)
  509. assert (self.side is SERVER or self.state is CONNECTING) == (self.eof_sent)
  510. while not (yield from self.reader.at_eof()):
  511. self.reader.discard()
  512. if self.debug:
  513. self.logger.debug("< EOF")
  514. # A server closes the TCP connection immediately, while a client
  515. # waits for the server to close the TCP connection.
  516. if self.side is CLIENT and self.state is not CONNECTING:
  517. self.send_eof()
  518. self.state = CLOSED
  519. # If discard() completes normally, execution ends here.
  520. yield
  521. # Once the reader reaches EOF, its feed_data/eof() methods raise an
  522. # error, so our receive_data/eof() methods don't step the generator.
  523. raise AssertionError("discard() shouldn't step after EOF")
  524. def recv_frame(self, frame: Frame) -> None:
  525. """
  526. Process an incoming frame.
  527. """
  528. if frame.opcode is OP_TEXT or frame.opcode is OP_BINARY:
  529. if self.current_size is not None:
  530. raise ProtocolError("expected a continuation frame")
  531. if not frame.fin:
  532. self.current_size = len(frame.data)
  533. elif frame.opcode is OP_CONT:
  534. if self.current_size is None:
  535. raise ProtocolError("unexpected continuation frame")
  536. if frame.fin:
  537. self.current_size = None
  538. else:
  539. self.current_size += len(frame.data)
  540. elif frame.opcode is OP_PING:
  541. # 5.5.2. Ping: "Upon receipt of a Ping frame, an endpoint MUST
  542. # send a Pong frame in response"
  543. pong_frame = Frame(OP_PONG, frame.data)
  544. self.send_frame(pong_frame)
  545. elif frame.opcode is OP_PONG:
  546. # 5.5.3 Pong: "A response to an unsolicited Pong frame is not
  547. # expected."
  548. pass
  549. elif frame.opcode is OP_CLOSE:
  550. # 7.1.5. The WebSocket Connection Close Code
  551. # 7.1.6. The WebSocket Connection Close Reason
  552. self.close_rcvd = Close.parse(frame.data)
  553. if self.state is CLOSING:
  554. assert self.close_sent is not None
  555. self.close_rcvd_then_sent = False
  556. if self.current_size is not None:
  557. raise ProtocolError("incomplete fragmented message")
  558. # 5.5.1 Close: "If an endpoint receives a Close frame and did
  559. # not previously send a Close frame, the endpoint MUST send a
  560. # Close frame in response. (When sending a Close frame in
  561. # response, the endpoint typically echos the status code it
  562. # received.)"
  563. if self.state is OPEN:
  564. # Echo the original data instead of re-serializing it with
  565. # Close.serialize() because that fails when the close frame
  566. # is empty and Close.parse() synthesizes a 1005 close code.
  567. # The rest is identical to send_close().
  568. self.send_frame(Frame(OP_CLOSE, frame.data))
  569. self.close_sent = self.close_rcvd
  570. self.close_rcvd_then_sent = True
  571. self.state = CLOSING
  572. # 7.1.2. Start the WebSocket Closing Handshake: "Once an
  573. # endpoint has both sent and received a Close control frame,
  574. # that endpoint SHOULD _Close the WebSocket Connection_"
  575. # A server closes the TCP connection immediately, while a client
  576. # waits for the server to close the TCP connection.
  577. if self.side is SERVER:
  578. self.send_eof()
  579. # 1.4. Closing Handshake: "after receiving a control frame
  580. # indicating the connection should be closed, a peer discards
  581. # any further data received."
  582. # RFC 6455 allows reading Ping and Pong frames after a Close frame.
  583. # However, that doesn't seem useful; websockets doesn't support it.
  584. self.parser = self.discard()
  585. next(self.parser) # start coroutine
  586. else:
  587. # This can't happen because Frame.parse() validates opcodes.
  588. raise AssertionError(f"unexpected opcode: {frame.opcode:02x}")
  589. self.events.append(frame)
  590. # Private methods for sending events.
  591. def send_frame(self, frame: Frame) -> None:
  592. if self.debug:
  593. self.logger.debug("> %s", frame)
  594. self.writes.append(
  595. frame.serialize(
  596. mask=self.side is CLIENT,
  597. extensions=self.extensions,
  598. )
  599. )
  600. def send_eof(self) -> None:
  601. assert not self.eof_sent
  602. self.eof_sent = True
  603. if self.debug:
  604. self.logger.debug("> EOF")
  605. self.writes.append(SEND_EOF)