diff --git a/docs/project/changelog.rst b/docs/project/changelog.rst index b74f84ce..fe0ba77e 100644 --- a/docs/project/changelog.rst +++ b/docs/project/changelog.rst @@ -39,6 +39,9 @@ Improvements .. _uvloop: https://uvloop.readthedocs.io/ +* Used non-blocking reads in the :mod:`threading` implementation. Depending on + platforms, this can make connection termination more reliable. + Bug fixes ......... diff --git a/src/websockets/sync/connection.py b/src/websockets/sync/connection.py index 15b6883e..51ccce80 100644 --- a/src/websockets/sync/connection.py +++ b/src/websockets/sync/connection.py @@ -4,6 +4,7 @@ import contextlib import logging import random +import selectors import socket import struct import threading @@ -23,7 +24,7 @@ ) from ..frames import DATA_OPCODES, PONG, CloseCode, Frame from ..http11 import Request, Response -from ..protocol import CLOSED, CONNECTING, OPEN, Event, Protocol, State +from ..protocol import CLOSED, OPEN, Event, Protocol, State from ..typing import BytesLike, Data, DataLike, LoggerLike, Subprotocol from .messages import Assembler from .utils import Deadline @@ -129,6 +130,9 @@ def __init__( # ConnectionClosed in order to show why the TCP connection dropped. self.recv_exc: BaseException | None = None + # Notification to wake up recv_events() when the connection is closed. + self.shutdown_watcher, self.shutdown_notifier = socket.socketpair() + # Receiving events from the socket. This thread is marked as daemon to # allow creating a connection in a non-daemon thread and using it in a # daemon thread. This mustn't prevent the interpreter from exiting. @@ -847,26 +851,33 @@ def recv_events(self) -> None: ``recv_events()`` exits immediately when ``self.socket`` is closed. """ - # When the opening handshake fails, we cannot trust rules in RFC 6455 - # for closing TCP connections will be followed. The HTTP server could - # keep the connection alive after sending the response. We attempt to - # close the connection immediately. Unfortunately, this is unreliable - # on macOS; recv() may block until close_timeout elapses: - # https://github.com/python-websockets/websockets/issues/1596 - # https://github.com/python/cpython/issues/154224 - # In that case, break out of the loop to prevent recv() from blocking - # until close_timeout elapses. - close_expected_while_connecting = False - + poller = selectors.DefaultSelector() try: - while not close_expected_while_connecting: + poller.register(self.shutdown_watcher, selectors.EVENT_READ) + try: + poller.register(self.socket, selectors.EVENT_READ, data="socket") + except (OSError, ValueError): + # close_socket() ran before poller.register(). + # * OSError: [Errno 9] Bad file descriptor + # (only seen on free-threaded Python for poller.register()) + # * ValueError: Invalid file descriptor: -1 + return + + while True: try: # If the assembler buffer is full, block until it drains. with self.recv_flow_control: pass - if self.close_deadline is not None: - self.socket.settimeout(self.close_deadline.timeout()) - data = self.socket.recv(self.recv_bufsize) + # Wait for the socket to be ready for reading. + for key, _ in poller.select(): + if key.data == "socket": + data = self.socket.recv(self.recv_bufsize) + break + else: + # If the socket wasn't ready for reading, then it's + # shutdown_watcher that was notified, meaning we're + # closing the socket. + data = b"" except Exception as exc: if self.debug: self.logger.debug( @@ -874,7 +885,7 @@ def recv_events(self) -> None: exc_info=True, ) # When the closing handshake is initiated by our side, - # recv() may block until send_context() closes the socket. + # select() may block until send_context() closes the socket. # In that case, send_context() already set recv_exc. # Calling set_recv_exc() avoids overwriting it. with self.protocol_mutex: @@ -903,7 +914,7 @@ def recv_events(self) -> None: ) # Similarly to the above, avoid overriding an exception # set by send_context(), in case of a race condition - # i.e. send_context() closes the socket after recv() + # i.e. send_context() closes the socket after select() # returns above but before send_data() calls send(). self.set_recv_exc(exc) break @@ -912,8 +923,6 @@ def recv_events(self) -> None: if self.protocol.close_expected(): if self.close_deadline is None: self.close_deadline = Deadline(self.close_timeout) - if self.protocol.state is CONNECTING: - close_expected_while_connecting = True # Unlock conn_mutex before processing events. Else, the # application can't send messages in response to events. @@ -926,8 +935,8 @@ def recv_events(self) -> None: # This isn't expected to raise an exception. self.process_event(event) - # Breaking out of the while not close_expected_while_connecting: ... - # loop means that we believe that the socket doesn't work anymore. + # Breaking out of the while True: ... loop means that we believe + # that the socket doesn't work anymore. with self.protocol_mutex: # Feed the end of the data stream to the protocol. @@ -957,6 +966,8 @@ def recv_events(self) -> None: finally: # This isn't expected to raise an exception. self.close_socket() + poller.close() + self.shutdown_watcher.close() @contextlib.contextmanager def send_context( @@ -1103,14 +1114,16 @@ def close_socket(self) -> None: Shutdown and close socket. Close message assembler. Calling close_socket() guarantees that recv_events() terminates. Indeed, - recv_events() may block only on socket.recv() or on recv_messages.put(). + recv_events() may block only on poller.select() or recv_messages.put(). """ - # shutdown() is required to interrupt recv() on Linux. try: - self.socket.shutdown(socket.SHUT_RDWR) - except OSError: # socket already closed + self.shutdown_notifier.send(b"x") + except OSError: # close_socket() was already called pass + finally: + self.shutdown_notifier.close() + try: self.socket.close() except OSError: # socket already closed # pragma: no cover diff --git a/src/websockets/trio/connection.py b/src/websockets/trio/connection.py index 26e03210..336d9b51 100644 --- a/src/websockets/trio/connection.py +++ b/src/websockets/trio/connection.py @@ -886,6 +886,8 @@ async def recv_events(self) -> None: if self.recv_flow_control is not None: await self.recv_flow_control.wait() data = await self.stream.receive_some() + except trio.ClosedResourceError: + break except Exception as exc: if self.debug: self.logger.debug( diff --git a/tests/asyncio/test_client.py b/tests/asyncio/test_client.py index a8c85846..293e41a2 100644 --- a/tests/asyncio/test_client.py +++ b/tests/asyncio/test_client.py @@ -516,12 +516,41 @@ def http_response(connection, request): self.assertEqual(raised.exception.response.status_code, 200) self.assertEqual(raised.exception.response.body.decode(), "👌") + async def test_http_response_with_keep_alive(self): + """Client reads HTTP response and server keeps connection alive.""" + + async def keep_alive(reader, writer): + # Wait for the client to send the handshake request. + await reader.read(4096) + writer.write( + b"HTTP/1.1 410 Gone\r\n" + b"Connection: keep-alive\r\n" # default in HTTP/1.1 + b"Content-Length: 0\r\n" # avoids "read body until EOF" + b"\r\n" + ) + # Wait for the client to close the connection. + await reader.read(4096) + writer.close() + + server = await asyncio.start_server(keep_alive, "localhost", 0) + host, port = get_host_port(server) + async with server: + with self.assertRaises(InvalidStatus) as raised: + async with asyncio.timeout(10 * MS): + async with connect(f"ws://{host}:{port}"): + self.fail("did not raise") + + self.assertEqual( + str(raised.exception), + "server rejected WebSocket connection: HTTP 410", + ) + async def test_junk_handshake_response(self): """Client closes the connection when receiving non-HTTP response from server.""" async def junk(reader, writer): # Wait for the client to send the handshake request. - await asyncio.sleep(MS) + await reader.read(4096) writer.write(b"220 smtp.invalid ESMTP Postfix\r\n") # Wait for the client to close the connection. await reader.read(4096) diff --git a/tests/asyncio/test_connection.py b/tests/asyncio/test_connection.py index ea4cbd20..25750c1f 100644 --- a/tests/asyncio/test_connection.py +++ b/tests/asyncio/test_connection.py @@ -1051,14 +1051,14 @@ async def test_pong_unsupported_type(self): async def test_keepalive(self, getrandbits): """keepalive sends pings at ping_interval and measures latency.""" getrandbits.side_effect = itertools.count(1918987876) - self.connection.ping_interval = 3 * MS + self.connection.ping_interval = 4 * MS self.connection.start_keepalive() self.assertIsNotNone(self.connection.keepalive_task) self.assertEqual(self.connection.latency, 0) - # 3 ms: keepalive() sends a ping frame. - # 3.x ms: a pong frame is received. - await asyncio.sleep(5 * MS) - # 5 ms: check that the ping frame was sent. + # 4 ms: keepalive() sends a ping frame. + # 4.x ms: a pong frame is received. + await asyncio.sleep(6 * MS) + # 6 ms: check that the ping frame was sent. await self.assertFrameSent(Frame(PING, b"rand")) self.assertGreater(self.connection.latency, 0) self.assertLess(self.connection.latency, MS) diff --git a/tests/sync/test_client.py b/tests/sync/test_client.py index bef30cfa..fdf198be 100644 --- a/tests/sync/test_client.py +++ b/tests/sync/test_client.py @@ -530,8 +530,8 @@ def test_http_response_with_keep_alive(self): class KeepAliveHandler(socketserver.BaseRequestHandler): def handle(self): - time.sleep(MS) # wait for the client to send the handshake request - self.request.recv(4096) # assume we read the full handshake request + # Wait for the client to send the handshake request. + self.request.recv(4096) self.request.send( b"HTTP/1.1 410 Gone\r\n" b"Connection: keep-alive\r\n" # default in HTTP/1.1 @@ -543,7 +543,7 @@ def handle(self): self.request.close() server = socketserver.TCPServer(("localhost", 0), KeepAliveHandler) - host, port = server.server_address + host, port = get_host_port(server) with server: thread = threading.Thread(target=server.serve_forever, args=(MS,)) thread.start() @@ -561,7 +561,7 @@ def handle(self): str(raised.exception), "server rejected WebSocket connection: HTTP 410", ) - self.assertLess(t1 - t0, 5 * MS) + self.assertLess(t1 - t0, 10 * MS) def test_junk_handshake_response(self): """Client closes the connection when receiving non-HTTP response from server.""" diff --git a/tests/sync/test_connection.py b/tests/sync/test_connection.py index 43e813a3..ab63dc1d 100644 --- a/tests/sync/test_connection.py +++ b/tests/sync/test_connection.py @@ -16,7 +16,7 @@ ConnectionClosedOK, ) from websockets.frames import BINARY, CLOSE, PING, PONG, TEXT, CloseCode, Frame -from websockets.protocol import CLIENT, CLOSED, CONNECTING, OPEN, SERVER, Protocol +from websockets.protocol import CLIENT, CLOSED, OPEN, SERVER, Protocol from websockets.sync.connection import * from websockets.sync.connection import broadcast @@ -787,14 +787,14 @@ def test_pong_unsupported_type(self): def test_keepalive(self, getrandbits): """keepalive sends pings at ping_interval and measures latency.""" getrandbits.side_effect = itertools.count(1918987876) - self.connection.ping_interval = 3 * MS + self.connection.ping_interval = 4 * MS self.connection.start_keepalive() self.assertIsNotNone(self.connection.keepalive_thread) self.assertEqual(self.connection.latency, 0) - # 3 ms: keepalive() sends a ping frame. - # 3.x ms: a pong frame is received. - time.sleep(5 * MS) - # 5 ms: check that the ping frame was sent. + # 4 ms: keepalive() sends a ping frame. + # 4.x ms: a pong frame is received. + time.sleep(6 * MS) + # 6 ms: check that the ping frame was sent. self.assertFrameSent(Frame(PING, b"rand")) self.assertGreater(self.connection.latency, 0) self.assertLess(self.connection.latency, MS) @@ -1033,6 +1033,15 @@ def test_writing_in_send_context_fails(self): # Test safety nets — catching all exceptions in case of bugs. + @patch("socket.socket.recv", side_effect=AssertionError) + def test_unexpected_failure_in_recv(self, recv): + """Unexpected internal error in recv() is correctly reported.""" + self.remote_connection.send("😀") + # Reading the message will trigger the injected fault. + with self.assertRaises(ConnectionClosedError) as raised: + self.connection.recv() + self.assertIsInstance(raised.exception.__cause__, AssertionError) + @patch("websockets.protocol.Protocol.events_received", side_effect=AssertionError) def test_unexpected_failure_in_recv_events(self, events_received): """Unexpected internal error in recv_events() is correctly reported.""" @@ -1050,30 +1059,6 @@ def test_unexpected_failure_in_send_context(self, send_text): self.connection.send("😀") self.assertIsInstance(raised.exception.__cause__, AssertionError) - # Test recv_events(). - - def test_close_expected_during_connecting(self): - """recv_events() terminates without waiting for close_timeout.""" - # See https://github.com/python-websockets/websockets/issues/1596. - socket_, remote_socket = socket.socketpair() - protocol = Protocol(self.LOCAL, state=CONNECTING) - # Simulate ClientProtocol / ServerProtocol deciding that the TCP - # connection must close because the opening handshake failed. - protocol.send_eof() - connection = Connection(socket_, protocol, close_timeout=5 * MS) - try: - t0 = time.time() - remote_socket.sendall(b"\x00") - connection.recv_events_thread.join(3 * MS) - t1 = time.time() - - self.assertFalse(connection.recv_events_thread.is_alive()) - self.assertLess(t1 - t0, 3 * MS) - finally: - remote_socket.close() - connection.close_socket() - connection.recv_events_thread.join() - # Test garbage collection. @unittest.skipUnless( diff --git a/tests/trio/test_client.py b/tests/trio/test_client.py index 34553ee4..a4d3d486 100644 --- a/tests/trio/test_client.py +++ b/tests/trio/test_client.py @@ -508,12 +508,48 @@ def http_response(connection, request): self.assertEqual(raised.exception.response.status_code, 200) self.assertEqual(raised.exception.response.body.decode(), "👌") + async def test_http_response_with_keep_alive(self): + """Client reads HTTP response and server keeps connection alive.""" + + async def keep_alive(stream): + # Wait for the client to send the handshake request. + await stream.receive_some() + await stream.send_all( + b"HTTP/1.1 410 Gone\r\n" + b"Connection: keep-alive\r\n" # default in HTTP/1.1 + b"Content-Length: 0\r\n" # avoids "read body until EOF" + b"\r\n" + ) + # Wait for the client to close the connection. + await stream.receive_some() + await stream.aclose() + + async with trio.open_nursery() as nursery: + listeners = await nursery.start(trio.serve_tcp, keep_alive, 0) + host, port = next( + listener + for listener in listeners + if listener.socket.family == socket.AF_INET + ).socket.getsockname() + try: + with self.assertRaises(InvalidStatus) as raised: + with trio.fail_after(10 * MS): + async with connect(f"ws://{host}:{port}"): + self.fail("did not raise") + finally: + nursery.cancel_scope.cancel() + + self.assertEqual( + str(raised.exception), + "server rejected WebSocket connection: HTTP 410", + ) + async def test_junk_handshake_response(self): """Client closes the connection when receiving non-HTTP response from server.""" async def junk(stream): # Wait for the client to send the handshake request. - await trio.testing.wait_all_tasks_blocked() + await stream.receive_some() await stream.send_all(b"220 smtp.invalid ESMTP Postfix\r\n") # Wait for the client to close the connection. await stream.receive_some() diff --git a/tests/trio/test_connection.py b/tests/trio/test_connection.py index 76e190ad..f3137a9d 100644 --- a/tests/trio/test_connection.py +++ b/tests/trio/test_connection.py @@ -978,14 +978,14 @@ def keepalive_task_is_running(self): async def test_keepalive(self, getrandbits): """keepalive sends pings at ping_interval and measures latency.""" getrandbits.side_effect = itertools.count(1918987876) - self.connection.ping_interval = 3 * MS + self.connection.ping_interval = 4 * MS self.connection.start_keepalive() self.assertTrue(self.keepalive_task_is_running()) self.assertEqual(self.connection.latency, 0) - # 3 ms: keepalive() sends a ping frame. - # 3.x ms: a pong frame is received. - await trio.sleep(5 * MS) - # 5 ms: check that the ping frame was sent. + # 4 ms: keepalive() sends a ping frame. + # 4.x ms: a pong frame is received. + await trio.sleep(6 * MS) + # 6 ms: check that the ping frame was sent. await self.assertFrameSent(Frame(PING, b"rand")) self.assertGreater(self.connection.latency, 0) self.assertLess(self.connection.latency, MS) @@ -1260,6 +1260,15 @@ async def test_writing_in_send_context_fails(self): # Test safety nets — catching all exceptions in case of bugs. + @patch("trio.testing.MemoryReceiveStream.receive_some", side_effect=AssertionError) + async def test_unexpected_failure_in_receive_some(self, receive_some): + """Unexpected internal error in receive_some() is correctly reported.""" + await self.remote_connection.send("😀") + # Reading the message will trigger the injected fault. + with self.assertRaises(ConnectionClosedError) as raised: + await self.connection.recv() + self.assertIsInstance(raised.exception.__cause__, AssertionError) + @patch("websockets.protocol.Protocol.events_received", side_effect=AssertionError) async def test_unexpected_failure_in_recv_events(self, events_received): """Unexpected internal error in recv_events() is correctly reported."""