codekingpro/portable-devtools
114k
1from __future__ import annotations
2
3import asyncio
4import collections
5import contextlib
6import logging
7import random
8import struct
9import sys
10import traceback
11import uuid
12from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Iterable, Mapping
13from types import TracebackType
14from typing import Any, Literal, cast, overload
15
16from ..exceptions import (
17 ConcurrencyError,
18 ConnectionClosed,
19 ConnectionClosedOK,
20 ProtocolError,
21)
22from ..frames import DATA_OPCODES, CloseCode, Frame, Opcode
23from ..http11 import Request, Response
24from ..protocol import CLOSED, OPEN, Event, Protocol, State
25from ..typing import BytesLike, Data, DataLike, LoggerLike, Subprotocol
26from .compatibility import (
27 TimeoutError,
28 aiter,
29 anext,
30 asyncio_timeout,
31 asyncio_timeout_at,
32)
33from .messages import Assembler
34
35
36__all__ = ["Connection"]
37
38
39class Connection(asyncio.Protocol):
40 """
41 :mod:`asyncio` implementation of a WebSocket connection.
42
43 :class:`Connection` provides APIs shared between WebSocket servers and
44 clients.
45
46 You shouldn't use it directly. Instead, use
47 :class:`~websockets.asyncio.client.ClientConnection` or
48 :class:`~websockets.asyncio.server.ServerConnection`.
49
50 """
51
52 def __init__(
53 self,
54 protocol: Protocol,
55 *,
56 ping_interval: float | None = 20,
57 ping_timeout: float | None = 20,
58 close_timeout: float | None = 10,
59 max_queue: int | None | tuple[int | None, int | None] = 16,
60 write_limit: int | tuple[int, int | None] = 2**15,
61 ) -> None:
62 self.protocol = protocol
63 self.ping_interval = ping_interval
64 self.ping_timeout = ping_timeout
65 self.close_timeout = close_timeout
66 if isinstance(max_queue, int) or max_queue is None:
67 self.max_queue_high, self.max_queue_low = max_queue, None
68 else:
69 self.max_queue_high, self.max_queue_low = max_queue
70 if isinstance(write_limit, int):
71 self.write_limit_high, self.write_limit_low = write_limit, None
72 else:
73 self.write_limit_high, self.write_limit_low = write_limit
74
75 # Inject reference to this instance in the protocol's logger.
76 self.protocol.logger = logging.LoggerAdapter(
77 self.protocol.logger,
78 {"websocket": self},
79 )
80
81 # Copy attributes from the protocol for convenience.
82 self.id: uuid.UUID = self.protocol.id
83 """Unique identifier of the connection. Useful in logs."""
84 self.logger: LoggerLike = self.protocol.logger
85 """Logger for this connection."""
86 self.debug = self.protocol.debug
87
88 # HTTP handshake request and response.
89 self.request: Request | None = None
90 """Opening handshake request."""
91 self.response: Response | None = None
92 """Opening handshake response."""
93
94 # Event loop running this connection.
95 self.loop = asyncio.get_running_loop()
96
97 # Assembler turning frames into messages and serializing reads.
98 self.recv_messages: Assembler # initialized in connection_made
99
100 # Deadline for the closing handshake.
101 self.close_deadline: float | None = None
102
103 # Whether we are busy sending a fragmented message.
104 self.send_in_progress: asyncio.Future[None] | None = None
105
106 # Mapping of ping IDs to pong waiters, in chronological order.
107 self.pending_pings: dict[bytes, tuple[asyncio.Future[float], float]] = {}
108
109 self.latency: float = 0.0
110 """
111 Latency of the connection, in seconds.
112
113 Latency is defined as the round-trip time of the connection. It is
114 measured by sending a Ping frame and waiting for a matching Pong frame.
115 Before the first measurement, :attr:`latency` is ``0.0``.
116
117 By default, websockets enables a :ref:`keepalive <keepalive>` mechanism
118 that sends Ping frames automatically at regular intervals. You can also
119 send Ping frames and measure latency with :meth:`ping`.
120 """
121
122 # Task that sends keepalive pings. None when ping_interval is None.
123 self.keepalive_task: asyncio.Task[None] | None = None
124
125 # Exception raised while reading from the connection, to be chained to
126 # ConnectionClosed in order to show why the TCP connection dropped.
127 self.recv_exc: BaseException | None = None
128
129 # Completed when the TCP connection is closed and the WebSocket
130 # connection state becomes CLOSED.
131 self.connection_lost_waiter: asyncio.Future[None] = self.loop.create_future()
132
133 # Adapted from asyncio.FlowControlMixin.
134 self.paused: bool = False
135 self.drain_waiters: collections.deque[asyncio.Future[None]] = (
136 collections.deque()
137 )
138
139 # Public attributes
140
141 @property
142 def local_address(self) -> Any:
143 """
144 Local address of the connection.
145
146 For IPv4 connections, this is a ``(host, port)`` tuple.
147
148 The format of the address depends on the address family.
149 See :meth:`~socket.socket.getsockname`.
150
151 """
152 return self.transport.get_extra_info("sockname")
153
154 @property
155 def remote_address(self) -> Any:
156 """
157 Remote address of the connection.
158
159 For IPv4 connections, this is a ``(host, port)`` tuple.
160
161 The format of the address depends on the address family.
162 See :meth:`~socket.socket.getpeername`.
163
164 """
165 return self.transport.get_extra_info("peername")
166
167 @property
168 def state(self) -> State:
169 """
170 State of the WebSocket connection, defined in :rfc:`6455`.
171
172 This attribute is provided for completeness. Typical applications
173 shouldn't check its value. Instead, they should call :meth:`~recv` or
174 :meth:`send` and handle :exc:`~websockets.exceptions.ConnectionClosed`
175 exceptions.
176
177 """
178 return self.protocol.state
179
180 @property
181 def subprotocol(self) -> Subprotocol | None:
182 """
183 Subprotocol negotiated during the opening handshake.
184
185 :obj:`None` if no subprotocol was negotiated.
186
187 """
188 return self.protocol.subprotocol
189
190 @property
191 def close_code(self) -> int | None:
192 """
193 State of the WebSocket connection, defined in :rfc:`6455`.
194
195 This attribute is provided for completeness. Typical applications
196 shouldn't check its value. Instead, they should inspect attributes
197 of :exc:`~websockets.exceptions.ConnectionClosed` exceptions.
198
199 """
200 return self.protocol.close_code
201
202 @property
203 def close_reason(self) -> str | None:
204 """
205 State of the WebSocket connection, defined in :rfc:`6455`.
206
207 This attribute is provided for completeness. Typical applications
208 shouldn't check its value. Instead, they should inspect attributes
209 of :exc:`~websockets.exceptions.ConnectionClosed` exceptions.
210
211 """
212 return self.protocol.close_reason
213
214 # Public methods
215
216 async def __aenter__(self) -> Connection:
217 return self
218
219 async def __aexit__(
220 self,
221 exc_type: type[BaseException] | None,
222 exc_value: BaseException | None,
223 traceback: TracebackType | None,
224 ) -> None:
225 if exc_type is None:
226 await self.close()
227 else:
228 await self.close(CloseCode.INTERNAL_ERROR)
229
230 async def __aiter__(self) -> AsyncIterator[Data]:
231 """
232 Iterate on incoming messages.
233
234 The iterator calls :meth:`recv` and yields messages asynchronously in an
235 infinite loop.
236
237 It exits when the connection is closed normally. It raises a
238 :exc:`~websockets.exceptions.ConnectionClosedError` exception after a
239 protocol error or a network failure.
240
241 """
242 try:
243 while True:
244 yield await self.recv()
245 except ConnectionClosedOK:
246 return
247
248 @overload
249 async def recv(self, decode: Literal[True]) -> str: ...
250
251 @overload
252 async def recv(self, decode: Literal[False]) -> bytes: ...
253
254 @overload
255 async def recv(self, decode: bool | None = None) -> Data: ...
256
257 async def recv(self, decode: bool | None = None) -> Data:
258 """
259 Receive the next message.
260
261 When the connection is closed, :meth:`recv` raises
262 :exc:`~websockets.exceptions.ConnectionClosed`. Specifically, it raises
263 :exc:`~websockets.exceptions.ConnectionClosedOK` after a normal closure
264 and :exc:`~websockets.exceptions.ConnectionClosedError` after a protocol
265 error or a network failure. This is how you detect the end of the
266 message stream.
267
268 Canceling :meth:`recv` is safe. There's no risk of losing data. The next
269 invocation of :meth:`recv` will return the next message.
270
271 This makes it possible to enforce a timeout by wrapping :meth:`recv` in
272 :func:`~asyncio.timeout` or :func:`~asyncio.wait_for`.
273
274 When the message is fragmented, :meth:`recv` waits until all fragments
275 are received, reassembles them, and returns the whole message.
276
277 Args:
278 decode: Set this flag to override the default behavior of returning
279 :class:`str` or :class:`bytes`. See below for details.
280
281 Returns:
282 A string (:class:`str`) for a Text_ frame or a bytestring
283 (:class:`bytes`) for a Binary_ frame.
284
285 .. _Text: https://datatracker.ietf.org/doc/html/rfc6455#section-5.6
286 .. _Binary: https://datatracker.ietf.org/doc/html/rfc6455#section-5.6
287
288 You may override this behavior with the ``decode`` argument:
289
290 * Set ``decode=False`` to disable UTF-8 decoding of Text_ frames and
291 return a bytestring (:class:`bytes`). This improves performance
292 when decoding isn't needed, for example if the message contains
293 JSON and you're using a JSON library that expects a bytestring.
294 * Set ``decode=True`` to force UTF-8 decoding of Binary_ frames and
295 return strings (:class:`str`). This may be useful for servers that
296 send binary frames instead of text frames.
297
298 Raises:
299 ConnectionClosed: When the connection is closed.
300 ConcurrencyError: If two coroutines call :meth:`recv` or
301 :meth:`recv_streaming` concurrently.
302
303 """
304 try:
305 return await self.recv_messages.get(decode)
306 except EOFError:
307 pass
308 # fallthrough
309 except ConcurrencyError:
310 raise ConcurrencyError(
311 "cannot call recv while another coroutine "
312 "is already running recv or recv_streaming"
313 ) from None
314 except UnicodeDecodeError as exc:
315 async with self.send_context():
316 self.protocol.fail(
317 CloseCode.INVALID_DATA,
318 f"{exc.reason} at position {exc.start}",
319 )
320 # fallthrough
321
322 # Wait for the protocol state to be CLOSED before accessing close_exc.
323 await asyncio.shield(self.connection_lost_waiter)
324 raise self.protocol.close_exc from self.recv_exc
325
326 @overload
327 def recv_streaming(self, decode: Literal[True]) -> AsyncIterator[str]: ...
328
329 @overload
330 def recv_streaming(self, decode: Literal[False]) -> AsyncIterator[bytes]: ...
331
332 @overload
333 def recv_streaming(self, decode: bool | None = None) -> AsyncIterator[Data]: ...
334
335 async def recv_streaming(self, decode: bool | None = None) -> AsyncIterator[Data]:
336 """
337 Receive the next message frame by frame.
338
339 This method is designed for receiving fragmented messages. It returns an
340 asynchronous iterator that yields each fragment as it is received. This
341 iterator must be fully consumed. Else, future calls to :meth:`recv` or
342 :meth:`recv_streaming` will raise
343 :exc:`~websockets.exceptions.ConcurrencyError`, making the connection
344 unusable.
345
346 :meth:`recv_streaming` raises the same exceptions as :meth:`recv`.
347
348 Canceling :meth:`recv_streaming` before receiving the first frame is
349 safe. Canceling it after receiving one or more frames leaves the
350 iterator in a partially consumed state, making the connection unusable.
351 Instead, you should close the connection with :meth:`close`.
352
353 Args:
354 decode: Set this flag to override the default behavior of returning
355 :class:`str` or :class:`bytes`. See below for details.
356
357 Returns:
358 An iterator of strings (:class:`str`) for a Text_ frame or
359 bytestrings (:class:`bytes`) for a Binary_ frame.
360
361 .. _Text: https://datatracker.ietf.org/doc/html/rfc6455#section-5.6
362 .. _Binary: https://datatracker.ietf.org/doc/html/rfc6455#section-5.6
363
364 You may override this behavior with the ``decode`` argument:
365
366 * Set ``decode=False`` to disable UTF-8 decoding of Text_ frames and
367 yield bytestrings (:class:`bytes`). This improves performance
368 when decoding isn't needed.
369 * Set ``decode=True`` to force UTF-8 decoding of Binary_ frames and
370 yield strings (:class:`str`). This may be useful for servers that
371 send binary frames instead of text frames.
372
373 Raises:
374 ConnectionClosed: When the connection is closed.
375 ConcurrencyError: If two coroutines call :meth:`recv` or
376 :meth:`recv_streaming` concurrently.
377
378 """
379 try:
380 async for frame in self.recv_messages.get_iter(decode):
381 yield frame
382 return
383 except EOFError:
384 pass
385 # fallthrough
386 except ConcurrencyError:
387 raise ConcurrencyError(
388 "cannot call recv_streaming while another coroutine "
389 "is already running recv or recv_streaming"
390 ) from None
391 except UnicodeDecodeError as exc:
392 async with self.send_context():
393 self.protocol.fail(
394 CloseCode.INVALID_DATA,
395 f"{exc.reason} at position {exc.start}",
396 )
397 # fallthrough
398
399 # Wait for the protocol state to be CLOSED before accessing close_exc.
400 await asyncio.shield(self.connection_lost_waiter)
401 raise self.protocol.close_exc from self.recv_exc
402
403 async def send(
404 self,
405 message: DataLike | Iterable[DataLike] | AsyncIterable[DataLike],
406 text: bool | None = None,
407 ) -> None:
408 """
409 Send a message.
410
411 A string (:class:`str`) is sent as a Text_ frame. A bytestring or
412 bytes-like object (:class:`bytes`, :class:`bytearray`, or
413 :class:`memoryview`) is sent as a Binary_ frame.
414
415 .. _Text: https://datatracker.ietf.org/doc/html/rfc6455#section-5.6
416 .. _Binary: https://datatracker.ietf.org/doc/html/rfc6455#section-5.6
417
418 You may override this behavior with the ``text`` argument:
419
420 * Set ``text=True`` to send an UTF-8 bytestring or bytes-like object
421 (:class:`bytes`, :class:`bytearray`, or :class:`memoryview`) in a
422 Text_ frame. This improves performance when the message is already
423 UTF-8 encoded, for example if the message contains JSON and you're
424 using a JSON library that produces a bytestring.
425 * Set ``text=False`` to send a string (:class:`str`) in a Binary_
426 frame. This may be useful for servers that expect binary frames
427 instead of text frames.
428
429 :meth:`send` also accepts an iterable or asynchronous iterable of
430 strings, bytestrings, or bytes-like objects to enable fragmentation_.
431 Each item is treated as a message fragment and sent in its own frame.
432 All items must be of the same type, or else :meth:`send` will raise a
433 :exc:`TypeError` and the connection will be closed.
434
435 .. _fragmentation: https://datatracker.ietf.org/doc/html/rfc6455#section-5.4
436
437 :meth:`send` rejects dict-like objects because this is often an error.
438 (If you really want to send the keys of a dict-like object as fragments,
439 call its :meth:`~dict.keys` method and pass the result to :meth:`send`.)
440
441 Canceling :meth:`send` is discouraged. Instead, you should close the
442 connection with :meth:`close`. Indeed, there are only two situations
443 where :meth:`send` may yield control to the event loop and then get
444 canceled; in both cases, :meth:`close` has the same effect and the
445 effect is more obvious:
446
447 1. The write buffer is full. If you don't want to wait until enough
448 data is sent, your only alternative is to close the connection.
449 :meth:`close` will likely time out then abort the TCP connection.
450 2. ``message`` is an asynchronous iterator that yields control.
451 Stopping in the middle of a fragmented message will cause a
452 protocol error and the connection will be closed.
453
454 When the connection is closed, :meth:`send` raises
455 :exc:`~websockets.exceptions.ConnectionClosed`. Specifically, it
456 raises :exc:`~websockets.exceptions.ConnectionClosedOK` after a normal
457 connection closure and
458 :exc:`~websockets.exceptions.ConnectionClosedError` after a protocol
459 error or a network failure.
460
461 Args:
462 message: Message to send.
463
464 Raises:
465 ConnectionClosed: When the connection is closed.
466 TypeError: If ``message`` doesn't have a supported type.
467
468 """
469 # While sending a fragmented message, prevent sending other messages
470 # until all fragments are sent.
471 while self.send_in_progress is not None:
472 await asyncio.shield(self.send_in_progress)
473
474 # Unfragmented message -- this case must be handled first because
475 # strings and bytes-like objects are iterable.
476
477 if isinstance(message, str):
478 async with self.send_context():
479 if text is False:
480 self.protocol.send_binary(message.encode())
481 else:
482 self.protocol.send_text(message.encode())
483
484 elif isinstance(message, BytesLike):
485 async with self.send_context():
486 if text is True:
487 self.protocol.send_text(message)
488 else:
489 self.protocol.send_binary(message)
490
491 # Catch a common mistake -- passing a dict to send().
492
493 elif isinstance(message, Mapping):
494 raise TypeError("data is a dict-like object")
495
496 # Fragmented message -- regular iterator.
497
498 elif isinstance(message, Iterable):
499 chunks = iter(message)
500 try:
501 chunk = next(chunks)
502 except StopIteration:
503 return
504
505 assert self.send_in_progress is None
506 self.send_in_progress = self.loop.create_future()
507 try:
508 # First fragment.
509 if isinstance(chunk, str):
510 async with self.send_context():
511 if text is False:
512 self.protocol.send_binary(chunk.encode(), fin=False)
513 else:
514 self.protocol.send_text(chunk.encode(), fin=False)
515 encode = True
516 elif isinstance(chunk, BytesLike):
517 async with self.send_context():
518 if text is True:
519 self.protocol.send_text(chunk, fin=False)
520 else:
521 self.protocol.send_binary(chunk, fin=False)
522 encode = False
523 else:
524 raise TypeError("iterable must contain bytes or str")
525
526 # Other fragments
527 for chunk in chunks:
528 if isinstance(chunk, str) and encode:
529 async with self.send_context():
530 self.protocol.send_continuation(chunk.encode(), fin=False)
531 elif isinstance(chunk, BytesLike) and not encode:
532 async with self.send_context():
533 self.protocol.send_continuation(chunk, fin=False)
534 else:
535 raise TypeError("iterable must contain uniform types")
536
537 # Final fragment.
538 async with self.send_context():
539 self.protocol.send_continuation(b"", fin=True)
540
541 except Exception:
542 # We're half-way through a fragmented message and we can't
543 # complete it. This makes the connection unusable.
544 async with self.send_context():
545 self.protocol.fail(
546 CloseCode.INTERNAL_ERROR,
547 "error in fragmented message",
548 )
549 raise
550
551 finally:
552 self.send_in_progress.set_result(None)
553 self.send_in_progress = None
554
555 # Fragmented message -- async iterator.
556
557 elif isinstance(message, AsyncIterable):
558 achunks = aiter(message)
559 try:
560 chunk = await anext(achunks)
561 except StopAsyncIteration:
562 return
563
564 assert self.send_in_progress is None
565 self.send_in_progress = self.loop.create_future()
566 try:
567 # First fragment.
568 if isinstance(chunk, str):
569 if text is False:
570 async with self.send_context():
571 self.protocol.send_binary(chunk.encode(), fin=False)
572 else:
573 async with self.send_context():
574 self.protocol.send_text(chunk.encode(), fin=False)
575 encode = True
576 elif isinstance(chunk, BytesLike):
577 if text is True:
578 async with self.send_context():
579 self.protocol.send_text(chunk, fin=False)
580 else:
581 async with self.send_context():
582 self.protocol.send_binary(chunk, fin=False)
583 encode = False
584 else:
585 raise TypeError("async iterable must contain bytes or str")
586
587 # Other fragments
588 async for chunk in achunks:
589 if isinstance(chunk, str) and encode:
590 async with self.send_context():
591 self.protocol.send_continuation(chunk.encode(), fin=False)
592 elif isinstance(chunk, BytesLike) and not encode:
593 async with self.send_context():
594 self.protocol.send_continuation(chunk, fin=False)
595 else:
596 raise TypeError("async iterable must contain uniform types")
597
598 # Final fragment.
599 async with self.send_context():
600 self.protocol.send_continuation(b"", fin=True)
601
602 except Exception:
603 # We're half-way through a fragmented message and we can't
604 # complete it. This makes the connection unusable.
605 async with self.send_context():
606 self.protocol.fail(
607 CloseCode.INTERNAL_ERROR,
608 "error in fragmented message",
609 )
610 raise
611
612 finally:
613 self.send_in_progress.set_result(None)
614 self.send_in_progress = None
615
616 else:
617 raise TypeError("data must be str, bytes, iterable, or async iterable")
618
619 async def close(
620 self,
621 code: CloseCode | int = CloseCode.NORMAL_CLOSURE,
622 reason: str = "",
623 ) -> None:
624 """
625 Perform the closing handshake.
626
627 :meth:`close` waits for the other end to complete the handshake and
628 for the TCP connection to terminate.
629
630 :meth:`close` is idempotent: it doesn't do anything once the
631 connection is closed.
632
633 Args:
634 code: WebSocket close code.
635 reason: WebSocket close reason.
636
637 """
638 try:
639 # The context manager takes care of waiting for the TCP connection
640 # to terminate after calling a method that sends a close frame.
641 async with self.send_context():
642 if self.send_in_progress is not None:
643 self.protocol.fail(
644 CloseCode.INTERNAL_ERROR,
645 "close during fragmented message",
646 )
647 else:
648 self.protocol.send_close(code, reason)
649 except ConnectionClosed:
650 # Ignore ConnectionClosed exceptions raised from send_context().
651 # They mean that the connection is closed, which was the goal.
652 pass
653
654 async def wait_closed(self) -> None:
655 """
656 Wait until the connection is closed.
657
658 :meth:`wait_closed` waits for the closing handshake to complete and for
659 the TCP connection to terminate.
660
661 """
662 await asyncio.shield(self.connection_lost_waiter)
663
664 async def ping(self, data: DataLike | None = None) -> Awaitable[float]:
665 """
666 Send a Ping_.
667
668 .. _Ping: https://datatracker.ietf.org/doc/html/rfc6455#section-5.5.2
669
670 A ping may serve as a keepalive or as a check that the remote endpoint
671 received all messages up to this point
672
673 Args:
674 data: Payload of the ping. A :class:`str` will be encoded to UTF-8.
675 If ``data`` is :obj:`None`, the payload is four random bytes.
676
677 Returns:
678 A future that will be completed when the corresponding pong is
679 received. You can ignore it if you don't intend to wait. The result
680 of the future is the latency of the connection in seconds.
681
682 ::
683
684 pong_received = await ws.ping()
685 # only if you want to wait for the corresponding pong
686 latency = await pong_received
687
688 Raises:
689 ConnectionClosed: When the connection is closed.
690 ConcurrencyError: If another ping was sent with the same data and
691 the corresponding pong wasn't received yet.
692
693 """
694 if isinstance(data, BytesLike):
695 data = bytes(data)
696 elif isinstance(data, str):
697 data = data.encode()
698 elif data is not None:
699 raise TypeError("data must be str or bytes-like")
700
701 async with self.send_context():
702 # Protect against duplicates if a payload is explicitly set.
703 if data in self.pending_pings:
704 raise ConcurrencyError("already waiting for a pong with the same data")
705
706 # Generate a unique random payload otherwise.
707 while data is None or data in self.pending_pings:
708 data = struct.pack("!I", random.getrandbits(32))
709
710 pong_received = self.loop.create_future()
711 ping_timestamp = self.loop.time()
712 # The event loop's default clock is time.monotonic(). Its resolution
713 # is a bit low on Windows (~16ms). This is improved in Python 3.13.
714 self.pending_pings[data] = (pong_received, ping_timestamp)
715 self.protocol.send_ping(data)
716 return pong_received
717
718 async def pong(self, data: DataLike = b"") -> None:
719 """
720 Send a Pong_.
721
722 .. _Pong: https://datatracker.ietf.org/doc/html/rfc6455#section-5.5.3
723
724 An unsolicited pong may serve as a unidirectional heartbeat.
725
726 Args:
727 data: Payload of the pong. A :class:`str` will be encoded to UTF-8.
728
729 Raises:
730 ConnectionClosed: When the connection is closed.
731
732 """
733 if isinstance(data, BytesLike):
734 data = bytes(data)
735 elif isinstance(data, str):
736 data = data.encode()
737 else:
738 raise TypeError("data must be str or bytes-like")
739
740 async with self.send_context():
741 self.protocol.send_pong(data)
742
743 # Private methods
744
745 def process_event(self, event: Event) -> None:
746 """
747 Process one incoming event.
748
749 This method is overridden in subclasses to handle the handshake.
750
751 """
752 assert isinstance(event, Frame)
753 if event.opcode in DATA_OPCODES:
754 self.recv_messages.put(event)
755
756 if event.opcode is Opcode.PONG:
757 self.acknowledge_pings(bytes(event.data))
758
759 def acknowledge_pings(self, data: bytes) -> None:
760 """
761 Acknowledge pings when receiving a pong.
762
763 """
764 # Ignore unsolicited pong.
765 if data not in self.pending_pings:
766 return
767
768 pong_timestamp = self.loop.time()
769
770 # Sending a pong for only the most recent ping is legal.
771 # Acknowledge all previous pings too in that case.
772 ping_id = None
773 ping_ids = []
774 for ping_id, (pong_received, ping_timestamp) in self.pending_pings.items():
775 ping_ids.append(ping_id)
776 latency = pong_timestamp - ping_timestamp
777 if not pong_received.done():
778 pong_received.set_result(latency)
779 if ping_id == data:
780 self.latency = latency
781 break
782 else:
783 raise AssertionError("solicited pong not found in pings")
784
785 # Remove acknowledged pings from self.pending_pings.
786 for ping_id in ping_ids:
787 del self.pending_pings[ping_id]
788
789 def terminate_pending_pings(self) -> None:
790 """
791 Raise ConnectionClosed in pending pings when the connection is closed.
792
793 """
794 assert self.protocol.state is CLOSED
795 exc = self.protocol.close_exc
796
797 for pong_received, _ping_timestamp in self.pending_pings.values():
798 if not pong_received.done():
799 pong_received.set_exception(exc)
800 # If the exception is never retrieved, it will be logged when ping
801 # is garbage-collected. This is confusing for users.
802 # Given that ping is done (with an exception), canceling it does
803 # nothing, but it prevents logging the exception.
804 pong_received.cancel()
805
806 self.pending_pings.clear()
807
808 async def keepalive(self) -> None:
809 """
810 Send a Ping frame and wait for a Pong frame at regular intervals.
811
812 """
813 assert self.ping_interval is not None
814 latency = 0.0
815 try:
816 while True:
817 # If self.ping_timeout > latency > self.ping_interval,
818 # pings will be sent immediately after receiving pongs.
819 # The period will be longer than self.ping_interval.
820 await asyncio.sleep(self.ping_interval - latency)
821
822 # This cannot raise ConnectionClosed when the connection is
823 # closing because ping(), via send_context(), waits for the
824 # connection to be closed before raising ConnectionClosed.
825 # However, connection_lost() cancels keepalive_task before
826 # it gets a chance to resume excuting.
827 pong_received = await self.ping()
828 if self.debug:
829 self.logger.debug("% sent keepalive ping")
830
831 if self.ping_timeout is not None:
832 try:
833 async with asyncio_timeout(self.ping_timeout):
834 # connection_lost cancels keepalive immediately
835 # after setting a ConnectionClosed exception on
836 # pong_received. A CancelledError is raised here,
837 # not a ConnectionClosed exception.
838 latency = await pong_received
839 if self.debug:
840 self.logger.debug("% received keepalive pong")
841 except asyncio.TimeoutError:
842 if self.debug:
843 self.logger.debug("- timed out waiting for keepalive pong")
844 async with self.send_context():
845 self.protocol.fail(
846 CloseCode.INTERNAL_ERROR,
847 "keepalive ping timeout",
848 )
849 raise AssertionError(
850 "send_context() should wait for connection_lost(), "
851 "which cancels keepalive()"
852 )
853 except Exception:
854 self.logger.error("keepalive ping failed", exc_info=True)
855
856 def start_keepalive(self) -> None:
857 """
858 Run :meth:`keepalive` in a task, unless keepalive is disabled.
859
860 """
861 if self.ping_interval is not None:
862 self.keepalive_task = self.loop.create_task(self.keepalive())
863
864 @contextlib.asynccontextmanager
865 async def send_context(
866 self,
867 *,
868 expected_state: State = OPEN, # CONNECTING during the opening handshake
869 ) -> AsyncIterator[None]:
870 """
871 Create a context for writing to the connection from user code.
872
873 On entry, :meth:`send_context` checks that the connection is open; on
874 exit, it writes outgoing data to the socket::
875
876 async with self.send_context():
877 self.protocol.send_text(message.encode())
878
879 When the connection isn't open on entry, when the connection is expected
880 to close on exit, or when an unexpected error happens, terminating the
881 connection, :meth:`send_context` waits until the connection is closed
882 then raises :exc:`~websockets.exceptions.ConnectionClosed`.
883
884 """
885 # Should we wait until the connection is closed?
886 wait_for_close = False
887 # Should we close the transport and raise ConnectionClosed?
888 raise_close_exc = False
889 # What exception should we chain ConnectionClosed to?
890 original_exc: BaseException | None = None
891
892 if self.protocol.state is expected_state:
893 # Let the caller interact with the protocol.
894 try:
895 yield
896 except (ProtocolError, ConcurrencyError):
897 # The protocol state wasn't changed. Exit immediately.
898 raise
899 except Exception as exc:
900 self.logger.error("unexpected internal error", exc_info=True)
901 # This branch should never run. It's a safety net in case of
902 # bugs. Since we don't know what happened, we will close the
903 # connection and raise the exception to the caller.
904 wait_for_close = False
905 raise_close_exc = True
906 original_exc = exc
907 else:
908 # Check if the connection is expected to close soon.
909 if self.protocol.close_expected():
910 wait_for_close = True
911 # Set the close deadline based on the close timeout.
912 # Since we tested earlier that protocol.state is OPEN
913 # (or CONNECTING), self.close_deadline is still None.
914 assert self.close_deadline is None
915 if self.close_timeout is not None:
916 self.close_deadline = self.loop.time() + self.close_timeout
917 # Write outgoing data to the socket with flow control.
918 try:
919 self.send_data()
920 await self.drain()
921 except Exception as exc:
922 if self.debug:
923 self.logger.debug(
924 "! error while sending data",
925 exc_info=True,
926 )
927 # While the only expected exception here is OSError,
928 # other exceptions would be treated identically.
929 wait_for_close = False
930 raise_close_exc = True
931 original_exc = exc
932
933 else: # self.protocol.state is not expected_state
934 # Minor layering violation: we assume that the connection
935 # will be closing soon if it isn't in the expected state.
936 wait_for_close = True
937 # Calculate close_deadline if it wasn't set yet.
938 if self.close_deadline is None:
939 if self.close_timeout is not None:
940 self.close_deadline = self.loop.time() + self.close_timeout
941 raise_close_exc = True
942
943 # If the connection is expected to close soon and the close timeout
944 # elapses, close the socket to terminate the connection.
945 if wait_for_close:
946 try:
947 async with asyncio_timeout_at(self.close_deadline):
948 await asyncio.shield(self.connection_lost_waiter)
949 except TimeoutError:
950 # There's no risk of overwriting another error because
951 # original_exc is never set when wait_for_close is True.
952 assert original_exc is None
953 original_exc = TimeoutError("timed out while closing connection")
954 # Set recv_exc before closing the transport in order to get
955 # proper exception reporting.
956 raise_close_exc = True
957 self.set_recv_exc(original_exc)
958
959 # If an error occurred, close the transport to terminate the connection and
960 # raise an exception.
961 if raise_close_exc:
962 self.transport.abort()
963 # Wait for the protocol state to be CLOSED before accessing close_exc.
964 await asyncio.shield(self.connection_lost_waiter)
965 raise self.protocol.close_exc from original_exc
966
967 def send_data(self) -> None:
968 """
969 Send outgoing data.
970
971 """
972 for data in self.protocol.data_to_send():
973 if data:
974 self.transport.write(data)
975 else:
976 # Half-close the TCP connection when possible i.e. no TLS.
977 if self.transport.can_write_eof():
978 if self.debug:
979 self.logger.debug("x half-closing TCP connection")
980 # write_eof() doesn't document which exceptions it raises.
981 # OSError is plausible. uvloop can raise RuntimeError here.
982 try:
983 self.transport.write_eof()
984 except Exception: # pragma: no cover
985 pass
986 # Else, close the TCP connection.
987 else: # pragma: no cover
988 if self.debug:
989 self.logger.debug("x closing TCP connection")
990 self.transport.close()
991
992 def set_recv_exc(self, exc: BaseException | None) -> None:
993 """
994 Set recv_exc, if not set yet.
995
996 This method must be called only from connection callbacks.
997
998 """
999 if self.recv_exc is None:
1000 self.recv_exc = exc
1001
1002 # asyncio.Protocol methods
1003
1004 # Connection callbacks
1005
1006 def connection_made(self, transport: asyncio.BaseTransport) -> None:
1007 transport = cast(asyncio.Transport, transport)
1008 self.recv_messages = Assembler(
1009 self.max_queue_high,
1010 self.max_queue_low,
1011 pause=transport.pause_reading,
1012 resume=transport.resume_reading,
1013 )
1014 transport.set_write_buffer_limits(
1015 self.write_limit_high,
1016 self.write_limit_low,
1017 )
1018 self.transport = transport
1019
1020 def connection_lost(self, exc: Exception | None) -> None:
1021 # Calling protocol.receive_eof() is safe because it's idempotent.
1022 # This guarantees that the protocol state becomes CLOSED.
1023 self.protocol.receive_eof()
1024 assert self.protocol.state is CLOSED
1025
1026 self.set_recv_exc(exc)
1027
1028 # Abort recv() and pending pings with a ConnectionClosed exception.
1029 self.recv_messages.close()
1030 self.terminate_pending_pings()
1031
1032 if self.keepalive_task is not None:
1033 self.keepalive_task.cancel()
1034
1035 # If self.connection_lost_waiter isn't pending, that's a bug, because:
1036 # - it's set only here in connection_lost() which is called only once;
1037 # - it must never be canceled.
1038 self.connection_lost_waiter.set_result(None)
1039
1040 # Adapted from asyncio.streams.FlowControlMixin
1041 if self.paused: # pragma: no cover
1042 self.paused = False
1043 for waiter in self.drain_waiters:
1044 if not waiter.done():
1045 if exc is None:
1046 waiter.set_result(None)
1047 else:
1048 waiter.set_exception(exc)
1049
1050 # Flow control callbacks
1051
1052 def pause_writing(self) -> None: # pragma: no cover
1053 # Adapted from asyncio.streams.FlowControlMixin
1054 assert not self.paused
1055 self.paused = True
1056
1057 def resume_writing(self) -> None: # pragma: no cover
1058 # Adapted from asyncio.streams.FlowControlMixin
1059 assert self.paused
1060 self.paused = False
1061 for waiter in self.drain_waiters:
1062 if not waiter.done():
1063 waiter.set_result(None)
1064
1065 async def drain(self) -> None: # pragma: no cover
1066 # We don't check if the connection is closed because we call drain()
1067 # immediately after write() and write() would fail in that case.
1068
1069 # Adapted from asyncio.streams.StreamWriter
1070 # Yield to the event loop so that connection_lost() may be called.
1071 if self.transport.is_closing():
1072 await asyncio.sleep(0)
1073
1074 # Adapted from asyncio.streams.FlowControlMixin
1075 if self.paused:
1076 waiter = self.loop.create_future()
1077 self.drain_waiters.append(waiter)
1078 try:
1079 await waiter
1080 finally:
1081 self.drain_waiters.remove(waiter)
1082
1083 # Streaming protocol callbacks
1084
1085 def data_received(self, data: bytes) -> None:
1086 # Feed incoming data to the protocol.
1087 self.protocol.receive_data(data)
1088
1089 # This isn't expected to raise an exception.
1090 events = self.protocol.events_received()
1091
1092 # Write outgoing data to the transport.
1093 try:
1094 self.send_data()
1095 except Exception as exc:
1096 if self.debug:
1097 self.logger.debug("! error while sending data", exc_info=True)
1098 self.set_recv_exc(exc)
1099
1100 # If needed, set the close deadline based on the close timeout.
1101 if self.protocol.close_expected():
1102 if self.close_deadline is None:
1103 if self.close_timeout is not None:
1104 self.close_deadline = self.loop.time() + self.close_timeout
1105
1106 # If self.send_data raised an exception, then events are lost.
1107 # Given that automatic responses write small amounts of data,
1108 # this should be uncommon, so we don't handle the edge case.
1109
1110 for event in events:
1111 # This isn't expected to raise an exception.
1112 self.process_event(event)
1113
1114 def eof_received(self) -> None:
1115 # Feed the end of the data stream to the protocol.
1116 self.protocol.receive_eof()
1117
1118 # This isn't expected to raise an exception.
1119 events = self.protocol.events_received()
1120
1121 # There is no error handling because send_data() can only write
1122 # the end of the data stream and it handles errors by itself.
1123 self.send_data()
1124
1125 # This code path is triggered when receiving an HTTP response
1126 # without a Content-Length header. This is the only case where
1127 # reading until EOF generates an event; all other events have
1128 # a known length. Ignore for coverage measurement because tests
1129 # are in test_client.py rather than test_connection.py.
1130 for event in events: # pragma: no cover
1131 # This isn't expected to raise an exception.
1132 self.process_event(event)
1133
1134 # The WebSocket protocol has its own closing handshake: endpoints close
1135 # the TCP or TLS connection after sending and receiving a close frame.
1136 # As a consequence, they never need to write after receiving EOF, so
1137 # there's no reason to keep the transport open by returning True.
1138 # Besides, that doesn't work on TLS connections.
1139
1140
1141# broadcast() is defined in the connection module even though it's primarily
1142# used by servers and documented in the server module because it works with
1143# client connections too and because it's easier to test together with the
1144# Connection class.
1145
1146
1147def broadcast(
1148 connections: Iterable[Connection],
1149 message: DataLike,
1150 raise_exceptions: bool = False,
1151) -> None:
1152 """
1153 Broadcast a message to several WebSocket connections.
1154
1155 A string (:class:`str`) is sent as a Text_ frame. A bytestring or bytes-like
1156 object (:class:`bytes`, :class:`bytearray`, or :class:`memoryview`) is sent
1157 as a Binary_ frame.
1158
1159 .. _Text: https://datatracker.ietf.org/doc/html/rfc6455#section-5.6
1160 .. _Binary: https://datatracker.ietf.org/doc/html/rfc6455#section-5.6
1161
1162 :func:`broadcast` pushes the message synchronously to all connections even
1163 if their write buffers are overflowing. There's no backpressure.
1164
1165 If you broadcast messages faster than a connection can handle them, messages
1166 will pile up in its write buffer until the connection times out. Keep
1167 ``ping_interval`` and ``ping_timeout`` low to prevent excessive memory usage
1168 from slow connections.
1169
1170 Unlike :meth:`~websockets.asyncio.connection.Connection.send`,
1171 :func:`broadcast` doesn't support sending fragmented messages. Indeed,
1172 fragmentation is useful for sending large messages without buffering them in
1173 memory, while :func:`broadcast` buffers one copy per connection as fast as
1174 possible.
1175
1176 :func:`broadcast` skips connections that aren't open in order to avoid
1177 errors on connections where the closing handshake is in progress.
1178
1179 :func:`broadcast` ignores failures to write the message on some connections.
1180 It continues writing to other connections. On Python 3.11 and above, you may
1181 set ``raise_exceptions`` to :obj:`True` to record failures and raise all
1182 exceptions in a :pep:`654` :exc:`ExceptionGroup`.
1183
1184 While :func:`broadcast` makes more sense for servers, it works identically
1185 with clients, if you have a use case for opening connections to many servers
1186 and broadcasting a message to them.
1187
1188 Args:
1189 websockets: WebSocket connections to which the message will be sent.
1190 message: Message to send.
1191 raise_exceptions: Whether to raise an exception in case of failures.
1192
1193 Raises:
1194 TypeError: If ``message`` doesn't have a supported type.
1195
1196 """
1197 if isinstance(message, str):
1198 send_method = "send_text"
1199 message = message.encode()
1200 elif isinstance(message, BytesLike):