您最多选择25个主题 主题必须以字母或数字开头,可以包含连字符 (-),并且长度不得超过35个字符
 
 
 
 

1362 行
52 KiB

  1. """
  2. :mod:`websockets.protocol` handles WebSocket control and data frames.
  3. See `sections 4 to 8 of RFC 6455`_.
  4. .. _sections 4 to 8 of RFC 6455: http://tools.ietf.org/html/rfc6455#section-4
  5. """
  6. import asyncio
  7. import codecs
  8. import collections
  9. import enum
  10. import logging
  11. import random
  12. import struct
  13. import warnings
  14. from typing import (
  15. Any,
  16. AsyncIterable,
  17. AsyncIterator,
  18. Awaitable,
  19. Deque,
  20. Dict,
  21. Iterable,
  22. List,
  23. Optional,
  24. Union,
  25. cast,
  26. )
  27. from .exceptions import (
  28. ConnectionClosed,
  29. ConnectionClosedError,
  30. ConnectionClosedOK,
  31. InvalidState,
  32. PayloadTooBig,
  33. ProtocolError,
  34. )
  35. from .extensions.base import Extension
  36. from .framing import *
  37. from .handshake import *
  38. from .http import Headers
  39. from .typing import Data
  40. __all__ = ["WebSocketCommonProtocol"]
  41. logger = logging.getLogger(__name__)
  42. # A WebSocket connection goes through the following four states, in order:
  43. class State(enum.IntEnum):
  44. CONNECTING, OPEN, CLOSING, CLOSED = range(4)
  45. # In order to ensure consistency, the code always checks the current value of
  46. # WebSocketCommonProtocol.state before assigning a new value and never yields
  47. # between the check and the assignment.
  48. class WebSocketCommonProtocol(asyncio.StreamReaderProtocol):
  49. """
  50. :class:`~asyncio.Protocol` subclass implementing the data transfer phase.
  51. Once the WebSocket connection is established, during the data transfer
  52. phase, the protocol is almost symmetrical between the server side and the
  53. client side. :class:`WebSocketCommonProtocol` implements logic that's
  54. shared between servers and clients..
  55. Subclasses such as :class:`~websockets.server.WebSocketServerProtocol` and
  56. :class:`~websockets.client.WebSocketClientProtocol` implement the opening
  57. handshake, which is different between servers and clients.
  58. :class:`WebSocketCommonProtocol` performs four functions:
  59. * It runs a task that stores incoming data frames in a queue and makes
  60. them available with the :meth:`recv` coroutine.
  61. * It sends outgoing data frames with the :meth:`send` coroutine.
  62. * It deals with control frames automatically.
  63. * It performs the closing handshake.
  64. :class:`WebSocketCommonProtocol` supports asynchronous iteration::
  65. async for message in websocket:
  66. await process(message)
  67. The iterator yields incoming messages. It exits normally when the
  68. connection is closed with the close code 1000 (OK) or 1001 (going away).
  69. It raises a :exc:`~websockets.exceptions.ConnectionClosedError` exception
  70. when the connection is closed with any other code.
  71. Once the connection is open, a `Ping frame`_ is sent every
  72. ``ping_interval`` seconds. This serves as a keepalive. It helps keeping
  73. the connection open, especially in the presence of proxies with short
  74. timeouts on inactive connections. Set ``ping_interval`` to ``None`` to
  75. disable this behavior.
  76. .. _Ping frame: https://tools.ietf.org/html/rfc6455#section-5.5.2
  77. If the corresponding `Pong frame`_ isn't received within ``ping_timeout``
  78. seconds, the connection is considered unusable and is closed with
  79. code 1011. This ensures that the remote endpoint remains responsive. Set
  80. ``ping_timeout`` to ``None`` to disable this behavior.
  81. .. _Pong frame: https://tools.ietf.org/html/rfc6455#section-5.5.3
  82. The ``close_timeout`` parameter defines a maximum wait time in seconds for
  83. completing the closing handshake and terminating the TCP connection.
  84. :meth:`close` completes in at most ``4 * close_timeout`` on the server
  85. side and ``5 * close_timeout`` on the client side.
  86. ``close_timeout`` needs to be a parameter of the protocol because
  87. ``websockets`` usually calls :meth:`close` implicitly:
  88. - on the server side, when the connection handler terminates,
  89. - on the client side, when exiting the context manager for the connection.
  90. To apply a timeout to any other API, wrap it in :func:`~asyncio.wait_for`.
  91. The ``max_size`` parameter enforces the maximum size for incoming messages
  92. in bytes. The default value is 1 MiB. ``None`` disables the limit. If a
  93. message larger than the maximum size is received, :meth:`recv` will
  94. raise :exc:`~websockets.exceptions.ConnectionClosedError` and the
  95. connection will be closed with code 1009.
  96. The ``max_queue`` parameter sets the maximum length of the queue that
  97. holds incoming messages. The default value is ``32``. ``None`` disables
  98. the limit. Messages are added to an in-memory queue when they're received;
  99. then :meth:`recv` pops from that queue. In order to prevent excessive
  100. memory consumption when messages are received faster than they can be
  101. processed, the queue must be bounded. If the queue fills up, the protocol
  102. stops processing incoming data until :meth:`recv` is called. In this
  103. situation, various receive buffers (at least in ``asyncio`` and in the OS)
  104. will fill up, then the TCP receive window will shrink, slowing down
  105. transmission to avoid packet loss.
  106. Since Python can use up to 4 bytes of memory to represent a single
  107. character, each connection may use up to ``4 * max_size * max_queue``
  108. bytes of memory to store incoming messages. By default, this is 128 MiB.
  109. You may want to lower the limits, depending on your application's
  110. requirements.
  111. The ``read_limit`` argument sets the high-water limit of the buffer for
  112. incoming bytes. The low-water limit is half the high-water limit. The
  113. default value is 64 KiB, half of asyncio's default (based on the current
  114. implementation of :class:`~asyncio.StreamReader`).
  115. The ``write_limit`` argument sets the high-water limit of the buffer for
  116. outgoing bytes. The low-water limit is a quarter of the high-water limit.
  117. The default value is 64 KiB, equal to asyncio's default (based on the
  118. current implementation of ``FlowControlMixin``).
  119. As soon as the HTTP request and response in the opening handshake are
  120. processed:
  121. * the request path is available in the :attr:`path` attribute;
  122. * the request and response HTTP headers are available in the
  123. :attr:`request_headers` and :attr:`response_headers` attributes,
  124. which are :class:`~websockets.http.Headers` instances.
  125. If a subprotocol was negotiated, it's available in the :attr:`subprotocol`
  126. attribute.
  127. Once the connection is closed, the code is available in the
  128. :attr:`close_code` attribute and the reason in :attr:`close_reason`.
  129. All these attributes must be treated as read-only.
  130. """
  131. # There are only two differences between the client-side and server-side
  132. # behavior: masking the payload and closing the underlying TCP connection.
  133. # Set is_client = True/False and side = "client"/"server" to pick a side.
  134. is_client: bool
  135. side: str = "undefined"
  136. def __init__(
  137. self,
  138. *,
  139. ping_interval: float = 20,
  140. ping_timeout: float = 20,
  141. close_timeout: Optional[float] = None,
  142. max_size: int = 2 ** 20,
  143. max_queue: int = 2 ** 5,
  144. read_limit: int = 2 ** 16,
  145. write_limit: int = 2 ** 16,
  146. loop: Optional[asyncio.AbstractEventLoop] = None,
  147. # The following arguments are kept only for backwards compatibility.
  148. host: Optional[str] = None,
  149. port: Optional[int] = None,
  150. secure: Optional[bool] = None,
  151. legacy_recv: bool = False,
  152. timeout: Optional[float] = None,
  153. ) -> None:
  154. # Backwards compatibility: close_timeout used to be called timeout.
  155. if timeout is None:
  156. timeout = 10
  157. else:
  158. warnings.warn("rename timeout to close_timeout", DeprecationWarning)
  159. # If both are specified, timeout is ignored.
  160. if close_timeout is None:
  161. close_timeout = timeout
  162. self.ping_interval = ping_interval
  163. self.ping_timeout = ping_timeout
  164. self.close_timeout = close_timeout
  165. self.max_size = max_size
  166. self.max_queue = max_queue
  167. self.read_limit = read_limit
  168. self.write_limit = write_limit
  169. # Store a reference to loop to avoid relying on self._loop, a private
  170. # attribute of StreamReaderProtocol, inherited from FlowControlMixin.
  171. if loop is None:
  172. loop = asyncio.get_event_loop()
  173. self.loop = loop
  174. self._host = host
  175. self._port = port
  176. self._secure = secure
  177. self.legacy_recv = legacy_recv
  178. # Configure read buffer limits. The high-water limit is defined by
  179. # ``self.read_limit``. The ``limit`` argument controls the line length
  180. # limit and half the buffer limit of :class:`~asyncio.StreamReader`.
  181. # That's why it must be set to half of ``self.read_limit``.
  182. stream_reader = asyncio.StreamReader(limit=read_limit // 2, loop=loop)
  183. super().__init__(stream_reader, self.client_connected, loop)
  184. self.reader: asyncio.StreamReader
  185. self.writer: asyncio.StreamWriter
  186. self._drain_lock = asyncio.Lock(loop=loop)
  187. # This class implements the data transfer and closing handshake, which
  188. # are shared between the client-side and the server-side.
  189. # Subclasses implement the opening handshake and, on success, execute
  190. # :meth:`connection_open` to change the state to OPEN.
  191. self.state = State.CONNECTING
  192. logger.debug("%s - state = CONNECTING", self.side)
  193. # HTTP protocol parameters.
  194. self.path: str
  195. self.request_headers: Headers
  196. self.response_headers: Headers
  197. # WebSocket protocol parameters.
  198. self.extensions: List[Extension] = []
  199. self.subprotocol: Optional[str] = None
  200. # The close code and reason are set when receiving a close frame or
  201. # losing the TCP connection.
  202. self.close_code: int
  203. self.close_reason: str
  204. # Completed when the connection state becomes CLOSED. Translates the
  205. # :meth:`connection_lost` callback to a :class:`~asyncio.Future`
  206. # that can be awaited. (Other :class:`~asyncio.Protocol` callbacks are
  207. # translated by ``self.stream_reader``).
  208. self.connection_lost_waiter: asyncio.Future[None] = loop.create_future()
  209. # Queue of received messages.
  210. self.messages: Deque[Data] = collections.deque()
  211. self._pop_message_waiter: Optional[asyncio.Future[None]] = None
  212. self._put_message_waiter: Optional[asyncio.Future[None]] = None
  213. # Protect sending fragmented messages.
  214. self._fragmented_message_waiter: Optional[asyncio.Future[None]] = None
  215. # Mapping of ping IDs to waiters, in chronological order.
  216. self.pings: Dict[bytes, asyncio.Future[None]] = {}
  217. # Task running the data transfer.
  218. self.transfer_data_task: asyncio.Task[None]
  219. # Exception that occurred during data transfer, if any.
  220. self.transfer_data_exc: Optional[BaseException] = None
  221. # Task sending keepalive pings.
  222. self.keepalive_ping_task: asyncio.Task[None]
  223. # Task closing the TCP connection.
  224. self.close_connection_task: asyncio.Task[None]
  225. def client_connected(
  226. self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter
  227. ) -> None:
  228. """
  229. Callback when the TCP connection is established.
  230. Record references to the stream reader and the stream writer to avoid
  231. using private attributes ``_stream_reader`` and ``_stream_writer`` of
  232. :class:`~asyncio.StreamReaderProtocol`.
  233. """
  234. self.reader = reader
  235. self.writer = writer
  236. def connection_open(self) -> None:
  237. """
  238. Callback when the WebSocket opening handshake completes.
  239. Enter the OPEN state and start the data transfer phase.
  240. """
  241. # 4.1. The WebSocket Connection is Established.
  242. assert self.state is State.CONNECTING
  243. self.state = State.OPEN
  244. logger.debug("%s - state = OPEN", self.side)
  245. # Start the task that receives incoming WebSocket messages.
  246. self.transfer_data_task = self.loop.create_task(self.transfer_data())
  247. # Start the task that sends pings at regular intervals.
  248. self.keepalive_ping_task = self.loop.create_task(self.keepalive_ping())
  249. # Start the task that eventually closes the TCP connection.
  250. self.close_connection_task = self.loop.create_task(self.close_connection())
  251. @property
  252. def host(self) -> Optional[str]:
  253. alternative = "remote_address" if self.is_client else "local_address"
  254. warnings.warn(f"use {alternative}[0] instead of host", DeprecationWarning)
  255. return self._host
  256. @property
  257. def port(self) -> Optional[int]:
  258. alternative = "remote_address" if self.is_client else "local_address"
  259. warnings.warn(f"use {alternative}[1] instead of port", DeprecationWarning)
  260. return self._port
  261. @property
  262. def secure(self) -> Optional[bool]:
  263. warnings.warn(f"don't use secure", DeprecationWarning)
  264. return self._secure
  265. # Public API
  266. @property
  267. def local_address(self) -> Any:
  268. """
  269. Local address of the connection.
  270. This is a ``(host, port)`` tuple or ``None`` if the connection hasn't
  271. been established yet.
  272. """
  273. if self.writer is None:
  274. return None
  275. return self.writer.get_extra_info("sockname")
  276. @property
  277. def remote_address(self) -> Any:
  278. """
  279. Remote address of the connection.
  280. This is a ``(host, port)`` tuple or ``None`` if the connection hasn't
  281. been established yet.
  282. """
  283. if self.writer is None:
  284. return None
  285. return self.writer.get_extra_info("peername")
  286. @property
  287. def open(self) -> bool:
  288. """
  289. ``True`` when the connection is usable.
  290. It may be used to detect disconnections. However, this approach is
  291. discouraged per the EAFP_ principle.
  292. When ``open`` is ``False``, using the connection raises a
  293. :exc:`~websockets.exceptions.ConnectionClosed` exception.
  294. .. _EAFP: https://docs.python.org/3/glossary.html#term-eafp
  295. """
  296. return self.state is State.OPEN and not self.transfer_data_task.done()
  297. @property
  298. def closed(self) -> bool:
  299. """
  300. ``True`` once the connection is closed.
  301. Be aware that both :attr:`open` and :attr:`closed` are ``False`` during
  302. the opening and closing sequences.
  303. """
  304. return self.state is State.CLOSED
  305. async def wait_closed(self) -> None:
  306. """
  307. Wait until the connection is closed.
  308. This is identical to :attr:`closed`, except it can be awaited.
  309. This can make it easier to handle connection termination, regardless
  310. of its cause, in tasks that interact with the WebSocket connection.
  311. """
  312. await asyncio.shield(self.connection_lost_waiter)
  313. async def __aiter__(self) -> AsyncIterator[Data]:
  314. """
  315. Iterate on received messages.
  316. Exit normally when the connection is closed with code 1000 or 1001.
  317. Raise an exception in other cases.
  318. """
  319. try:
  320. while True:
  321. yield await self.recv()
  322. except ConnectionClosedOK:
  323. return
  324. async def recv(self) -> Data:
  325. """
  326. Receive the next message.
  327. Return a :class:`str` for a text frame and :class:`bytes` for a binary
  328. frame.
  329. When the end of the message stream is reached, :meth:`recv` raises
  330. :exc:`~websockets.exceptions.ConnectionClosed`. Specifically, it
  331. raises :exc:`~websockets.exceptions.ConnectionClosedOK` after a normal
  332. connection closure and
  333. :exc:`~websockets.exceptions.ConnectionClosedError` after a protocol
  334. error or a network failure.
  335. .. versionchanged:: 3.0
  336. :meth:`recv` used to return ``None`` instead. Refer to the
  337. changelog for details.
  338. Canceling :meth:`recv` is safe. There's no risk of losing the next
  339. message. The next invocation of :meth:`recv` will return it. This
  340. makes it possible to enforce a timeout by wrapping :meth:`recv` in
  341. :func:`~asyncio.wait_for`.
  342. :raises ~websockets.exceptions.ConnectionClosed: when the
  343. connection is closed
  344. :raises RuntimeError: if two coroutines call :meth:`recv` concurrently
  345. """
  346. if self._pop_message_waiter is not None:
  347. raise RuntimeError(
  348. "cannot call recv while another coroutine "
  349. "is already waiting for the next message"
  350. )
  351. # Don't await self.ensure_open() here:
  352. # - messages could be available in the queue even if the connection
  353. # is closed;
  354. # - messages could be received before the closing frame even if the
  355. # connection is closing.
  356. # Wait until there's a message in the queue (if necessary) or the
  357. # connection is closed.
  358. while len(self.messages) <= 0:
  359. pop_message_waiter: asyncio.Future[None] = self.loop.create_future()
  360. self._pop_message_waiter = pop_message_waiter
  361. try:
  362. # If asyncio.wait() is canceled, it doesn't cancel
  363. # pop_message_waiter and self.transfer_data_task.
  364. await asyncio.wait(
  365. [pop_message_waiter, self.transfer_data_task],
  366. loop=self.loop,
  367. return_when=asyncio.FIRST_COMPLETED,
  368. )
  369. finally:
  370. self._pop_message_waiter = None
  371. # If asyncio.wait(...) exited because self.transfer_data_task
  372. # completed before receiving a new message, raise a suitable
  373. # exception (or return None if legacy_recv is enabled).
  374. if not pop_message_waiter.done():
  375. if self.legacy_recv:
  376. return None # type: ignore
  377. else:
  378. # Wait until the connection is closed to raise
  379. # ConnectionClosed with the correct code and reason.
  380. await self.ensure_open()
  381. # Pop a message from the queue.
  382. message = self.messages.popleft()
  383. # Notify transfer_data().
  384. if self._put_message_waiter is not None:
  385. self._put_message_waiter.set_result(None)
  386. self._put_message_waiter = None
  387. return message
  388. async def send(
  389. self, message: Union[Data, Iterable[Data], AsyncIterable[Data]]
  390. ) -> None:
  391. """
  392. Send a message.
  393. A string (:class:`str`) is sent as a `Text frame`_. A bytestring or
  394. bytes-like object (:class:`bytes`, :class:`bytearray`, or
  395. :class:`memoryview`) is sent as a `Binary frame`_.
  396. .. _Text frame: https://tools.ietf.org/html/rfc6455#section-5.6
  397. .. _Binary frame: https://tools.ietf.org/html/rfc6455#section-5.6
  398. :meth:`send` also accepts an iterable or an asynchronous iterable of
  399. strings, bytestrings, or bytes-like objects. In that case the message
  400. is fragmented. Each item is treated as a message fragment and sent in
  401. its own frame. All items must be of the same type, or else
  402. :meth:`send` will raise a :exc:`TypeError` and the connection will be
  403. closed.
  404. Canceling :meth:`send` is discouraged. Instead, you should close the
  405. connection with :meth:`close`. Indeed, there only two situations where
  406. :meth:`send` yields control to the event loop:
  407. 1. The write buffer is full. If you don't want to wait until enough
  408. data is sent, your only alternative is to close the connection.
  409. :meth:`close` will likely time out then abort the TCP connection.
  410. 2. ``message`` is an asynchronous iterator. Stopping in the middle of
  411. a fragmented message will cause a protocol error. Closing the
  412. connection has the same effect.
  413. :raises TypeError: for unsupported inputs
  414. """
  415. await self.ensure_open()
  416. # While sending a fragmented message, prevent sending other messages
  417. # until all fragments are sent.
  418. while self._fragmented_message_waiter is not None:
  419. await asyncio.shield(self._fragmented_message_waiter)
  420. # Unfragmented message -- this case must be handled first because
  421. # strings and bytes-like objects are iterable.
  422. if isinstance(message, (str, bytes, bytearray, memoryview)):
  423. opcode, data = prepare_data(message)
  424. await self.write_frame(True, opcode, data)
  425. # Fragmented message -- regular iterator.
  426. elif isinstance(message, Iterable):
  427. # Work around https://github.com/python/mypy/issues/6227
  428. message = cast(Iterable[Data], message)
  429. iter_message = iter(message)
  430. try:
  431. message_chunk = next(iter_message)
  432. except StopIteration:
  433. return
  434. opcode, data = prepare_data(message_chunk)
  435. self._fragmented_message_waiter = asyncio.Future()
  436. try:
  437. # First fragment.
  438. await self.write_frame(False, opcode, data)
  439. # Other fragments.
  440. for message_chunk in iter_message:
  441. confirm_opcode, data = prepare_data(message_chunk)
  442. if confirm_opcode != opcode:
  443. raise TypeError("data contains inconsistent types")
  444. await self.write_frame(False, OP_CONT, data)
  445. # Final fragment.
  446. await self.write_frame(True, OP_CONT, b"")
  447. except Exception:
  448. # We're half-way through a fragmented message and we can't
  449. # complete it. This makes the connection unusable.
  450. self.fail_connection(1011)
  451. raise
  452. finally:
  453. self._fragmented_message_waiter.set_result(None)
  454. self._fragmented_message_waiter = None
  455. # Fragmented message -- asynchronous iterator
  456. elif isinstance(message, AsyncIterable):
  457. # aiter_message = aiter(message) without aiter
  458. aiter_message = type(message).__aiter__(message)
  459. try:
  460. # message_chunk = anext(aiter_message) without anext
  461. message_chunk = await type(aiter_message).__anext__(aiter_message)
  462. except StopAsyncIteration:
  463. return
  464. opcode, data = prepare_data(message_chunk)
  465. self._fragmented_message_waiter = asyncio.Future()
  466. try:
  467. # First fragment.
  468. await self.write_frame(False, opcode, data)
  469. # Other fragments.
  470. async for message_chunk in aiter_message:
  471. confirm_opcode, data = prepare_data(message_chunk)
  472. if confirm_opcode != opcode:
  473. raise TypeError("data contains inconsistent types")
  474. await self.write_frame(False, OP_CONT, data)
  475. # Final fragment.
  476. await self.write_frame(True, OP_CONT, b"")
  477. except Exception:
  478. # We're half-way through a fragmented message and we can't
  479. # complete it. This makes the connection unusable.
  480. self.fail_connection(1011)
  481. raise
  482. finally:
  483. self._fragmented_message_waiter.set_result(None)
  484. self._fragmented_message_waiter = None
  485. else:
  486. raise TypeError("data must be bytes, str, or iterable")
  487. async def close(self, code: int = 1000, reason: str = "") -> None:
  488. """
  489. Perform the closing handshake.
  490. :meth:`close` waits for the other end to complete the handshake and
  491. for the TCP connection to terminate. As a consequence, there's no need
  492. to await :meth:`wait_closed`; :meth:`close` already does it.
  493. :meth:`close` is idempotent: it doesn't do anything once the
  494. connection is closed.
  495. Wrapping :func:`close` in :func:`~asyncio.create_task` is safe, given
  496. that errors during connection termination aren't particularly useful.
  497. Canceling :meth:`close` is discouraged. If it takes too long, you can
  498. set a shorter ``close_timeout``. If you don't want to wait, let the
  499. Python process exit, then the OS will close the TCP connection.
  500. :param code: WebSocket close code
  501. :param reason: WebSocket close reason
  502. """
  503. try:
  504. await asyncio.wait_for(
  505. self.write_close_frame(serialize_close(code, reason)),
  506. self.close_timeout,
  507. loop=self.loop,
  508. )
  509. except asyncio.TimeoutError:
  510. # If the close frame cannot be sent because the send buffers
  511. # are full, the closing handshake won't complete anyway.
  512. # Fail the connection to shut down faster.
  513. self.fail_connection()
  514. # If no close frame is received within the timeout, wait_for() cancels
  515. # the data transfer task and raises TimeoutError.
  516. # If close() is called multiple times concurrently and one of these
  517. # calls hits the timeout, the data transfer task will be cancelled.
  518. # Other calls will receive a CancelledError here.
  519. try:
  520. # If close() is canceled during the wait, self.transfer_data_task
  521. # is canceled before the timeout elapses.
  522. await asyncio.wait_for(
  523. self.transfer_data_task, self.close_timeout, loop=self.loop
  524. )
  525. except (asyncio.TimeoutError, asyncio.CancelledError):
  526. pass
  527. # Wait for the close connection task to close the TCP connection.
  528. await asyncio.shield(self.close_connection_task)
  529. async def ping(self, data: Optional[Data] = None) -> Awaitable[None]:
  530. """
  531. Send a ping.
  532. Return a :class:`~asyncio.Future` which will be completed when the
  533. corresponding pong is received and which you may ignore if you don't
  534. want to wait.
  535. A ping may serve as a keepalive or as a check that the remote endpoint
  536. received all messages up to this point::
  537. pong_waiter = await ws.ping()
  538. await pong_waiter # only if you want to wait for the pong
  539. By default, the ping contains four random bytes. This payload may be
  540. overridden with the optional ``data`` argument which must be a string
  541. (which will be encoded to UTF-8) or a bytes-like object.
  542. Canceling :meth:`ping` is discouraged. If :meth:`ping` doesn't return
  543. immediately, it means the write buffer is full. If you don't want to
  544. wait, you should close the connection.
  545. Canceling the :class:`~asyncio.Future` returned by :meth:`ping` has no
  546. effect.
  547. """
  548. await self.ensure_open()
  549. if data is not None:
  550. data = encode_data(data)
  551. # Protect against duplicates if a payload is explicitly set.
  552. if data in self.pings:
  553. raise ValueError("already waiting for a pong with the same data")
  554. # Generate a unique random payload otherwise.
  555. while data is None or data in self.pings:
  556. data = struct.pack("!I", random.getrandbits(32))
  557. self.pings[data] = self.loop.create_future()
  558. await self.write_frame(True, OP_PING, data)
  559. return asyncio.shield(self.pings[data])
  560. async def pong(self, data: Data = b"") -> None:
  561. """
  562. Send a pong.
  563. An unsolicited pong may serve as a unidirectional heartbeat.
  564. The payload may be set with the optional ``data`` argument which must
  565. be a string (which will be encoded to UTF-8) or a bytes-like object.
  566. Canceling :meth:`pong` is discouraged for the same reason as
  567. :meth:`ping`.
  568. """
  569. await self.ensure_open()
  570. data = encode_data(data)
  571. await self.write_frame(True, OP_PONG, data)
  572. # Private methods - no guarantees.
  573. def connection_closed_exc(self) -> ConnectionClosed:
  574. exception: ConnectionClosed
  575. if self.close_code == 1000 or self.close_code == 1001:
  576. exception = ConnectionClosedOK(self.close_code, self.close_reason)
  577. else:
  578. exception = ConnectionClosedError(self.close_code, self.close_reason)
  579. # Chain to the exception that terminated data transfer, if any.
  580. exception.__cause__ = self.transfer_data_exc
  581. return exception
  582. async def ensure_open(self) -> None:
  583. """
  584. Check that the WebSocket connection is open.
  585. Raise :exc:`~websockets.exceptions.ConnectionClosed` if it isn't.
  586. """
  587. # Handle cases from most common to least common for performance.
  588. if self.state is State.OPEN:
  589. # If self.transfer_data_task exited without a closing handshake,
  590. # self.close_connection_task may be closing the connection, going
  591. # straight from OPEN to CLOSED.
  592. if self.transfer_data_task.done():
  593. await asyncio.shield(self.close_connection_task)
  594. raise self.connection_closed_exc()
  595. else:
  596. return
  597. if self.state is State.CLOSED:
  598. raise self.connection_closed_exc()
  599. if self.state is State.CLOSING:
  600. # If we started the closing handshake, wait for its completion to
  601. # get the proper close code and reason. self.close_connection_task
  602. # will complete within 4 or 5 * close_timeout after close(). The
  603. # CLOSING state also occurs when failing the connection. In that
  604. # case self.close_connection_task will complete even faster.
  605. await asyncio.shield(self.close_connection_task)
  606. raise self.connection_closed_exc()
  607. # Control may only reach this point in buggy third-party subclasses.
  608. assert self.state is State.CONNECTING
  609. raise InvalidState("WebSocket connection isn't established yet")
  610. async def transfer_data(self) -> None:
  611. """
  612. Read incoming messages and put them in a queue.
  613. This coroutine runs in a task until the closing handshake is started.
  614. """
  615. try:
  616. while True:
  617. message = await self.read_message()
  618. # Exit the loop when receiving a close frame.
  619. if message is None:
  620. break
  621. # Wait until there's room in the queue (if necessary).
  622. if self.max_queue is not None:
  623. while len(self.messages) >= self.max_queue:
  624. self._put_message_waiter = self.loop.create_future()
  625. try:
  626. await asyncio.shield(self._put_message_waiter)
  627. finally:
  628. self._put_message_waiter = None
  629. # Put the message in the queue.
  630. self.messages.append(message)
  631. # Notify recv().
  632. if self._pop_message_waiter is not None:
  633. self._pop_message_waiter.set_result(None)
  634. self._pop_message_waiter = None
  635. except asyncio.CancelledError as exc:
  636. self.transfer_data_exc = exc
  637. # If fail_connection() cancels this task, avoid logging the error
  638. # twice and failing the connection again.
  639. raise
  640. except ProtocolError as exc:
  641. self.transfer_data_exc = exc
  642. self.fail_connection(1002)
  643. except (ConnectionError, EOFError) as exc:
  644. # Reading data with self.reader.readexactly may raise:
  645. # - most subclasses of ConnectionError if the TCP connection
  646. # breaks, is reset, or is aborted;
  647. # - IncompleteReadError, a subclass of EOFError, if fewer
  648. # bytes are available than requested.
  649. self.transfer_data_exc = exc
  650. self.fail_connection(1006)
  651. except UnicodeDecodeError as exc:
  652. self.transfer_data_exc = exc
  653. self.fail_connection(1007)
  654. except PayloadTooBig as exc:
  655. self.transfer_data_exc = exc
  656. self.fail_connection(1009)
  657. except Exception as exc:
  658. # This shouldn't happen often because exceptions expected under
  659. # regular circumstances are handled above. If it does, consider
  660. # catching and handling more exceptions.
  661. logger.error("Error in data transfer", exc_info=True)
  662. self.transfer_data_exc = exc
  663. self.fail_connection(1011)
  664. async def read_message(self) -> Optional[Data]:
  665. """
  666. Read a single message from the connection.
  667. Re-assemble data frames if the message is fragmented.
  668. Return ``None`` when the closing handshake is started.
  669. """
  670. frame = await self.read_data_frame(max_size=self.max_size)
  671. # A close frame was received.
  672. if frame is None:
  673. return None
  674. if frame.opcode == OP_TEXT:
  675. text = True
  676. elif frame.opcode == OP_BINARY:
  677. text = False
  678. else: # frame.opcode == OP_CONT
  679. raise ProtocolError("unexpected opcode")
  680. # Shortcut for the common case - no fragmentation
  681. if frame.fin:
  682. return frame.data.decode("utf-8") if text else frame.data
  683. # 5.4. Fragmentation
  684. chunks: List[Data] = []
  685. max_size = self.max_size
  686. if text:
  687. decoder_factory = codecs.getincrementaldecoder("utf-8")
  688. # https://github.com/python/typeshed/pull/2752
  689. decoder = decoder_factory(errors="strict") # type: ignore
  690. if max_size is None:
  691. def append(frame: Frame) -> None:
  692. nonlocal chunks
  693. chunks.append(decoder.decode(frame.data, frame.fin))
  694. else:
  695. def append(frame: Frame) -> None:
  696. nonlocal chunks, max_size
  697. chunks.append(decoder.decode(frame.data, frame.fin))
  698. max_size -= len(frame.data)
  699. else:
  700. if max_size is None:
  701. def append(frame: Frame) -> None:
  702. nonlocal chunks
  703. chunks.append(frame.data)
  704. else:
  705. def append(frame: Frame) -> None:
  706. nonlocal chunks, max_size
  707. chunks.append(frame.data)
  708. max_size -= len(frame.data)
  709. append(frame)
  710. while not frame.fin:
  711. frame = await self.read_data_frame(max_size=max_size)
  712. if frame is None:
  713. raise ProtocolError("incomplete fragmented message")
  714. if frame.opcode != OP_CONT:
  715. raise ProtocolError("unexpected opcode")
  716. append(frame)
  717. # mypy cannot figure out that chunks have the proper type.
  718. return ("" if text else b"").join(chunks) # type: ignore
  719. async def read_data_frame(self, max_size: int) -> Optional[Frame]:
  720. """
  721. Read a single data frame from the connection.
  722. Process control frames received before the next data frame.
  723. Return ``None`` if a close frame is encountered before any data frame.
  724. """
  725. # 6.2. Receiving Data
  726. while True:
  727. frame = await self.read_frame(max_size)
  728. # 5.5. Control Frames
  729. if frame.opcode == OP_CLOSE:
  730. # 7.1.5. The WebSocket Connection Close Code
  731. # 7.1.6. The WebSocket Connection Close Reason
  732. self.close_code, self.close_reason = parse_close(frame.data)
  733. try:
  734. # Echo the original data instead of re-serializing it with
  735. # serialize_close() because that fails when the close frame
  736. # is empty and parse_close() synthetizes a 1005 close code.
  737. await self.write_close_frame(frame.data)
  738. except ConnectionClosed:
  739. # It doesn't really matter if the connection was closed
  740. # before we could send back a close frame.
  741. pass
  742. return None
  743. elif frame.opcode == OP_PING:
  744. # Answer pings.
  745. ping_hex = frame.data.hex() or "[empty]"
  746. logger.debug(
  747. "%s - received ping, sending pong: %s", self.side, ping_hex
  748. )
  749. await self.pong(frame.data)
  750. elif frame.opcode == OP_PONG:
  751. # Acknowledge pings on solicited pongs.
  752. if frame.data in self.pings:
  753. logger.debug(
  754. "%s - received solicited pong: %s",
  755. self.side,
  756. frame.data.hex() or "[empty]",
  757. )
  758. # Acknowledge all pings up to the one matching this pong.
  759. ping_id = None
  760. ping_ids = []
  761. for ping_id, ping in self.pings.items():
  762. ping_ids.append(ping_id)
  763. if not ping.done():
  764. ping.set_result(None)
  765. if ping_id == frame.data:
  766. break
  767. else: # pragma: no cover
  768. assert False, "ping_id is in self.pings"
  769. # Remove acknowledged pings from self.pings.
  770. for ping_id in ping_ids:
  771. del self.pings[ping_id]
  772. ping_ids = ping_ids[:-1]
  773. if ping_ids:
  774. pings_hex = ", ".join(
  775. ping_id.hex() or "[empty]" for ping_id in ping_ids
  776. )
  777. plural = "s" if len(ping_ids) > 1 else ""
  778. logger.debug(
  779. "%s - acknowledged previous ping%s: %s",
  780. self.side,
  781. plural,
  782. pings_hex,
  783. )
  784. else:
  785. logger.debug(
  786. "%s - received unsolicited pong: %s",
  787. self.side,
  788. frame.data.hex() or "[empty]",
  789. )
  790. # 5.6. Data Frames
  791. else:
  792. return frame
  793. async def read_frame(self, max_size: int) -> Frame:
  794. """
  795. Read a single frame from the connection.
  796. """
  797. frame = await Frame.read(
  798. self.reader.readexactly,
  799. mask=not self.is_client,
  800. max_size=max_size,
  801. extensions=self.extensions,
  802. )
  803. logger.debug("%s < %r", self.side, frame)
  804. return frame
  805. async def write_frame(
  806. self, fin: bool, opcode: int, data: bytes, *, _expected_state: int = State.OPEN
  807. ) -> None:
  808. # Defensive assertion for protocol compliance.
  809. if self.state is not _expected_state: # pragma: no cover
  810. raise InvalidState(
  811. f"Cannot write to a WebSocket in the {self.state.name} state"
  812. )
  813. frame = Frame(fin, opcode, data)
  814. logger.debug("%s > %r", self.side, frame)
  815. frame.write(self.writer.write, mask=self.is_client, extensions=self.extensions)
  816. try:
  817. # drain() cannot be called concurrently by multiple coroutines:
  818. # http://bugs.python.org/issue29930. Remove this lock when no
  819. # version of Python where this bugs exists is supported anymore.
  820. async with self._drain_lock:
  821. # Handle flow control automatically.
  822. await self.writer.drain()
  823. except ConnectionError:
  824. # Terminate the connection if the socket died.
  825. self.fail_connection()
  826. # Wait until the connection is closed to raise ConnectionClosed
  827. # with the correct code and reason.
  828. await self.ensure_open()
  829. async def write_close_frame(self, data: bytes = b"") -> None:
  830. """
  831. Write a close frame if and only if the connection state is OPEN.
  832. This dedicated coroutine must be used for writing close frames to
  833. ensure that at most one close frame is sent on a given connection.
  834. """
  835. # Test and set the connection state before sending the close frame to
  836. # avoid sending two frames in case of concurrent calls.
  837. if self.state is State.OPEN:
  838. # 7.1.3. The WebSocket Closing Handshake is Started
  839. self.state = State.CLOSING
  840. logger.debug("%s - state = CLOSING", self.side)
  841. # 7.1.2. Start the WebSocket Closing Handshake
  842. await self.write_frame(True, OP_CLOSE, data, _expected_state=State.CLOSING)
  843. async def keepalive_ping(self) -> None:
  844. """
  845. Send a Ping frame and wait for a Pong frame at regular intervals.
  846. This coroutine exits when the connection terminates and one of the
  847. following happens:
  848. - :meth:`ping` raises :exc:`ConnectionClosed`, or
  849. - :meth:`close_connection` cancels :attr:`keepalive_ping_task`.
  850. """
  851. if self.ping_interval is None:
  852. return
  853. try:
  854. while True:
  855. await asyncio.sleep(self.ping_interval, loop=self.loop)
  856. # ping() raises CancelledError if the connection is closed,
  857. # when close_connection() cancels self.keepalive_ping_task.
  858. # ping() raises ConnectionClosed if the connection is lost,
  859. # when connection_lost() calls abort_pings().
  860. ping_waiter = await self.ping()
  861. if self.ping_timeout is not None:
  862. try:
  863. await asyncio.wait_for(
  864. ping_waiter, self.ping_timeout, loop=self.loop
  865. )
  866. except asyncio.TimeoutError:
  867. logger.debug("%s ! timed out waiting for pong", self.side)
  868. self.fail_connection(1011)
  869. break
  870. except asyncio.CancelledError:
  871. raise
  872. except ConnectionClosed:
  873. pass
  874. except Exception:
  875. logger.warning("Unexpected exception in keepalive ping task", exc_info=True)
  876. async def close_connection(self) -> None:
  877. """
  878. 7.1.1. Close the WebSocket Connection
  879. When the opening handshake succeeds, :meth:`connection_open` starts
  880. this coroutine in a task. It waits for the data transfer phase to
  881. complete then it closes the TCP connection cleanly.
  882. When the opening handshake fails, :meth:`fail_connection` does the
  883. same. There's no data transfer phase in that case.
  884. """
  885. try:
  886. # Wait for the data transfer phase to complete.
  887. if hasattr(self, "transfer_data_task"):
  888. try:
  889. await self.transfer_data_task
  890. except asyncio.CancelledError:
  891. pass
  892. # Cancel the keepalive ping task.
  893. if hasattr(self, "keepalive_ping_task"):
  894. self.keepalive_ping_task.cancel()
  895. # A client should wait for a TCP close from the server.
  896. if self.is_client and hasattr(self, "transfer_data_task"):
  897. if await self.wait_for_connection_lost():
  898. return
  899. logger.debug("%s ! timed out waiting for TCP close", self.side)
  900. # Half-close the TCP connection if possible (when there's no TLS).
  901. if self.writer.can_write_eof():
  902. logger.debug("%s x half-closing TCP connection", self.side)
  903. self.writer.write_eof()
  904. if await self.wait_for_connection_lost():
  905. return
  906. logger.debug("%s ! timed out waiting for TCP close", self.side)
  907. finally:
  908. # The try/finally ensures that the transport never remains open,
  909. # even if this coroutine is canceled (for example).
  910. # If connection_lost() was called, the TCP connection is closed.
  911. # However, if TLS is enabled, the transport still needs closing.
  912. # Else asyncio complains: ResourceWarning: unclosed transport.
  913. try:
  914. writer_is_closing = self.writer.is_closing # type: ignore
  915. except AttributeError: # pragma: no cover
  916. # Python < 3.7
  917. writer_is_closing = self.writer.transport.is_closing
  918. if self.connection_lost_waiter.done() and writer_is_closing():
  919. return
  920. # Close the TCP connection. Buffers are flushed asynchronously.
  921. logger.debug("%s x closing TCP connection", self.side)
  922. self.writer.close()
  923. if await self.wait_for_connection_lost():
  924. return
  925. logger.debug("%s ! timed out waiting for TCP close", self.side)
  926. # Abort the TCP connection. Buffers are discarded.
  927. logger.debug("%s x aborting TCP connection", self.side)
  928. # mypy thinks self.writer.transport is a BaseTransport, not a Transport.
  929. self.writer.transport.abort() # type: ignore
  930. # connection_lost() is called quickly after aborting.
  931. await self.wait_for_connection_lost()
  932. async def wait_for_connection_lost(self) -> bool:
  933. """
  934. Wait until the TCP connection is closed or ``self.close_timeout`` elapses.
  935. Return ``True`` if the connection is closed and ``False`` otherwise.
  936. """
  937. if not self.connection_lost_waiter.done():
  938. try:
  939. await asyncio.wait_for(
  940. asyncio.shield(self.connection_lost_waiter),
  941. self.close_timeout,
  942. loop=self.loop,
  943. )
  944. except asyncio.TimeoutError:
  945. pass
  946. # Re-check self.connection_lost_waiter.done() synchronously because
  947. # connection_lost() could run between the moment the timeout occurs
  948. # and the moment this coroutine resumes running.
  949. return self.connection_lost_waiter.done()
  950. def fail_connection(self, code: int = 1006, reason: str = "") -> None:
  951. """
  952. 7.1.7. Fail the WebSocket Connection
  953. This requires:
  954. 1. Stopping all processing of incoming data, which means cancelling
  955. :attr:`transfer_data_task`. The close code will be 1006 unless a
  956. close frame was received earlier.
  957. 2. Sending a close frame with an appropriate code if the opening
  958. handshake succeeded and the other side is likely to process it.
  959. 3. Closing the connection. :meth:`close_connection` takes care of
  960. this once :attr:`transfer_data_task` exits after being canceled.
  961. (The specification describes these steps in the opposite order.)
  962. """
  963. logger.debug(
  964. "%s ! failing %s WebSocket connection with code %d",
  965. self.side,
  966. self.state.name,
  967. code,
  968. )
  969. # Cancel transfer_data_task if the opening handshake succeeded.
  970. # cancel() is idempotent and ignored if the task is done already.
  971. if hasattr(self, "transfer_data_task"):
  972. self.transfer_data_task.cancel()
  973. # Send a close frame when the state is OPEN (a close frame was already
  974. # sent if it's CLOSING), except when failing the connection because of
  975. # an error reading from or writing to the network.
  976. # Don't send a close frame if the connection is broken.
  977. if code != 1006 and self.state is State.OPEN:
  978. frame_data = serialize_close(code, reason)
  979. # Write the close frame without draining the write buffer.
  980. # Keeping fail_connection() synchronous guarantees it can't
  981. # get stuck and simplifies the implementation of the callers.
  982. # Not drainig the write buffer is acceptable in this context.
  983. # This duplicates a few lines of code from write_close_frame()
  984. # and write_frame().
  985. self.state = State.CLOSING
  986. logger.debug("%s - state = CLOSING", self.side)
  987. frame = Frame(True, OP_CLOSE, frame_data)
  988. logger.debug("%s > %r", self.side, frame)
  989. frame.write(
  990. self.writer.write, mask=self.is_client, extensions=self.extensions
  991. )
  992. # Start close_connection_task if the opening handshake didn't succeed.
  993. if not hasattr(self, "close_connection_task"):
  994. self.close_connection_task = self.loop.create_task(self.close_connection())
  995. def abort_pings(self) -> None:
  996. """
  997. Raise ConnectionClosed in pending keepalive pings.
  998. They'll never receive a pong once the connection is closed.
  999. """
  1000. assert self.state is State.CLOSED
  1001. exc = self.connection_closed_exc()
  1002. for ping in self.pings.values():
  1003. ping.set_exception(exc)
  1004. # If the exception is never retrieved, it will be logged when ping
  1005. # is garbage-collected. This is confusing for users.
  1006. # Given that ping is done (with an exception), canceling it does
  1007. # nothing, but it prevents logging the exception.
  1008. ping.cancel()
  1009. if self.pings:
  1010. pings_hex = ", ".join(ping_id.hex() or "[empty]" for ping_id in self.pings)
  1011. plural = "s" if len(self.pings) > 1 else ""
  1012. logger.debug(
  1013. "%s - aborted pending ping%s: %s", self.side, plural, pings_hex
  1014. )
  1015. # asyncio.StreamReaderProtocol methods
  1016. def connection_made(self, transport: asyncio.BaseTransport) -> None:
  1017. """
  1018. Configure write buffer limits.
  1019. The high-water limit is defined by ``self.write_limit``.
  1020. The low-water limit currently defaults to ``self.write_limit // 4`` in
  1021. :meth:`~asyncio.WriteTransport.set_write_buffer_limits`, which should
  1022. be all right for reasonable use cases of this library.
  1023. This is the earliest point where we can get hold of the transport,
  1024. which means it's the best point for configuring it.
  1025. """
  1026. logger.debug("%s - event = connection_made(%s)", self.side, transport)
  1027. # mypy thinks transport is a BaseTransport, not a Transport.
  1028. transport.set_write_buffer_limits(self.write_limit) # type: ignore
  1029. super().connection_made(transport)
  1030. def eof_received(self) -> bool:
  1031. """
  1032. Close the transport after receiving EOF.
  1033. Since Python 3.5, `:meth:~StreamReaderProtocol.eof_received` returns
  1034. ``True`` on non-TLS connections.
  1035. See http://bugs.python.org/issue24539 for more information.
  1036. This is inappropriate for ``websockets`` for at least three reasons:
  1037. 1. The use case is to read data until EOF with self.reader.read(-1).
  1038. Since WebSocket is a TLV protocol, this never happens.
  1039. 2. It doesn't work on TLS connections. A falsy value must be
  1040. returned to have the same behavior on TLS and plain connections.
  1041. 3. The WebSocket protocol has its own closing handshake. Endpoints
  1042. close the TCP connection after sending a close frame.
  1043. As a consequence we revert to the previous, more useful behavior.
  1044. """
  1045. logger.debug("%s - event = eof_received()", self.side)
  1046. super().eof_received()
  1047. return False
  1048. def connection_lost(self, exc: Optional[Exception]) -> None:
  1049. """
  1050. 7.1.4. The WebSocket Connection is Closed.
  1051. """
  1052. logger.debug("%s - event = connection_lost(%s)", self.side, exc)
  1053. self.state = State.CLOSED
  1054. logger.debug("%s - state = CLOSED", self.side)
  1055. if not hasattr(self, "close_code"):
  1056. self.close_code = 1006
  1057. if not hasattr(self, "close_reason"):
  1058. self.close_reason = ""
  1059. logger.debug(
  1060. "%s x code = %d, reason = %s",
  1061. self.side,
  1062. self.close_code,
  1063. self.close_reason or "[no reason]",
  1064. )
  1065. self.abort_pings()
  1066. # If self.connection_lost_waiter isn't pending, that's a bug, because:
  1067. # - it's set only here in connection_lost() which is called only once;
  1068. # - it must never be canceled.
  1069. self.connection_lost_waiter.set_result(None)
  1070. super().connection_lost(exc)