codekingpro/portable-devtools
114k
1from __future__ import annotations
2
3import contextlib
4import logging
5import random
6import socket
7import struct
8import threading
9import time
10import uuid
11from collections.abc import Iterable, Iterator, Mapping
12from types import TracebackType
13from typing import Any, Literal, overload
14
15from ..exceptions import (
16 ConcurrencyError,
17 ConnectionClosed,
18 ConnectionClosedOK,
19 ProtocolError,
20)
21from ..frames import DATA_OPCODES, CloseCode, Frame, Opcode
22from ..http11 import Request, Response
23from ..protocol import CLOSED, OPEN, Event, Protocol, State
24from ..typing import BytesLike, Data, DataLike, LoggerLike, Subprotocol
25from .messages import Assembler
26from .utils import Deadline
27
28
29__all__ = ["Connection"]
30
31
32class Connection:
33 """
34 :mod:`threading` implementation of a WebSocket connection.
35
36 :class:`Connection` provides APIs shared between WebSocket servers and
37 clients.
38
39 You shouldn't use it directly. Instead, use
40 :class:`~websockets.sync.client.ClientConnection` or
41 :class:`~websockets.sync.server.ServerConnection`.
42
43 """
44
45 recv_bufsize = 65536
46
47 def __init__(
48 self,
49 socket: socket.socket,
50 protocol: Protocol,
51 *,
52 ping_interval: float | None = 20,
53 ping_timeout: float | None = 20,
54 close_timeout: float | None = 10,
55 max_queue: int | None | tuple[int | None, int | None] = 16,
56 ) -> None:
57 self.socket = socket
58 self.protocol = protocol
59 self.ping_interval = ping_interval
60 self.ping_timeout = ping_timeout
61 self.close_timeout = close_timeout
62 if isinstance(max_queue, int) or max_queue is None:
63 max_queue_high, max_queue_low = max_queue, None
64 else:
65 max_queue_high, max_queue_low = max_queue
66
67 # Inject reference to this instance in the protocol's logger.
68 self.protocol.logger = logging.LoggerAdapter(
69 self.protocol.logger,
70 {"websocket": self},
71 )
72
73 # Copy attributes from the protocol for convenience.
74 self.id: uuid.UUID = self.protocol.id
75 """Unique identifier of the connection. Useful in logs."""
76 self.logger: LoggerLike = self.protocol.logger
77 """Logger for this connection."""
78 self.debug = self.protocol.debug
79
80 # HTTP handshake request and response.
81 self.request: Request | None = None
82 """Opening handshake request."""
83 self.response: Response | None = None
84 """Opening handshake response."""
85
86 # Mutex serializing interactions with the protocol.
87 self.protocol_mutex = threading.Lock()
88
89 # Lock stopping reads when the assembler buffer is full.
90 self.recv_flow_control = threading.Lock()
91
92 # Assembler turning frames into messages and serializing reads.
93 self.recv_messages = Assembler(
94 max_queue_high,
95 max_queue_low,
96 pause=self.recv_flow_control.acquire,
97 resume=self.recv_flow_control.release,
98 )
99
100 # Deadline for the closing handshake.
101 self.close_deadline: Deadline | None = None
102
103 # Whether we are busy sending a fragmented message.
104 self.send_in_progress = False
105
106 # Mapping of ping IDs to pong waiters, in chronological order.
107 self.pending_pings: dict[bytes, tuple[threading.Event, float, bool]] = {}
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 # Thread that sends keepalive pings. None when ping_interval is None.
123 self.keepalive_thread: threading.Thread | 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 # Receiving events from the socket. This thread is marked as daemon to
130 # allow creating a connection in a non-daemon thread and using it in a
131 # daemon thread. This mustn't prevent the interpreter from exiting.
132 self.recv_events_thread = threading.Thread(
133 target=self.recv_events,
134 daemon=True,
135 )
136
137 # Start recv_events only after all attributes are initialized.
138 self.recv_events_thread.start()
139
140 # Public attributes
141
142 @property
143 def local_address(self) -> Any:
144 """
145 Local address of the connection.
146
147 For IPv4 connections, this is a ``(host, port)`` tuple.
148
149 The format of the address depends on the address family.
150 See :meth:`~socket.socket.getsockname`.
151
152 """
153 return self.socket.getsockname()
154
155 @property
156 def remote_address(self) -> Any:
157 """
158 Remote address of the connection.
159
160 For IPv4 connections, this is a ``(host, port)`` tuple.
161
162 The format of the address depends on the address family.
163 See :meth:`~socket.socket.getpeername`.
164
165 """
166 return self.socket.getpeername()
167
168 @property
169 def state(self) -> State:
170 """
171 State of the WebSocket connection, defined in :rfc:`6455`.
172
173 This attribute is provided for completeness. Typical applications
174 shouldn't check its value. Instead, they should call :meth:`~recv` or
175 :meth:`send` and handle :exc:`~websockets.exceptions.ConnectionClosed`
176 exceptions.
177
178 """
179 return self.protocol.state
180
181 @property
182 def subprotocol(self) -> Subprotocol | None:
183 """
184 Subprotocol negotiated during the opening handshake.
185
186 :obj:`None` if no subprotocol was negotiated.
187
188 """
189 return self.protocol.subprotocol
190
191 @property
192 def close_code(self) -> int | None:
193 """
194 State of the WebSocket connection, defined in :rfc:`6455`.
195
196 This attribute is provided for completeness. Typical applications
197 shouldn't check its value. Instead, they should inspect attributes
198 of :exc:`~websockets.exceptions.ConnectionClosed` exceptions.
199
200 """
201 return self.protocol.close_code
202
203 @property
204 def close_reason(self) -> str | None:
205 """
206 State of the WebSocket connection, defined in :rfc:`6455`.
207
208 This attribute is provided for completeness. Typical applications
209 shouldn't check its value. Instead, they should inspect attributes
210 of :exc:`~websockets.exceptions.ConnectionClosed` exceptions.
211
212 """
213 return self.protocol.close_reason
214
215 # Public methods
216
217 def __enter__(self) -> Connection:
218 return self
219
220 def __exit__(
221 self,
222 exc_type: type[BaseException] | None,
223 exc_value: BaseException | None,
224 traceback: TracebackType | None,
225 ) -> None:
226 if exc_type is None:
227 self.close()
228 else:
229 self.close(CloseCode.INTERNAL_ERROR)
230
231 def __iter__(self) -> Iterator[Data]:
232 """
233 Iterate on incoming messages.
234
235 The iterator calls :meth:`recv` and yields messages in an 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 self.recv()
245 except ConnectionClosedOK:
246 return
247
248 # This overload structure is required to avoid the error:
249 # "parameter without a default follows parameter with a default"
250
251 @overload
252 def recv(self, timeout: float | None, decode: Literal[True]) -> str: ...
253
254 @overload
255 def recv(self, timeout: float | None, decode: Literal[False]) -> bytes: ...
256
257 @overload
258 def recv(self, timeout: float | None = None, *, decode: Literal[True]) -> str: ...
259
260 @overload
261 def recv(
262 self, timeout: float | None = None, *, decode: Literal[False]
263 ) -> bytes: ...
264
265 @overload
266 def recv(
267 self, timeout: float | None = None, decode: bool | None = None
268 ) -> Data: ...
269
270 def recv(self, timeout: float | None = None, decode: bool | None = None) -> Data:
271 """
272 Receive the next message.
273
274 When the connection is closed, :meth:`recv` raises
275 :exc:`~websockets.exceptions.ConnectionClosed`. Specifically, it raises
276 :exc:`~websockets.exceptions.ConnectionClosedOK` after a normal closure
277 and :exc:`~websockets.exceptions.ConnectionClosedError` after a protocol
278 error or a network failure. This is how you detect the end of the
279 message stream.
280
281 If ``timeout`` is :obj:`None`, block until a message is received. If
282 ``timeout`` is set, wait up to ``timeout`` seconds for a message to be
283 received and return it, else raise :exc:`TimeoutError`. If ``timeout``
284 is ``0`` or negative, check if a message has been received already and
285 return it, else raise :exc:`TimeoutError`.
286
287 When the message is fragmented, :meth:`recv` waits until all fragments
288 are received, reassembles them, and returns the whole message.
289
290 Args:
291 timeout: Timeout for receiving a message in seconds.
292 decode: Set this flag to override the default behavior of returning
293 :class:`str` or :class:`bytes`. See below for details.
294
295 Returns:
296 A string (:class:`str`) for a Text_ frame or a bytestring
297 (:class:`bytes`) for a Binary_ frame.
298
299 .. _Text: https://datatracker.ietf.org/doc/html/rfc6455#section-5.6
300 .. _Binary: https://datatracker.ietf.org/doc/html/rfc6455#section-5.6
301
302 You may override this behavior with the ``decode`` argument:
303
304 * Set ``decode=False`` to disable UTF-8 decoding of Text_ frames and
305 return a bytestring (:class:`bytes`). This improves performance
306 when decoding isn't needed, for example if the message contains
307 JSON and you're using a JSON library that expects a bytestring.
308 * Set ``decode=True`` to force UTF-8 decoding of Binary_ frames and
309 return strings (:class:`str`). This may be useful for servers that
310 send binary frames instead of text frames.
311
312 Raises:
313 ConnectionClosed: When the connection is closed.
314 ConcurrencyError: If two threads call :meth:`recv` or
315 :meth:`recv_streaming` concurrently.
316
317 """
318 try:
319 return self.recv_messages.get(timeout, decode)
320 except EOFError:
321 pass
322 # fallthrough
323 except ConcurrencyError:
324 raise ConcurrencyError(
325 "cannot call recv while another thread "
326 "is already running recv or recv_streaming"
327 ) from None
328 except UnicodeDecodeError as exc:
329 with self.send_context():
330 self.protocol.fail(
331 CloseCode.INVALID_DATA,
332 f"{exc.reason} at position {exc.start}",
333 )
334 # fallthrough
335
336 # Wait for the protocol state to be CLOSED before accessing close_exc.
337 self.recv_events_thread.join()
338 raise self.protocol.close_exc from self.recv_exc
339
340 @overload
341 def recv_streaming(self, decode: Literal[True]) -> Iterator[str]: ...
342
343 @overload
344 def recv_streaming(self, decode: Literal[False]) -> Iterator[bytes]: ...
345
346 @overload
347 def recv_streaming(self, decode: bool | None = None) -> Iterator[Data]: ...
348
349 def recv_streaming(self, decode: bool | None = None) -> Iterator[Data]:
350 """
351 Receive the next message frame by frame.
352
353 This method is designed for receiving fragmented messages. It returns an
354 iterator that yields each fragment as it is received. This iterator must
355 be fully consumed. Else, future calls to :meth:`recv` or
356 :meth:`recv_streaming` will raise
357 :exc:`~websockets.exceptions.ConcurrencyError`, making the connection
358 unusable.
359
360 :meth:`recv_streaming` raises the same exceptions as :meth:`recv`.
361
362 Args:
363 decode: Set this flag to override the default behavior of returning
364 :class:`str` or :class:`bytes`. See below for details.
365
366 Returns:
367 An iterator of strings (:class:`str`) for a Text_ frame or
368 bytestrings (:class:`bytes`) for a Binary_ frame.
369
370 .. _Text: https://datatracker.ietf.org/doc/html/rfc6455#section-5.6
371 .. _Binary: https://datatracker.ietf.org/doc/html/rfc6455#section-5.6
372
373 You may override this behavior with the ``decode`` argument:
374
375 * Set ``decode=False`` to disable UTF-8 decoding of Text_ frames and
376 yield bytestrings (:class:`bytes`). This improves performance
377 when decoding isn't needed.
378 * Set ``decode=True`` to force UTF-8 decoding of Binary_ frames and
379 yield strings (:class:`str`). This may be useful for servers that
380 send binary frames instead of text frames.
381
382 Raises:
383 ConnectionClosed: When the connection is closed.
384 ConcurrencyError: If two threads call :meth:`recv` or
385 :meth:`recv_streaming` concurrently.
386
387 """
388 try:
389 yield from self.recv_messages.get_iter(decode)
390 return
391 except EOFError:
392 pass
393 # fallthrough
394 except ConcurrencyError:
395 raise ConcurrencyError(
396 "cannot call recv_streaming while another thread "
397 "is already running recv or recv_streaming"
398 ) from None
399 except UnicodeDecodeError as exc:
400 with self.send_context():
401 self.protocol.fail(
402 CloseCode.INVALID_DATA,
403 f"{exc.reason} at position {exc.start}",
404 )
405 # fallthrough
406
407 # Wait for the protocol state to be CLOSED before accessing close_exc.
408 self.recv_events_thread.join()
409 raise self.protocol.close_exc from self.recv_exc
410
411 def send(
412 self,
413 message: DataLike | Iterable[DataLike],
414 text: bool | None = None,
415 ) -> None:
416 """
417 Send a message.
418
419 A string (:class:`str`) is sent as a Text_ frame. A bytestring or
420 bytes-like object (:class:`bytes`, :class:`bytearray`, or
421 :class:`memoryview`) is sent as a Binary_ frame.
422
423 .. _Text: https://datatracker.ietf.org/doc/html/rfc6455#section-5.6
424 .. _Binary: https://datatracker.ietf.org/doc/html/rfc6455#section-5.6
425
426 You may override this behavior with the ``text`` argument:
427
428 * Set ``text=True`` to send an UTF-8 bytestring or bytes-like object
429 (:class:`bytes`, :class:`bytearray`, or :class:`memoryview`) in a
430 Text_ frame. This improves performance when the message is already
431 UTF-8 encoded, for example if the message contains JSON and you're
432 using a JSON library that produces a bytestring.
433 * Set ``text=False`` to send a string (:class:`str`) in a Binary_
434 frame. This may be useful for servers that expect binary frames
435 instead of text frames.
436
437 :meth:`send` also accepts an iterable of strings, bytestrings, or
438 bytes-like objects to enable fragmentation_. Each item is treated as a
439 message fragment and sent in its own frame. All items must be of the
440 same type, or else :meth:`send` will raise a :exc:`TypeError` and the
441 connection will be closed.
442
443 .. _fragmentation: https://datatracker.ietf.org/doc/html/rfc6455#section-5.4
444
445 :meth:`send` rejects dict-like objects because this is often an error.
446 (If you really want to send the keys of a dict-like object as fragments,
447 call its :meth:`~dict.keys` method and pass the result to :meth:`send`.)
448
449 When the connection is closed, :meth:`send` raises
450 :exc:`~websockets.exceptions.ConnectionClosed`. Specifically, it
451 raises :exc:`~websockets.exceptions.ConnectionClosedOK` after a normal
452 connection closure and
453 :exc:`~websockets.exceptions.ConnectionClosedError` after a protocol
454 error or a network failure.
455
456 Args:
457 message: Message to send.
458
459 Raises:
460 ConnectionClosed: When the connection is closed.
461 ConcurrencyError: If the connection is sending a fragmented message.
462 TypeError: If ``message`` doesn't have a supported type.
463
464 """
465 # Unfragmented message -- this case must be handled first because
466 # strings and bytes-like objects are iterable.
467
468 if isinstance(message, str):
469 with self.send_context():
470 if self.send_in_progress:
471 raise ConcurrencyError(
472 "cannot call send while another thread is already running send"
473 )
474 if text is False:
475 self.protocol.send_binary(message.encode())
476 else:
477 self.protocol.send_text(message.encode())
478
479 elif isinstance(message, BytesLike):
480 with self.send_context():
481 if self.send_in_progress:
482 raise ConcurrencyError(
483 "cannot call send while another thread is already running send"
484 )
485 if text is True:
486 self.protocol.send_text(message)
487 else:
488 self.protocol.send_binary(message)
489
490 # Catch a common mistake -- passing a dict to send().
491
492 elif isinstance(message, Mapping):
493 raise TypeError("data is a dict-like object")
494
495 # Fragmented message -- regular iterator.
496
497 elif isinstance(message, Iterable):
498 chunks = iter(message)
499 try:
500 chunk = next(chunks)
501 except StopIteration:
502 return
503
504 try:
505 # First fragment.
506 if isinstance(chunk, str):
507 with self.send_context():
508 if self.send_in_progress:
509 raise ConcurrencyError(
510 "cannot call send while another thread "
511 "is already running send"
512 )
513 self.send_in_progress = True
514 if text is False:
515 self.protocol.send_binary(chunk.encode(), fin=False)
516 else:
517 self.protocol.send_text(chunk.encode(), fin=False)
518 encode = True
519 elif isinstance(chunk, BytesLike):
520 with self.send_context():
521 if self.send_in_progress:
522 raise ConcurrencyError(
523 "cannot call send while another thread "
524 "is already running send"
525 )
526 self.send_in_progress = True
527 if text is True:
528 self.protocol.send_text(chunk, fin=False)
529 else:
530 self.protocol.send_binary(chunk, fin=False)
531 encode = False
532 else:
533 raise TypeError("iterable must contain bytes or str")
534
535 # Other fragments
536 for chunk in chunks:
537 if isinstance(chunk, str) and encode:
538 with self.send_context():
539 assert self.send_in_progress
540 self.protocol.send_continuation(chunk.encode(), fin=False)
541 elif isinstance(chunk, BytesLike) and not encode:
542 with self.send_context():
543 assert self.send_in_progress
544 self.protocol.send_continuation(chunk, fin=False)
545 else:
546 raise TypeError("iterable must contain uniform types")
547
548 # Final fragment.
549 with self.send_context():
550 self.protocol.send_continuation(b"", fin=True)
551 self.send_in_progress = False
552
553 except ConcurrencyError:
554 # We didn't start sending a fragmented message.
555 # The connection is still usable.
556 raise
557
558 except Exception:
559 # We're half-way through a fragmented message and we can't
560 # complete it. This makes the connection unusable.
561 with self.send_context():
562 self.protocol.fail(
563 CloseCode.INTERNAL_ERROR,
564 "error in fragmented message",
565 )
566 raise
567
568 else:
569 raise TypeError("data must be str, bytes, or iterable")
570
571 def close(
572 self,
573 code: CloseCode | int = CloseCode.NORMAL_CLOSURE,
574 reason: str = "",
575 ) -> None:
576 """
577 Perform the closing handshake.
578
579 :meth:`close` waits for the other end to complete the handshake and
580 for the TCP connection to terminate.
581
582 :meth:`close` is idempotent: it doesn't do anything once the
583 connection is closed.
584
585 Args:
586 code: WebSocket close code.
587 reason: WebSocket close reason.
588
589 """
590 try:
591 # The context manager takes care of waiting for the TCP connection
592 # to terminate after calling a method that sends a close frame.
593 with self.send_context():
594 if self.send_in_progress:
595 self.protocol.fail(
596 CloseCode.INTERNAL_ERROR,
597 "close during fragmented message",
598 )
599 else:
600 self.protocol.send_close(code, reason)
601 except ConnectionClosed:
602 # Ignore ConnectionClosed exceptions raised from send_context().
603 # They mean that the connection is closed, which was the goal.
604 pass
605
606 def ping(
607 self,
608 data: DataLike | None = None,
609 ack_on_close: bool = False,
610 ) -> threading.Event:
611 """
612 Send a Ping_.
613
614 .. _Ping: https://datatracker.ietf.org/doc/html/rfc6455#section-5.5.2
615
616 A ping may serve as a keepalive or as a check that the remote endpoint
617 received all messages up to this point
618
619 Args:
620 data: Payload of the ping. A :class:`str` will be encoded to UTF-8.
621 If ``data`` is :obj:`None`, the payload is four random bytes.
622 ack_on_close: when this option is :obj:`True`, the event will also
623 be set when the connection is closed. While this avoids getting
624 stuck waiting for a pong that will never arrive, it requires
625 checking that the state of the connection is still ``OPEN`` to
626 confirm that a pong was received, rather than the connection
627 being closed.
628
629 Returns:
630 An event that will be set when the corresponding pong is received.
631 You can ignore it if you don't intend to wait.
632
633 ::
634
635 pong_received = ws.ping()
636 # only if you want to wait for the corresponding pong
637 pong_received.wait()
638
639 Raises:
640 ConnectionClosed: When the connection is closed.
641 ConcurrencyError: If another ping was sent with the same data and
642 the corresponding pong wasn't received yet.
643
644 """
645 if isinstance(data, BytesLike):
646 data = bytes(data)
647 elif isinstance(data, str):
648 data = data.encode()
649 elif data is not None:
650 raise TypeError("data must be str or bytes-like")
651
652 with self.send_context():
653 # Protect against duplicates if a payload is explicitly set.
654 if data in self.pending_pings:
655 raise ConcurrencyError("already waiting for a pong with the same data")
656
657 # Generate a unique random payload otherwise.
658 while data is None or data in self.pending_pings:
659 data = struct.pack("!I", random.getrandbits(32))
660
661 pong_received = threading.Event()
662 ping_timestamp = time.monotonic()
663 self.pending_pings[data] = (pong_received, ping_timestamp, ack_on_close)
664 self.protocol.send_ping(data)
665 return pong_received
666
667 def pong(self, data: DataLike = b"") -> None:
668 """
669 Send a Pong_.
670
671 .. _Pong: https://datatracker.ietf.org/doc/html/rfc6455#section-5.5.3
672
673 An unsolicited pong may serve as a unidirectional heartbeat.
674
675 Args:
676 data: Payload of the pong. A :class:`str` will be encoded to UTF-8.
677
678 Raises:
679 ConnectionClosed: When the connection is closed.
680
681 """
682 if isinstance(data, BytesLike):
683 data = bytes(data)
684 elif isinstance(data, str):
685 data = data.encode()
686 else:
687 raise TypeError("data must be str or bytes-like")
688
689 with self.send_context():
690 self.protocol.send_pong(data)
691
692 # Private methods
693
694 def process_event(self, event: Event) -> None:
695 """
696 Process one incoming event.
697
698 This method is overridden in subclasses to handle the handshake.
699
700 """
701 assert isinstance(event, Frame)
702 if event.opcode in DATA_OPCODES:
703 self.recv_messages.put(event)
704
705 if event.opcode is Opcode.PONG:
706 self.acknowledge_pings(bytes(event.data))
707
708 def acknowledge_pings(self, data: bytes) -> None:
709 """
710 Acknowledge pings when receiving a pong.
711
712 """
713 with self.protocol_mutex:
714 # Ignore unsolicited pong.
715 if data not in self.pending_pings:
716 return
717
718 pong_timestamp = time.monotonic()
719
720 # Sending a pong for only the most recent ping is legal.
721 # Acknowledge all previous pings too in that case.
722 ping_id = None
723 ping_ids = []
724 for ping_id, (
725 pong_received,
726 ping_timestamp,
727 _ack_on_close,
728 ) in self.pending_pings.items():
729 ping_ids.append(ping_id)
730 pong_received.set()
731 if ping_id == data:
732 self.latency = pong_timestamp - ping_timestamp
733 break
734 else:
735 raise AssertionError("solicited pong not found in pings")
736
737 # Remove acknowledged pings from self.pending_pings.
738 for ping_id in ping_ids:
739 del self.pending_pings[ping_id]
740
741 def terminate_pending_pings(self) -> None:
742 """
743 Acknowledge pending pings when the connection is closed.
744
745 """
746 assert self.protocol_mutex.locked()
747 assert self.protocol.state is CLOSED
748
749 for pong_received, _ping_timestamp, ack_on_close in self.pending_pings.values():
750 if ack_on_close:
751 pong_received.set()
752
753 self.pending_pings.clear()
754
755 def keepalive(self) -> None:
756 """
757 Send a Ping frame and wait for a Pong frame at regular intervals.
758
759 """
760 assert self.ping_interval is not None
761 try:
762 while True:
763 # If self.ping_timeout > self.latency > self.ping_interval,
764 # pings will be sent immediately after receiving pongs.
765 # The period will be longer than self.ping_interval.
766 self.recv_events_thread.join(self.ping_interval - self.latency)
767 if not self.recv_events_thread.is_alive():
768 break
769
770 try:
771 pong_received = self.ping(ack_on_close=True)
772 except ConnectionClosed:
773 break
774 if self.debug:
775 self.logger.debug("% sent keepalive ping")
776
777 if self.ping_timeout is not None:
778 if pong_received.wait(self.ping_timeout):
779 if self.debug:
780 self.logger.debug("% received keepalive pong")
781 else:
782 if self.debug:
783 self.logger.debug("- timed out waiting for keepalive pong")
784 with self.send_context():
785 self.protocol.fail(
786 CloseCode.INTERNAL_ERROR,
787 "keepalive ping timeout",
788 )
789 break
790 except Exception:
791 self.logger.error("keepalive ping failed", exc_info=True)
792
793 def start_keepalive(self) -> None:
794 """
795 Run :meth:`keepalive` in a thread, unless keepalive is disabled.
796
797 """
798 if self.ping_interval is not None:
799 # This thread is marked as daemon like self.recv_events_thread.
800 self.keepalive_thread = threading.Thread(
801 target=self.keepalive,
802 daemon=True,
803 )
804 self.keepalive_thread.start()
805
806 def recv_events(self) -> None:
807 """
808 Read incoming data from the socket and process events.
809
810 Run this method in a thread as long as the connection is alive.
811
812 ``recv_events()`` exits immediately when ``self.socket`` is closed.
813
814 """
815 try:
816 while True:
817 try:
818 # If the assembler buffer is full, block until it drains.
819 with self.recv_flow_control:
820 pass
821 if self.close_deadline is not None:
822 self.socket.settimeout(self.close_deadline.timeout())
823 data = self.socket.recv(self.recv_bufsize)
824 except Exception as exc:
825 if self.debug:
826 self.logger.debug(
827 "! error while receiving data",
828 exc_info=True,
829 )
830 # When the closing handshake is initiated by our side,
831 # recv() may block until send_context() closes the socket.
832 # In that case, send_context() already set recv_exc.
833 # Calling set_recv_exc() avoids overwriting it.
834 with self.protocol_mutex:
835 self.set_recv_exc(exc)
836 break
837
838 if data == b"":
839 break
840
841 # Acquire the connection lock.
842 with self.protocol_mutex:
843 # Feed incoming data to the protocol.
844 self.protocol.receive_data(data)
845
846 # This isn't expected to raise an exception.
847 events = self.protocol.events_received()
848
849 # Write outgoing data to the socket.
850 try:
851 self.send_data()
852 except Exception as exc:
853 if self.debug:
854 self.logger.debug(
855 "! error while sending data",
856 exc_info=True,
857 )
858 # Similarly to the above, avoid overriding an exception
859 # set by send_context(), in case of a race condition
860 # i.e. send_context() closes the socket after recv()
861 # returns above but before send_data() calls send().
862 self.set_recv_exc(exc)
863 break
864
865 # If needed, set the close deadline based on the close timeout.
866 if self.protocol.close_expected():
867 if self.close_deadline is None:
868 self.close_deadline = Deadline(self.close_timeout)
869
870 # Unlock conn_mutex before processing events. Else, the
871 # application can't send messages in response to events.
872
873 # If self.send_data raised an exception, then events are lost.
874 # Given that automatic responses write small amounts of data,
875 # this should be uncommon, so we don't handle the edge case.
876
877 for event in events:
878 # This isn't expected to raise an exception.
879 self.process_event(event)
880
881 # Breaking out of the while True: ... loop means that we believe
882 # that the socket doesn't work anymore.
883
884 with self.protocol_mutex:
885 # Feed the end of the data stream to the protocol.
886 self.protocol.receive_eof()
887
888 # This isn't expected to raise an exception.
889 events = self.protocol.events_received()
890
891 # There is no error handling because send_data() can only write
892 # the end of the data stream and it handles errors by itself.
893 self.send_data()
894
895 # This code path is triggered when receiving an HTTP response
896 # without a Content-Length header. This is the only case where
897 # reading until EOF generates an event; all other events have
898 # a known length. Ignore for coverage measurement because tests
899 # are in test_client.py rather than test_connection.py.
900 for event in events: # pragma: no cover
901 # This isn't expected to raise an exception.
902 self.process_event(event)
903
904 except Exception as exc:
905 # This branch should never run. It's a safety net in case of bugs.
906 self.logger.error("unexpected internal error", exc_info=True)
907 with self.protocol_mutex:
908 self.set_recv_exc(exc)
909 finally:
910 # This isn't expected to raise an exception.
911 self.close_socket()
912
913 @contextlib.contextmanager
914 def send_context(
915 self,
916 *,
917 expected_state: State = OPEN, # CONNECTING during the opening handshake
918 ) -> Iterator[None]:
919 """
920 Create a context for writing to the connection from user code.
921
922 On entry, :meth:`send_context` acquires the connection lock and checks
923 that the connection is open; on exit, it writes outgoing data to the
924 socket and releases the connection lock::
925
926 with self.send_context():
927 self.protocol.send_text(message.encode())
928
929 When the connection isn't open on entry, when the connection is expected
930 to close on exit, or when an unexpected error happens, terminating the
931 connection, :meth:`send_context` waits until the connection is closed
932 then raises :exc:`~websockets.exceptions.ConnectionClosed`.
933
934 """
935 # Should we wait until the connection is closed?
936 wait_for_close = False
937 # Should we close the socket and raise ConnectionClosed?
938 raise_close_exc = False
939 # What exception should we chain ConnectionClosed to?
940 original_exc: BaseException | None = None
941
942 # Acquire the protocol lock.
943 with self.protocol_mutex:
944 if self.protocol.state is expected_state:
945 # Let the caller interact with the protocol.
946 try:
947 yield
948 except (ProtocolError, ConcurrencyError):
949 # The protocol state wasn't changed. Exit immediately.
950 raise
951 except Exception as exc:
952 self.logger.error("unexpected internal error", exc_info=True)
953 # This branch should never run. It's a safety net in case of
954 # bugs. Since we don't know what happened, we will close the
955 # connection and raise the exception to the caller.
956 wait_for_close = False
957 raise_close_exc = True
958 original_exc = exc
959 else:
960 # Check if the connection is expected to close soon.
961 if self.protocol.close_expected():
962 wait_for_close = True
963 # Set the close deadline based on the close timeout.
964 # Since we tested earlier that protocol.state is OPEN
965 # (or CONNECTING) and we didn't release protocol_mutex,
966 # self.close_deadline is still None.
967 assert self.close_deadline is None
968 self.close_deadline = Deadline(self.close_timeout)
969 # Write outgoing data to the socket.
970 try:
971 self.send_data()
972 except Exception as exc:
973 if self.debug:
974 self.logger.debug(
975 "! error while sending data",
976 exc_info=True,
977 )
978 # While the only expected exception here is OSError,
979 # other exceptions would be treated identically.
980 wait_for_close = False
981 raise_close_exc = True
982 original_exc = exc
983
984 else: # self.protocol.state is not expected_state
985 # Minor layering violation: we assume that the connection
986 # will be closing soon if it isn't in the expected state.
987 wait_for_close = True
988 # Calculate close_deadline if it wasn't set yet.
989 if self.close_deadline is None:
990 self.close_deadline = Deadline(self.close_timeout)
991 raise_close_exc = True
992
993 # To avoid a deadlock, release the connection lock by exiting the
994 # context manager before waiting for recv_events() to terminate.
995
996 # If the connection is expected to close soon and the close timeout
997 # elapses, close the socket to terminate the connection.
998 if wait_for_close:
999 # Thread.join() returns immediately if timeout is negative.
1000 assert self.close_deadline is not None
1001 timeout = self.close_deadline.timeout(raise_if_elapsed=False)
1002 self.recv_events_thread.join(timeout)
1003 if self.recv_events_thread.is_alive():
1004 # There's no risk of overwriting another error because
1005 # original_exc is never set when wait_for_close is True.
1006 assert original_exc is None
1007 original_exc = TimeoutError("timed out while closing connection")
1008 # Set recv_exc before closing the socket in order to get
1009 # proper exception reporting.
1010 raise_close_exc = True
1011 with self.protocol_mutex:
1012 self.set_recv_exc(original_exc)
1013
1014 # If an error occurred, close the socket to terminate the connection and
1015 # raise an exception.
1016 if raise_close_exc:
1017 self.close_socket()
1018 # Wait for the protocol state to be CLOSED before accessing close_exc.
1019 self.recv_events_thread.join()
1020 raise self.protocol.close_exc from original_exc
1021
1022 def send_data(self) -> None:
1023 """
1024 Send outgoing data.
1025
1026 This method requires holding protocol_mutex.
1027
1028 """
1029 assert self.protocol_mutex.locked()
1030 for data in self.protocol.data_to_send():
1031 if data:
1032 if self.close_deadline is not None:
1033 self.socket.settimeout(self.close_deadline.timeout())
1034 self.socket.sendall(data)
1035 else:
1036 try:
1037 self.socket.shutdown(socket.SHUT_WR)
1038 except OSError: # socket already closed
1039 pass
1040
1041 def set_recv_exc(self, exc: BaseException | None) -> None:
1042 """
1043 Set recv_exc, if not set yet.
1044
1045 This method requires holding protocol_mutex and must be called only from
1046 the thread running recv_events().
1047
1048 """
1049 assert self.protocol_mutex.locked()
1050 if self.recv_exc is None:
1051 self.recv_exc = exc
1052
1053 def close_socket(self) -> None:
1054 """
1055 Shutdown and close socket. Close message assembler.
1056
1057 Calling close_socket() guarantees that recv_events() terminates. Indeed,
1058 recv_events() may block only on socket.recv() or on recv_messages.put().
1059
1060 """
1061 # shutdown() is required to interrupt recv() on Linux.
1062 try:
1063 self.socket.shutdown(socket.SHUT_RDWR)
1064 except OSError: # socket already closed
1065 pass
1066 self.socket.close()
1067
1068 # Calling protocol.receive_eof() is safe because it's idempotent.
1069 # This guarantees that the protocol state becomes CLOSED.
1070 with self.protocol_mutex:
1071 self.protocol.receive_eof()
1072 assert self.protocol.state is CLOSED
1073
1074 # Abort recv() with a ConnectionClosed exception.
1075 self.recv_messages.close()
1076
1077 # Acknowledge pings sent with the ack_on_close option.
1078 self.terminate_pending_pings()
1079