From 1eb023c85ce95e49d16425a8171a7ad1afae5a54 Mon Sep 17 00:00:00 2001 From: Aymeric Augustin Date: Tue, 11 Aug 2026 08:47:11 +0200 Subject: [PATCH 1/7] Simplify declaration of werkzeug Map in tests. --- tests/asyncio/test_router.py | 40 ++++++++++++++++-------------------- tests/sync/test_router.py | 36 +++++++++++++++----------------- tests/trio/test_router.py | 38 ++++++++++++++++------------------ 3 files changed, 52 insertions(+), 62 deletions(-) diff --git a/tests/asyncio/test_router.py b/tests/asyncio/test_router.py index ea8e14bf5..c9651e9cc 100644 --- a/tests/asyncio/test_router.py +++ b/tests/asyncio/test_router.py @@ -16,6 +16,13 @@ from werkzeug.routing import Map, Rule except ImportError: pass +else: + url_map = Map( + [ + Rule("/", endpoint=handler), + Rule("/r", redirect_to="/"), + ] + ) async def echo(websocket, count): @@ -48,37 +55,28 @@ async def test_router_matches_paths_and_extracts_parameters(self): messages = await alist(client) self.assertEqual(messages, ["hello", "hello", "hello"]) - @property # avoids an import-time dependency on werkzeug - def url_map(self): - return Map( - [ - Rule("/", endpoint=handler), - Rule("/r", redirect_to="/"), - ] - ) - async def test_route_with_query_string(self): """Router ignores query strings when matching paths.""" - async with route(self.url_map, "localhost", 0) as server: + async with route(url_map, "localhost", 0) as server: async with connect(get_uri(server) + "/?a=b") as client: await self.assertEval(client, "ws.request.path", "/?a=b") async def test_redirect(self): """Router redirects connections according to redirect_to.""" - async with route(self.url_map, "localhost", 0) as server: + async with route(url_map, "localhost", 0) as server: async with connect(get_uri(server) + "/r") as client: await self.assertEval(client, "ws.request.path", "/") async def test_secure_redirect(self): """Router redirects connections according to redirect_to when TLS is enabled.""" - async with route(self.url_map, "localhost", 0, ssl=SERVER_CONTEXT) as server: + async with route(url_map, "localhost", 0, ssl=SERVER_CONTEXT) as server: async with connect(get_uri(server) + "/r", ssl=CLIENT_CONTEXT) as client: await self.assertEval(client, "ws.request.path", "/") @patch("websockets.asyncio.client.connect.process_redirect", lambda _, exc: exc) async def test_force_secure_redirect(self): """Router redirects ws:// connections to a wss:// URI when ssl=True.""" - async with route(self.url_map, "localhost", 0, ssl=True) as server: + async with route(url_map, "localhost", 0, ssl=True) as server: redirect_uri = get_uri(server, secure=True) with self.assertRaises(InvalidStatus) as raised: async with connect(get_uri(server) + "/r"): @@ -91,7 +89,7 @@ async def test_force_secure_redirect(self): @patch("websockets.asyncio.client.connect.process_redirect", lambda _, exc: exc) async def test_force_redirect_server_name(self): """Router redirects connections to the host declared in server_name.""" - async with route(self.url_map, "localhost", 0, server_name="other") as server: + async with route(url_map, "localhost", 0, server_name="other") as server: with self.assertRaises(InvalidStatus) as raised: async with connect(get_uri(server) + "/r"): self.fail("did not raise") @@ -102,7 +100,7 @@ async def test_force_redirect_server_name(self): async def test_not_found(self): """Router rejects requests to unknown paths with an HTTP 404 error.""" - async with route(self.url_map, "localhost", 0) as server: + async with route(url_map, "localhost", 0) as server: with self.assertRaises(InvalidStatus) as raised: async with connect(get_uri(server) + "/n"): self.fail("did not raise") @@ -118,7 +116,7 @@ def process_request(ws, request): ws.process_request_ran = True async with route( - self.url_map, "localhost", 0, process_request=process_request + url_map, "localhost", 0, process_request=process_request ) as server: async with connect(get_uri(server) + "/") as client: await self.assertEval(client, "ws.process_request_ran", "True") @@ -130,7 +128,7 @@ async def process_request(ws, request): ws.process_request_ran = True async with route( - self.url_map, "localhost", 0, process_request=process_request + url_map, "localhost", 0, process_request=process_request ) as server: async with connect(get_uri(server) + "/") as client: await self.assertEval(client, "ws.process_request_ran", "True") @@ -142,7 +140,7 @@ def process_request(ws, request): return ws.respond(http.HTTPStatus.FORBIDDEN, "Forbidden") async with route( - self.url_map, "localhost", 0, process_request=process_request + url_map, "localhost", 0, process_request=process_request ) as server: with self.assertRaises(InvalidStatus) as raised: async with connect(get_uri(server) + "/"): @@ -159,7 +157,7 @@ async def process_request(ws, request): return ws.respond(http.HTTPStatus.FORBIDDEN, "Forbidden") async with route( - self.url_map, "localhost", 0, process_request=process_request + url_map, "localhost", 0, process_request=process_request ) as server: with self.assertRaises(InvalidStatus) as raised: async with connect(get_uri(server) + "/"): @@ -177,9 +175,7 @@ async def handler(self, connection): connection.my_router_ran = True return await super().handler(connection) - async with route( - self.url_map, "localhost", 0, create_router=MyRouter - ) as server: + async with route(url_map, "localhost", 0, create_router=MyRouter) as server: async with connect(get_uri(server)) as client: await self.assertEval(client, "ws.my_router_ran", "True") diff --git a/tests/sync/test_router.py b/tests/sync/test_router.py index cf04a8689..5c8b4de2d 100644 --- a/tests/sync/test_router.py +++ b/tests/sync/test_router.py @@ -16,6 +16,13 @@ from werkzeug.routing import Map, Rule except ImportError: pass +else: + url_map = Map( + [ + Rule("/", endpoint=handler), + Rule("/r", redirect_to="/"), + ] + ) def echo(websocket, count): @@ -48,24 +55,15 @@ def test_router_matches_paths_and_extracts_parameters(self): messages = list(client) self.assertEqual(messages, ["hello", "hello", "hello"]) - @property # avoids an import-time dependency on werkzeug - def url_map(self): - return Map( - [ - Rule("/", endpoint=handler), - Rule("/r", redirect_to="/"), - ] - ) - def test_route_with_query_string(self): """Router ignores query strings when matching paths.""" - with run_router(self.url_map) as server: + with run_router(url_map) as server: with connect(get_uri(server) + "/?a=b") as client: self.assertEval(client, "ws.request.path", "/?a=b") def test_redirect(self): """Router redirects connections according to redirect_to.""" - with run_router(self.url_map, server_name="localhost") as server: + with run_router(url_map, server_name="localhost") as server: with self.assertRaises(InvalidStatus) as raised: with connect(get_uri(server) + "/r"): self.fail("did not raise") @@ -76,9 +74,7 @@ def test_redirect(self): def test_secure_redirect(self): """Router redirects connections to a wss:// URI when TLS is enabled.""" - with run_router( - self.url_map, server_name="localhost", ssl=SERVER_CONTEXT - ) as server: + with run_router(url_map, server_name="localhost", ssl=SERVER_CONTEXT) as server: with self.assertRaises(InvalidStatus) as raised: with connect(get_uri(server) + "/r", ssl=CLIENT_CONTEXT): self.fail("did not raise") @@ -90,7 +86,7 @@ def test_secure_redirect(self): @patch("websockets.asyncio.client.connect.process_redirect", lambda _, exc: exc) def test_force_secure_redirect(self): """Router redirects ws:// connections to a wss:// URI when ssl=True.""" - with run_router(self.url_map, ssl=True) as server: + with run_router(url_map, ssl=True) as server: redirect_uri = get_uri(server, secure=True) with self.assertRaises(InvalidStatus) as raised: with connect(get_uri(server) + "/r"): @@ -103,7 +99,7 @@ def test_force_secure_redirect(self): @patch("websockets.asyncio.client.connect.process_redirect", lambda _, exc: exc) def test_force_redirect_server_name(self): """Router redirects connections to the host declared in server_name.""" - with run_router(self.url_map, server_name="other") as server: + with run_router(url_map, server_name="other") as server: with self.assertRaises(InvalidStatus) as raised: with connect(get_uri(server) + "/r"): self.fail("did not raise") @@ -114,7 +110,7 @@ def test_force_redirect_server_name(self): def test_not_found(self): """Router rejects requests to unknown paths with an HTTP 404 error.""" - with run_router(self.url_map) as server: + with run_router(url_map) as server: with self.assertRaises(InvalidStatus) as raised: with connect(get_uri(server) + "/n"): self.fail("did not raise") @@ -129,7 +125,7 @@ def test_process_request_returning_none(self): def process_request(ws, request): ws.process_request_ran = True - with run_router(self.url_map, process_request=process_request) as server: + with run_router(url_map, process_request=process_request) as server: with connect(get_uri(server) + "/") as client: self.assertEval(client, "ws.process_request_ran", "True") @@ -139,7 +135,7 @@ def test_process_request_returning_response(self): def process_request(ws, request): return ws.respond(http.HTTPStatus.FORBIDDEN, "Forbidden") - with run_router(self.url_map, process_request=process_request) as server: + with run_router(url_map, process_request=process_request) as server: with self.assertRaises(InvalidStatus) as raised: with connect(get_uri(server) + "/"): self.fail("did not raise") @@ -156,7 +152,7 @@ def handler(self, connection): connection.my_router_ran = True return super().handler(connection) - with run_router(self.url_map, create_router=MyRouter) as server: + with run_router(url_map, create_router=MyRouter) as server: with connect(get_uri(server)) as client: self.assertEval(client, "ws.my_router_ran", "True") diff --git a/tests/trio/test_router.py b/tests/trio/test_router.py index a85f5bf01..e303a3bd6 100644 --- a/tests/trio/test_router.py +++ b/tests/trio/test_router.py @@ -16,6 +16,13 @@ from werkzeug.routing import Map, Rule except ImportError: pass +else: + url_map = Map( + [ + Rule("/", endpoint=handler), + Rule("/r", redirect_to="/"), + ] + ) async def echo(websocket, count): @@ -48,37 +55,28 @@ async def test_router_matches_paths_and_extracts_parameters(self): messages = await alist(client) self.assertEqual(messages, ["hello", "hello", "hello"]) - @property # avoids an import-time dependency on werkzeug - def url_map(self): - return Map( - [ - Rule("/", endpoint=handler), - Rule("/r", redirect_to="/"), - ] - ) - async def test_route_with_query_string(self): """Router ignores query strings when matching paths.""" - async with run_router(self.url_map) as server: + async with run_router(url_map) as server: async with connect(get_uri(server) + "/?a=b") as client: await self.assertEval(client, "ws.request.path", "/?a=b") async def test_redirect(self): """Router redirects connections according to redirect_to.""" - async with run_router(self.url_map) as server: + async with run_router(url_map) as server: async with connect(get_uri(server) + "/r") as client: await self.assertEval(client, "ws.request.path", "/") async def test_secure_redirect(self): """Router redirects connections according to redirect_to when TLS is enabled.""" - async with run_router(self.url_map, ssl=SERVER_CONTEXT) as server: + async with run_router(url_map, ssl=SERVER_CONTEXT) as server: async with connect(get_uri(server) + "/r", ssl=CLIENT_CONTEXT) as client: await self.assertEval(client, "ws.request.path", "/") @patch("websockets.trio.client.connect.process_redirect", lambda _, exc: exc) async def test_force_secure_redirect(self): """Router redirects ws:// connections to a wss:// URI when ssl=True.""" - async with run_router(self.url_map, ssl=True) as server: + async with run_router(url_map, ssl=True) as server: redirect_uri = get_uri(server, secure=True) with self.assertRaises(InvalidStatus) as raised: async with connect(get_uri(server) + "/r"): @@ -91,7 +89,7 @@ async def test_force_secure_redirect(self): @patch("websockets.trio.client.connect.process_redirect", lambda _, exc: exc) async def test_force_redirect_server_name(self): """Router redirects connections to the host declared in server_name.""" - async with run_router(self.url_map, server_name="other") as server: + async with run_router(url_map, server_name="other") as server: with self.assertRaises(InvalidStatus) as raised: async with connect(get_uri(server) + "/r"): self.fail("did not raise") @@ -102,7 +100,7 @@ async def test_force_redirect_server_name(self): async def test_not_found(self): """Router rejects requests to unknown paths with an HTTP 404 error.""" - async with run_router(self.url_map) as server: + async with run_router(url_map) as server: with self.assertRaises(InvalidStatus) as raised: async with connect(get_uri(server) + "/n"): self.fail("did not raise") @@ -117,7 +115,7 @@ async def test_process_request_function_returning_none(self): def process_request(ws, request): ws.process_request_ran = True - async with run_router(self.url_map, process_request=process_request) as server: + async with run_router(url_map, process_request=process_request) as server: async with connect(get_uri(server) + "/") as client: await self.assertEval(client, "ws.process_request_ran", "True") @@ -127,7 +125,7 @@ async def test_process_request_coroutine_returning_none(self): async def process_request(ws, request): ws.process_request_ran = True - async with run_router(self.url_map, process_request=process_request) as server: + async with run_router(url_map, process_request=process_request) as server: async with connect(get_uri(server) + "/") as client: await self.assertEval(client, "ws.process_request_ran", "True") @@ -137,7 +135,7 @@ async def test_process_request_function_returning_response(self): def process_request(ws, request): return ws.respond(http.HTTPStatus.FORBIDDEN, "Forbidden") - async with run_router(self.url_map, process_request=process_request) as server: + async with run_router(url_map, process_request=process_request) as server: with self.assertRaises(InvalidStatus) as raised: async with connect(get_uri(server) + "/"): self.fail("did not raise") @@ -152,7 +150,7 @@ async def test_process_request_coroutine_returning_response(self): async def process_request(ws, request): return ws.respond(http.HTTPStatus.FORBIDDEN, "Forbidden") - async with run_router(self.url_map, process_request=process_request) as server: + async with run_router(url_map, process_request=process_request) as server: with self.assertRaises(InvalidStatus) as raised: async with connect(get_uri(server) + "/"): self.fail("did not raise") @@ -169,6 +167,6 @@ async def handler(self, connection): connection.my_router_ran = True return await super().handler(connection) - async with run_router(self.url_map, create_router=MyRouter) as server: + async with run_router(url_map, create_router=MyRouter) as server: async with connect(get_uri(server)) as client: await self.assertEval(client, "ws.my_router_ran", "True") From 46817d2a6f7f7aad7b78ae668922ee204286ecc2 Mon Sep 17 00:00:00 2001 From: Aymeric Augustin Date: Tue, 11 Aug 2026 08:09:42 +0200 Subject: [PATCH 2/7] Make implementations of test servers consistent. --- tests/sync/server.py | 12 ++++++++---- tests/sync/test_client.py | 6 +++--- tests/sync/test_server.py | 11 ++++++----- tests/trio/server.py | 30 +++++++++++++++++------------- tests/trio/test_client.py | 22 +++++++++++++--------- tests/trio/test_server.py | 17 +++++++++++------ 6 files changed, 58 insertions(+), 40 deletions(-) diff --git a/tests/sync/server.py b/tests/sync/server.py index 78d1c6745..194829fea 100644 --- a/tests/sync/server.py +++ b/tests/sync/server.py @@ -8,11 +8,15 @@ from websockets.sync.server import serve, unix_serve +def get_host_port(server): + return server.socket.getsockname() + + def get_uri(server, secure=None): if secure is None: secure = isinstance(server.socket, ssl.SSLSocket) # hack protocol = "wss" if secure else "ws" - host, port = server.socket.getsockname() + host, port = get_host_port(server) return f"{protocol}://{host}:{port}" @@ -69,9 +73,9 @@ def run_router(url_map, **kwargs): @contextlib.contextmanager def run_unix_server_or_router( - path, unix_serve_or_route, handler_or_url_map, + path, **kwargs, ): with unix_serve_or_route(handler_or_url_map, path, **kwargs) as server: @@ -85,8 +89,8 @@ def run_unix_server_or_router( def run_unix_server(path, handler=handler, **kwargs): - return run_unix_server_or_router(path, unix_serve, handler, **kwargs) + return run_unix_server_or_router(unix_serve, handler, path, **kwargs) def run_unix_router(path, url_map, **kwargs): - return run_unix_server_or_router(path, unix_route, url_map, **kwargs) + return run_unix_server_or_router(unix_route, url_map, path, **kwargs) diff --git a/tests/sync/test_client.py b/tests/sync/test_client.py index d4d42c318..1ecb649b9 100644 --- a/tests/sync/test_client.py +++ b/tests/sync/test_client.py @@ -30,7 +30,7 @@ DeprecationTestCase, temp_unix_socket_path, ) -from .server import get_uri, run_server, run_unix_server +from .server import get_host_port, get_uri, run_server, run_unix_server class ClientTests(unittest.TestCase): @@ -43,7 +43,7 @@ def test_connection(self): def test_existing_socket(self): """Client connects using a pre-existing socket.""" with run_server() as server: - with socket.create_connection(server.socket.getsockname()) as sock: + with socket.create_connection(get_host_port(server)) as sock: # Use a non-existing domain to ensure we connect via sock. with connect("ws://invalid/", sock=sock) as client: self.assertEqual(client.protocol.state.name, "OPEN") @@ -440,7 +440,7 @@ def test_explicit_socks_proxy(self): def test_ignore_proxy_with_existing_socket(self): """Client connects using a pre-existing socket.""" with run_server() as server: - with socket.create_connection(server.socket.getsockname()) as sock: + with socket.create_connection(get_host_port(server)) as sock: # Use a non-existing domain to ensure we connect via sock. with connect("ws://invalid/", sock=sock) as client: self.assertEqual(client.protocol.state.name, "OPEN") diff --git a/tests/sync/test_server.py b/tests/sync/test_server.py index ee01a5eef..92280ac3f 100644 --- a/tests/sync/test_server.py +++ b/tests/sync/test_server.py @@ -29,6 +29,7 @@ ) from .server import ( EvalShellMixin, + get_host_port, get_uri, handler, run_server, @@ -320,7 +321,7 @@ def test_timeout_before_handshake_request(self): """Server times out before receiving handshake request from client.""" with self.assertLogs("websockets", logging.DEBUG) as logs: with run_server(open_timeout=MS) as server: - with socket.create_connection(server.socket.getsockname()) as sock: + with socket.create_connection(get_host_port(server)) as sock: # Wait for the server to close the connection. self.assertEqual(sock.recv(4096), b"") @@ -334,7 +335,7 @@ def test_connection_closed_before_handshake_request(self): """Server reads EOF before receiving handshake request from client.""" with self.assertLogs("websockets", logging.DEBUG) as logs: with run_server() as server: - with socket.create_connection(server.socket.getsockname()): + with socket.create_connection(get_host_port(server)): # Wait for the server to receive the connection, then close it. time.sleep(MS) @@ -360,7 +361,7 @@ def test_junk_handshake_request(self): """Server closes the connection when receiving non-HTTP request from client.""" with self.assertLogs("websockets.server", logging.DEBUG) as logs: with run_server() as server: - with socket.create_connection(server.socket.getsockname()) as sock: + with socket.create_connection(get_host_port(server)) as sock: sock.send(b"HELO relay.invalid\r\n") # Wait for the server to close the connection. self.assertEqual(sock.recv(4096), b"") @@ -502,14 +503,14 @@ def test_connection(self): def test_timeout_during_tls_handshake(self): """Server times out before receiving TLS handshake request from client.""" with run_server(ssl=SERVER_CONTEXT, open_timeout=MS) as server: - with socket.create_connection(server.socket.getsockname()) as sock: + with socket.create_connection(get_host_port(server)) as sock: # Wait for the server to close the connection. self.assertEqual(sock.recv(4096), b"") def test_connection_closed_during_tls_handshake(self): """Server reads EOF before receiving TLS handshake request from client.""" with run_server(ssl=SERVER_CONTEXT) as server: - with socket.create_connection(server.socket.getsockname()): + with socket.create_connection(get_host_port(server)): # Wait for the server to receive the connection, then close it. time.sleep(MS) diff --git a/tests/trio/server.py b/tests/trio/server.py index 6e9ef417e..7b69e491b 100644 --- a/tests/trio/server.py +++ b/tests/trio/server.py @@ -10,8 +10,8 @@ from websockets.trio.server import serve -def get_host_port(listeners): - for listener in listeners: +def get_host_port(server): + for listener in server.listeners: if listener.socket.family == socket.AF_INET: # pragma: no branch return listener.socket.getsockname() raise AssertionError("expected at least one IPv4 socket") @@ -24,7 +24,7 @@ def get_uri(server, secure=None): for cell in server.handler.__closure__ ) # l33t hack protocol = "wss" if secure else "ws" - host, port = get_host_port(server.listeners) + host, port = get_host_port(server) return f"{protocol}://{host}:{port}" @@ -53,19 +53,23 @@ async def assertEval(self, client, expr, value): self.assertEqual(await client.recv(), value) -kwargs = {"port": 0, "host": "localhost"} - - @contextlib.asynccontextmanager async def run_server_or_route( serve_or_route, handler_or_url_map, - **overrides, + port=0, + host="localhost", + **kwargs, ): - merged_kwargs = {**kwargs, **overrides} async with trio.open_nursery() as nursery: server = await nursery.start( - functools.partial(serve_or_route, handler_or_url_map, **merged_kwargs) + functools.partial( + serve_or_route, + handler_or_url_map, + port, + host=host, + **kwargs, + ) ) try: yield server @@ -76,9 +80,9 @@ async def run_server_or_route( nursery.cancel_scope.cancel() -def run_server(handler=handler, **overrides): - return run_server_or_route(serve, handler, **overrides) +def run_server(handler=handler, **kwargs): + return run_server_or_route(serve, handler, **kwargs) -def run_router(url_map, **overrides): - return run_server_or_route(route, url_map, **overrides) +def run_router(url_map, **kwargs): + return run_server_or_route(route, url_map, **kwargs) diff --git a/tests/trio/test_client.py b/tests/trio/test_client.py index 6afd35f81..b8b74fa3c 100644 --- a/tests/trio/test_client.py +++ b/tests/trio/test_client.py @@ -63,14 +63,14 @@ async def test_connection(self): async def test_explicit_host_port(self): """Client connects using an explicit host / port.""" async with run_server() as server: - host, port = get_host_port(server.listeners) + host, port = get_host_port(server) async with connect("ws://overridden/", host=host, port=port) as client: self.assertEqual(client.protocol.state.name, "OPEN") async def test_existing_stream(self): """Client connects using a pre-existing stream.""" async with run_server() as server: - stream = await trio.open_tcp_stream(*get_host_port(server.listeners)) + stream = await trio.open_tcp_stream(*get_host_port(server)) # Use a non-existing domain to ensure we connect via stream. async with connect("ws://invalid/", stream=stream) as client: self.assertEqual(client.protocol.state.name, "OPEN") @@ -306,7 +306,7 @@ def redirect(connection, request): return response async with run_server(process_request=redirect) as server: - host, port = get_host_port(server.listeners) + host, port = get_host_port(server) async with connect( "ws://overridden/redirect", host=host, port=port ) as client: @@ -321,7 +321,7 @@ def redirect(connection, request): return response async with run_server(process_request=redirect) as server: - host, port = get_host_port(server.listeners) + host, port = get_host_port(server) with self.assertRaises(ValueError) as raised: async with connect("ws://overridden/", host=host, port=port): self.fail("did not raise") @@ -341,9 +341,9 @@ def redirect(connection, request): return response async with run_server(process_request=redirect) as server: - stream = await trio.open_tcp_stream(*get_host_port(server.listeners)) + stream = await trio.open_tcp_stream(*get_host_port(server)) with self.assertRaises(ValueError) as raised: - # Use a non-existing domain to ensure we connect via sock. + # Use a non-existing domain to ensure we connect via stream. async with connect("ws://invalid/redirect", stream=stream): self.fail("did not raise") @@ -505,7 +505,11 @@ async def junk(stream): async with trio.open_nursery() as nursery: try: listeners = await nursery.start(trio.serve_tcp, junk, 0) - host, port = get_host_port(listeners) + host, port = next( + listener + for listener in listeners + if listener.socket.family == socket.AF_INET + ).socket.getsockname() with self.assertRaises(InvalidMessage) as raised: async with connect(f"ws://{host}:{port}"): self.fail("did not raise") @@ -537,7 +541,7 @@ async def test_connection(self): async def test_set_server_hostname_implicitly(self): """Client sets server_hostname to the host in the WebSocket URI.""" async with run_server(ssl=SERVER_CONTEXT) as server: - host, port = get_host_port(server.listeners) + host, port = get_host_port(server) async with connect( "wss://overridden/", host=host, port=port, ssl=CLIENT_CONTEXT ) as client: @@ -721,7 +725,7 @@ async def test_explicit_socks_proxy(self): async def test_ignore_proxy_with_existing_stream(self): """Cli ent connects using a pre-existing stream.""" async with run_server() as server: - stream = await trio.open_tcp_stream(*get_host_port(server.listeners)) + stream = await trio.open_tcp_stream(*get_host_port(server)) # Use a non-existing domain to ensure we connect via stream. async with connect("ws://invalid/", stream=stream) as client: self.assertEqual(client.protocol.state.name, "OPEN") diff --git a/tests/trio/test_server.py b/tests/trio/test_server.py index 8018175e4..ffd7337c9 100644 --- a/tests/trio/test_server.py +++ b/tests/trio/test_server.py @@ -2,6 +2,7 @@ import hmac import http import logging +import socket import trio @@ -64,7 +65,11 @@ async def test_connection_handler_raises_exception(self): async def test_existing_listeners(self): """Server receives connection using pre-existing listeners.""" listeners = await trio.open_tcp_listeners(0, host="localhost") - host, port = get_host_port(listeners) + host, port = next( + listener + for listener in listeners + if listener.socket.family == socket.AF_INET + ).socket.getsockname() # Unset the default values of port and host set by run_server. async with run_server(port=None, host=None, listeners=listeners): async with connect(f"ws://{host}:{port}/") as client: # type: ignore @@ -416,7 +421,7 @@ async def test_timeout_before_handshake_request(self): """Server times out before receiving handshake request from client.""" with self.assertLogs("websockets", logging.DEBUG) as logs: async with run_server(open_timeout=MS) as server: - stream = await trio.open_tcp_stream(*get_host_port(server.listeners)) + stream = await trio.open_tcp_stream(*get_host_port(server)) try: # Wait for the server to close the connection. self.assertEqual(await stream.receive_some(4096), b"") @@ -433,7 +438,7 @@ async def test_connection_closed_before_handshake_request(self): """Server reads EOF before receiving handshake request from client.""" with self.assertLogs("websockets", logging.DEBUG) as logs: async with run_server() as server: - stream = await trio.open_tcp_stream(*get_host_port(server.listeners)) + stream = await trio.open_tcp_stream(*get_host_port(server)) await stream.aclose() self.assertExceptionLogged( @@ -459,7 +464,7 @@ async def test_junk_handshake_request(self): """Server closes the connection when receiving non-HTTP request from client.""" with self.assertLogs("websockets", logging.DEBUG) as logs: async with run_server() as server: - stream = await trio.open_tcp_stream(*get_host_port(server.listeners)) + stream = await trio.open_tcp_stream(*get_host_port(server)) await stream.send_all(b"HELO relay.invalid\r\n") try: # Wait for the server to close the connection. @@ -612,7 +617,7 @@ async def test_connection(self): async def test_timeout_during_tls_handshake(self): """Server times out before receiving TLS handshake request from client.""" async with run_server(ssl=SERVER_CONTEXT, open_timeout=MS) as server: - stream = await trio.open_tcp_stream(*get_host_port(server.listeners)) + stream = await trio.open_tcp_stream(*get_host_port(server)) try: # Wait for the server to close the connection. self.assertEqual(await stream.receive_some(4096), b"") @@ -622,7 +627,7 @@ async def test_timeout_during_tls_handshake(self): async def test_connection_closed_during_tls_handshake(self): """Server reads EOF before receiving TLS handshake request from client.""" async with run_server(ssl=SERVER_CONTEXT) as server: - stream = await trio.open_tcp_stream(*get_host_port(server.listeners)) + stream = await trio.open_tcp_stream(*get_host_port(server)) await stream.aclose() From fab6be379acacb36f3f9b1d7eb4245da304af03c Mon Sep 17 00:00:00 2001 From: Aymeric Augustin Date: Tue, 18 Aug 2026 09:19:20 +0200 Subject: [PATCH 3/7] Add tests for connecting without a context manager. --- tests/asyncio/test_client.py | 54 ++++++++++++++++++++++++++++++------ tests/sync/test_client.py | 52 ++++++++++++++++++++++++++++------ tests/trio/test_client.py | 20 +++++++------ 3 files changed, 102 insertions(+), 24 deletions(-) diff --git a/tests/asyncio/test_client.py b/tests/asyncio/test_client.py index 6f6070c60..b011a1fd8 100644 --- a/tests/asyncio/test_client.py +++ b/tests/asyncio/test_client.py @@ -52,11 +52,19 @@ async def few_redirects(): class ClientTests(unittest.IsolatedAsyncioTestCase): - async def test_connection(self): - """Client connects to server.""" + async def test_context_manager(self): + """Client connects to server and disconnects automatically.""" async with serve(*args) as server: async with connect(get_uri(server)) as client: self.assertEqual(client.protocol.state.name, "OPEN") + self.assertEqual(client.protocol.state.name, "CLOSED") + + async def test_direct_connection(self): + """Client connects to server directly.""" + async with serve(*args) as server: + client = await connect(get_uri(server)) + self.addAsyncCleanup(client.close) + self.assertEqual(client.protocol.state.name, "OPEN") async def test_explicit_host_port(self): """Client connects using an explicit host / port.""" @@ -519,13 +527,23 @@ async def junk(reader, writer): class SecureClientTests(unittest.IsolatedAsyncioTestCase): - async def test_connection(self): - """Client connects to server securely.""" + async def test_context_manager(self): + """Client connects to server securely and disconnects automatically.""" async with serve(*args, ssl=SERVER_CONTEXT) as server: async with connect(get_uri(server), ssl=CLIENT_CONTEXT) as client: self.assertEqual(client.protocol.state.name, "OPEN") ssl_object = client.transport.get_extra_info("ssl_object") self.assertEqual(ssl_object.version()[:3], "TLS") + self.assertEqual(client.protocol.state.name, "CLOSED") + + async def test_direct_connection(self): + """Client connects to server securely and directly.""" + async with serve(*args, ssl=SERVER_CONTEXT) as server: + client = await connect(get_uri(server), ssl=CLIENT_CONTEXT) + self.addAsyncCleanup(client.close) + self.assertEqual(client.protocol.state.name, "OPEN") + ssl_object = client.transport.get_extra_info("ssl_object") + self.assertEqual(ssl_object.version()[:3], "TLS") async def test_set_server_hostname_implicitly(self): """Client sets server_hostname to the host in the WebSocket URI.""" @@ -899,12 +917,21 @@ async def test_https_proxy_invalid_server_certificate(self): @unittest.skipUnless(hasattr(socket, "AF_UNIX"), "this test requires Unix sockets") class UnixClientTests(unittest.IsolatedAsyncioTestCase): - async def test_connection(self): - """Client connects to server over a Unix socket.""" + async def test_context_manager(self): + """Client connects to Unix server and disconnects automatically.""" with temp_unix_socket_path() as path: async with unix_serve(handler, path): async with unix_connect(path) as client: self.assertEqual(client.protocol.state.name, "OPEN") + self.assertEqual(client.protocol.state.name, "CLOSED") + + async def test_direct_connection(self): + """Client connects to Unix server directly.""" + with temp_unix_socket_path() as path: + async with unix_serve(handler, path): + client = await unix_connect(path) + self.addAsyncCleanup(client.close) + self.assertEqual(client.protocol.state.name, "OPEN") async def test_set_host_header(self): """Client sets the Host header to the host in the WebSocket URI.""" @@ -933,14 +960,25 @@ def redirect(connection, request): "cannot follow cross-origin redirect to ws://other/ with a Unix socket", ) - async def test_secure_connection(self): - """Client connects to server securely over a Unix socket.""" + async def test_secure_context_manager(self): + """Client connects to Unix server securely and disconnects automatically.""" with temp_unix_socket_path() as path: async with unix_serve(handler, path, ssl=SERVER_CONTEXT): async with unix_connect(path, ssl=CLIENT_CONTEXT) as client: self.assertEqual(client.protocol.state.name, "OPEN") ssl_object = client.transport.get_extra_info("ssl_object") self.assertEqual(ssl_object.version()[:3], "TLS") + self.assertEqual(client.protocol.state.name, "CLOSED") + + async def test_secure_direct_connection(self): + """Client connects to Unix server securely and directly.""" + with temp_unix_socket_path() as path: + async with unix_serve(handler, path, ssl=SERVER_CONTEXT): + client = await unix_connect(path, ssl=CLIENT_CONTEXT) + self.addAsyncCleanup(client.close) + self.assertEqual(client.protocol.state.name, "OPEN") + ssl_object = client.transport.get_extra_info("ssl_object") + self.assertEqual(ssl_object.version()[:3], "TLS") async def test_set_server_hostname(self): """Client sets server_hostname to the host in the WebSocket URI.""" diff --git a/tests/sync/test_client.py b/tests/sync/test_client.py index 1ecb649b9..fd299dace 100644 --- a/tests/sync/test_client.py +++ b/tests/sync/test_client.py @@ -34,11 +34,19 @@ class ClientTests(unittest.TestCase): - def test_connection(self): - """Client connects to server and the handshake succeeds.""" + def test_context_manager(self): + """Client connects to server and disconnects automatically.""" with run_server() as server: with connect(get_uri(server)) as client: self.assertEqual(client.protocol.state.name, "OPEN") + self.assertEqual(client.protocol.state.name, "CLOSED") + + def test_direct_connection(self): + """Client connects to server directly.""" + with run_server() as server: + client = connect(get_uri(server)) + self.addCleanup(client.close) + self.assertEqual(client.protocol.state.name, "OPEN") def test_existing_socket(self): """Client connects using a pre-existing socket.""" @@ -295,12 +303,21 @@ def handle(self): class SecureClientTests(unittest.TestCase): - def test_connection(self): - """Client connects to server securely.""" + def test_context_manager(self): + """Client connects to server securely and disconnects automatically.""" with run_server(ssl=SERVER_CONTEXT) as server: with connect(get_uri(server), ssl=CLIENT_CONTEXT) as client: self.assertEqual(client.protocol.state.name, "OPEN") self.assertEqual(client.socket.version()[:3], "TLS") + self.assertEqual(client.protocol.state.name, "CLOSED") + + def test_direct_connection(self): + """Client connects to server directly.""" + with run_server(ssl=SERVER_CONTEXT) as server: + client = connect(get_uri(server), ssl=CLIENT_CONTEXT) + self.addCleanup(client.close) + self.assertEqual(client.protocol.state.name, "OPEN") + self.assertEqual(client.socket.version()[:3], "TLS") def test_set_server_hostname_implicitly(self): """Client sets server_hostname to the host in the WebSocket URI.""" @@ -637,12 +654,21 @@ def test_https_proxy_invalid_server_certificate(self): @unittest.skipUnless(hasattr(socket, "AF_UNIX"), "this test requires Unix sockets") class UnixClientTests(unittest.TestCase): - def test_connection(self): - """Client connects to server over a Unix socket.""" + def test_context_manager(self): + """Client connects to Unix server and disconnects automatically.""" with temp_unix_socket_path() as path: with run_unix_server(path): with unix_connect(path) as client: self.assertEqual(client.protocol.state.name, "OPEN") + self.assertEqual(client.protocol.state.name, "CLOSED") + + def test_direct_connection(self): + """Client connects to Unix server directly.""" + with temp_unix_socket_path() as path: + with run_unix_server(path): + client = unix_connect(path) + self.addCleanup(client.close) + self.assertEqual(client.protocol.state.name, "OPEN") def test_set_host_header(self): """Client sets the Host header to the host in the WebSocket URI.""" @@ -652,13 +678,23 @@ def test_set_host_header(self): with unix_connect(path, uri="ws://overridden/") as client: self.assertEqual(client.request.headers["Host"], "overridden") - def test_secure_connection(self): - """Client connects to server securely over a Unix socket.""" + def test_secure_context_manager(self): + """Client connects to Unix server securely and disconnects automatically.""" with temp_unix_socket_path() as path: with run_unix_server(path, ssl=SERVER_CONTEXT): with unix_connect(path, ssl=CLIENT_CONTEXT) as client: self.assertEqual(client.protocol.state.name, "OPEN") self.assertEqual(client.socket.version()[:3], "TLS") + self.assertEqual(client.protocol.state.name, "CLOSED") + + def test_secure_direct_connection(self): + """Client connects to Unix server securely and directly.""" + with temp_unix_socket_path() as path: + with run_unix_server(path, ssl=SERVER_CONTEXT): + client = unix_connect(path, ssl=CLIENT_CONTEXT) + self.addCleanup(client.close) + self.assertEqual(client.protocol.state.name, "OPEN") + self.assertEqual(client.socket.version()[:3], "TLS") def test_set_server_hostname(self): """Client sets server_hostname to the host in the WebSocket URI.""" diff --git a/tests/trio/test_client.py b/tests/trio/test_client.py index b8b74fa3c..cedd078cf 100644 --- a/tests/trio/test_client.py +++ b/tests/trio/test_client.py @@ -54,11 +54,12 @@ async def few_redirects(): class ClientTests(IsolatedTrioTestCase): - async def test_connection(self): - """Client connects to server.""" + async def test_context_manager(self): + """Client connects to server and disconnects automatically.""" async with run_server() as server: async with connect(get_uri(server)) as client: self.assertEqual(client.protocol.state.name, "OPEN") + self.assertEqual(client.protocol.state.name, "CLOSED") async def test_explicit_host_port(self): """Client connects using an explicit host / port.""" @@ -528,8 +529,8 @@ async def junk(stream): class SecureClientTests(IsolatedTrioTestCase): - async def test_connection(self): - """Client connects to server securely.""" + async def test_context_manager(self): + """Client connects to server securely and disconnects automatically.""" async with run_server(ssl=SERVER_CONTEXT) as server: async with connect( get_uri(server, secure=True), ssl=CLIENT_CONTEXT @@ -537,6 +538,7 @@ async def test_connection(self): self.assertEqual(client.protocol.state.name, "OPEN") ssl_object = client.stream._ssl_object self.assertEqual(ssl_object.version()[:3], "TLS") + self.assertEqual(client.protocol.state.name, "CLOSED") async def test_set_server_hostname_implicitly(self): """Client sets server_hostname to the host in the WebSocket URI.""" @@ -918,12 +920,13 @@ async def test_https_proxy_invalid_server_certificate(self): @unittest.skipUnless(hasattr(socket, "AF_UNIX"), "this test requires Unix sockets") class UnixClientTests(IsolatedTrioTestCase): - async def test_connection(self): - """Client connects to server over a Unix socket.""" + async def test_context_manager(self): + """Client connects to Unix server and disconnects automatically.""" with temp_unix_socket_path() as path: with run_unix_server(path): async with unix_connect(path) as client: self.assertEqual(client.protocol.state.name, "OPEN") + self.assertEqual(client.protocol.state.name, "CLOSED") async def test_set_host_header(self): """Client sets the Host header to the host in the WebSocket URI.""" @@ -952,14 +955,15 @@ def redirect(connection, request): "cannot follow cross-origin redirect to ws://other/ with a Unix socket", ) - async def test_secure_connection(self): - """Client connects to server securely over a Unix socket.""" + async def test_secure_context_manager(self): + """Client connects to Unix server securely.""" with temp_unix_socket_path() as path: with run_unix_server(path, ssl=SERVER_CONTEXT): async with unix_connect(path, ssl=CLIENT_CONTEXT) as client: self.assertEqual(client.protocol.state.name, "OPEN") ssl_object = client.stream._ssl_object self.assertEqual(ssl_object.version()[:3], "TLS") + self.assertEqual(client.protocol.state.name, "CLOSED") async def test_set_server_hostname(self): """Client sets server_hostname to the host in the WebSocket URI.""" From 92ca138bb13717469aa52f6b67ff226e17d0d798 Mon Sep 17 00:00:00 2001 From: Aymeric Augustin Date: Tue, 18 Aug 2026 09:20:42 +0200 Subject: [PATCH 4/7] Support overriding host/post in the sync client. --- docs/project/changelog.rst | 6 ++++++ src/websockets/sync/client.py | 10 ++++++---- src/websockets/trio/client.py | 4 ++++ tests/sync/test_client.py | 7 +++++++ 4 files changed, 23 insertions(+), 4 deletions(-) diff --git a/docs/project/changelog.rst b/docs/project/changelog.rst index fa2540c11..ab375b481 100644 --- a/docs/project/changelog.rst +++ b/docs/project/changelog.rst @@ -32,6 +32,12 @@ notice. *In development* +New features +............ + +* :func:`~sync.client.connect` can connect to another host and port than those + specified in the URI in the :mod:`threading` implementation. + .. _17.0.1: 17.0.1 diff --git a/src/websockets/sync/client.py b/src/websockets/sync/client.py index 509eb9df8..85b9a1a56 100644 --- a/src/websockets/sync/client.py +++ b/src/websockets/sync/client.py @@ -229,6 +229,10 @@ def connect( connection handling. Any other keyword arguments are passed to :func:`~socket.create_connection`. + For example, you can set ``address`` to a ``(host, port)`` tuple to connect + to a different host and port from those found in ``uri``. This only changes + the destination of the TCP connection. The host name from ``uri`` is still + used in the TLS handshake for secure connections and in the ``Host`` header. Raises: InvalidURI: If ``uri`` isn't a valid WebSocket URI. @@ -331,11 +335,9 @@ def connect( raise AssertionError("parse_proxy returned unsupported proxy") else: # proxy is None + kwargs.setdefault("address", (ws_uri.host, ws_uri.port)) kwargs.setdefault("timeout", deadline.timeout()) - sock = socket.create_connection( - (ws_uri.host, ws_uri.port), - **kwargs, - ) + sock = socket.create_connection(**kwargs) sock.settimeout(None) diff --git a/src/websockets/trio/client.py b/src/websockets/trio/client.py index 1bdde69c8..9c92d2f67 100644 --- a/src/websockets/trio/client.py +++ b/src/websockets/trio/client.py @@ -220,6 +220,10 @@ class connect: connection handling. Any other keyword arguments are passed to :func:`~trio.open_tcp_stream`. + For example, you can set ``host`` and ``port`` to connect to a different + host and port from those found in ``uri``. This only changes the destination + of the TCP connection. The host name from ``uri`` is still used in the TLS + handshake for secure connections and in the ``Host`` header. Raises: InvalidURI: If ``uri`` isn't a valid WebSocket URI. diff --git a/tests/sync/test_client.py b/tests/sync/test_client.py index fd299dace..2e94aeef6 100644 --- a/tests/sync/test_client.py +++ b/tests/sync/test_client.py @@ -48,6 +48,13 @@ def test_direct_connection(self): self.addCleanup(client.close) self.assertEqual(client.protocol.state.name, "OPEN") + def test_explicit_host_port(self): + """Client connects using an explicit host / port.""" + with run_server() as server: + address = get_host_port(server) + with connect("ws://overridden/", address=address) as client: + self.assertEqual(client.protocol.state.name, "OPEN") + def test_existing_socket(self): """Client connects using a pre-existing socket.""" with run_server() as server: From c0139941a45ed6038d02bfc2386d6a7b56e91750 Mon Sep 17 00:00:00 2001 From: Aymeric Augustin Date: Tue, 18 Aug 2026 15:17:04 +0200 Subject: [PATCH 5/7] Deprecate calling connect() directly. This usage was never shown in examples, but it was clearly described in the documentation and it must exist in the wild. Add a flag for backwards compatibility. While the name implies that it should go away in the long run, it isn't immediately deprecated. Its name provide light deterrence, which is enough for now. --- src/websockets/asyncio/client.py | 11 ++++----- src/websockets/sync/client.py | 35 +++++++++++++++++++++++---- src/websockets/sync/connection.py | 9 +++++++ src/websockets/trio/client.py | 17 ++++++++------ tests/sync/test_client.py | 39 ++++++++++++++++++++++++++++--- 5 files changed, 91 insertions(+), 20 deletions(-) diff --git a/src/websockets/asyncio/client.py b/src/websockets/asyncio/client.py index 6c6c11495..f49af72da 100644 --- a/src/websockets/asyncio/client.py +++ b/src/websockets/asyncio/client.py @@ -177,10 +177,9 @@ class connect: """ Connect to the WebSocket server at ``uri``. - This coroutine returns a :class:`ClientConnection` instance, which you can - use to send and receive messages. - - :func:`connect` may be used as an asynchronous context manager:: + :func:`connect` should be treated as an asynchronous context manager + yielding a :class:`ClientConnection`, which can then receive and send + messages:: from websockets.asyncio.client import connect @@ -189,8 +188,8 @@ class connect: The connection is closed automatically when exiting the context. - :func:`connect` can be used as an infinite asynchronous iterator to - reconnect automatically on errors:: + :func:`connect` can also be treated as an infinite asynchronous iterator + to reconnect automatically on errors:: async for websocket in connect(...): try: diff --git a/src/websockets/sync/client.py b/src/websockets/sync/client.py index 85b9a1a56..6aebf027a 100644 --- a/src/websockets/sync/client.py +++ b/src/websockets/sync/client.py @@ -65,6 +65,7 @@ def __init__( ) -> None: self.protocol: ClientProtocol self.response_rcvd = threading.Event() + self.pending_legacy_warning = True super().__init__( sock, protocol, @@ -74,6 +75,21 @@ def __init__( max_queue=max_queue, ) + def __enter__(self) -> ClientConnection: + self.pending_legacy_warning = False + return super().__enter__() + + def maybe_raise_legacy_warning(self) -> None: + if self.pending_legacy_warning: + self.pending_legacy_warning = False + warnings.warn( # deprecated in 17.1 + "connect() must be used as a context manager: " + "with connect(...) as websocket: ...; alternatively, use " + "websocket = connect(..., legacy=True) to connect directly", + DeprecationWarning, + stacklevel=3, + ) + def handshake( self, additional_headers: HeadersLike | None = None, @@ -164,10 +180,8 @@ def connect( """ Connect to the WebSocket server at ``uri``. - This function returns a :class:`ClientConnection` instance, which you can - use to send and receive messages. - - :func:`connect` may be used as a context manager:: + :func:`connect` should be treated as a context manager yielding a + :class:`ClientConnection`, which can then receive and send messages:: from websockets.sync.client import connect @@ -176,6 +190,13 @@ def connect( The connection is closed automatically when exiting the context. + For backwards compatibility, :func:`connect` may be called directly:: + + websocket = await connect(..., legacy=True) + + In that case, you're responsible for closing the connection with + :meth:`ClientConnection.close` when no longer needed. + Args: uri: URI of the WebSocket server. sock: Preexisting TCP socket. ``sock`` overrides the host and port @@ -253,6 +274,9 @@ def connect( DeprecationWarning, ) + # Backwards compatibility: connect can return a ClientConnection. + legacy = kwargs.pop("legacy", False) + ws_uri = parse_uri(uri) if not ws_uri.secure and ssl is not None: raise ValueError("ssl argument is incompatible with a ws:// URI") @@ -389,6 +413,9 @@ def connect( sock.close() raise + if legacy: + connection.pending_legacy_warning = False + try: connection.handshake( additional_headers, diff --git a/src/websockets/sync/connection.py b/src/websockets/sync/connection.py index 74ac997ab..b2b7415cb 100644 --- a/src/websockets/sync/connection.py +++ b/src/websockets/sync/connection.py @@ -317,6 +317,7 @@ def recv(self, timeout: float | None = None, decode: bool | None = None) -> Data :meth:`recv_streaming` concurrently. """ + self.maybe_raise_legacy_warning() try: return self.recv_messages.get(timeout, decode) except EOFError: @@ -387,6 +388,7 @@ def recv_streaming(self, decode: bool | None = None) -> Iterator[Data]: :meth:`recv_streaming` concurrently. """ + self.maybe_raise_legacy_warning() try: yield from self.recv_messages.get_iter(decode) return @@ -466,6 +468,7 @@ def send( TypeError: If ``message`` doesn't have a supported type. """ + self.maybe_raise_legacy_warning() # Unfragmented message — this case must be handled first because # strings and bytes-like objects are iterable. @@ -591,6 +594,7 @@ def close( reason: WebSocket close reason. """ + self.maybe_raise_legacy_warning() try: # The context manager takes care of waiting for the TCP connection # to terminate after calling a method that sends a close frame. @@ -647,6 +651,7 @@ def ping( the corresponding pong wasn't received yet. """ + self.maybe_raise_legacy_warning() if isinstance(data, BytesLike): data = bytes(data) elif isinstance(data, str): @@ -684,6 +689,7 @@ def pong(self, data: DataLike = b"") -> None: ConnectionClosed: When the connection is closed. """ + self.maybe_raise_legacy_warning() if isinstance(data, BytesLike): data = bytes(data) elif isinstance(data, str): @@ -696,6 +702,9 @@ def pong(self, data: DataLike = b"") -> None: # Private methods + def maybe_raise_legacy_warning(self) -> None: + pass # see override in ClientConnection + def process_event(self, event: Event) -> None: """ Process one incoming event. diff --git a/src/websockets/trio/client.py b/src/websockets/trio/client.py index 9c92d2f67..18a348dc9 100644 --- a/src/websockets/trio/client.py +++ b/src/websockets/trio/client.py @@ -135,10 +135,9 @@ class connect: """ Connect to the WebSocket server at ``uri``. - This coroutine returns a :class:`ClientConnection` instance, which you can - use to send and receive messages. - - :func:`connect` may be used as an asynchronous context manager:: + :func:`connect` is designed to be called as an asynchronous context manager + yielding a :class:`ClientConnection`, which you can then use to receive and + send messages:: from websockets.trio.client import connect @@ -147,8 +146,8 @@ class connect: The connection is closed automatically when exiting the context. - :func:`connect` can be used as an infinite asynchronous iterator to - reconnect automatically on errors:: + :func:`connect` can also be treated as an infinite asynchronous iterator + to reconnect automatically on errors:: async for websocket in connect(...): try: @@ -162,6 +161,10 @@ class connect: The connection is closed automatically after each iteration of the loop. + :func:`connect` cannot be awaited directly. This is because it runs a task + to manage the connection and Trio doesn't support spawning tasks without a + context that ensures completion. + Args: uri: URI of the WebSocket server. stream: Preexisting TCP stream. ``stream`` overrides the host and port @@ -536,7 +539,7 @@ async def connect(self, nursery: trio.Nursery) -> ClientConnection: # Re-raise exception with an informative error message. raise TimeoutError("timed out during opening handshake") from exc - # Do not define __await__ for... = await nursery.start(connect, ...) + # Do not define __await__ for ... = await nursery.start(connect, ...) # because it doesn't look idiomatic in Trio. # async with connect(...) as ...: ... diff --git a/tests/sync/test_client.py b/tests/sync/test_client.py index 2e94aeef6..214b13b42 100644 --- a/tests/sync/test_client.py +++ b/tests/sync/test_client.py @@ -44,7 +44,7 @@ def test_context_manager(self): def test_direct_connection(self): """Client connects to server directly.""" with run_server() as server: - client = connect(get_uri(server)) + client = connect(get_uri(server), legacy=True) self.addCleanup(client.close) self.assertEqual(client.protocol.state.name, "OPEN") @@ -321,7 +321,7 @@ def test_context_manager(self): def test_direct_connection(self): """Client connects to server directly.""" with run_server(ssl=SERVER_CONTEXT) as server: - client = connect(get_uri(server), ssl=CLIENT_CONTEXT) + client = connect(get_uri(server), ssl=CLIENT_CONTEXT, legacy=True) self.addCleanup(client.close) self.assertEqual(client.protocol.state.name, "OPEN") self.assertEqual(client.socket.version()[:3], "TLS") @@ -673,7 +673,7 @@ def test_direct_connection(self): """Client connects to Unix server directly.""" with temp_unix_socket_path() as path: with run_unix_server(path): - client = unix_connect(path) + client = unix_connect(path, legacy=True) self.addCleanup(client.close) self.assertEqual(client.protocol.state.name, "OPEN") @@ -793,3 +793,36 @@ def test_ssl_context_argument(self): with self.assertDeprecationWarning("ssl_context was renamed to ssl"): with connect(get_uri(server), ssl_context=CLIENT_CONTEXT): pass + + def test_direct_connection_without_legacy_flag(self): + """Client connects to server without legacy=True.""" + with run_server() as server: + client = connect(get_uri(server)) + self.addCleanup(client.close) + self.assertEqual(client.protocol.state.name, "OPEN") + # First call of a public API triggers a warning + with self.assertDeprecationWarning( + "connect() must be used as a context manager: " + "with connect(...) as websocket: ...; alternatively, use " + "websocket = connect(..., legacy=True) to connect directly" + ): + client.ping() + # Later calls don't trigger a warning + client.pong() + + def test_direct_unix_connection_without_legacy_flag(self): + """Client connects to Unix server without legacy=True.""" + with temp_unix_socket_path() as path: + with run_unix_server(path): + client = unix_connect(path) + self.addCleanup(client.close) + self.assertEqual(client.protocol.state.name, "OPEN") + # First call of a public API triggers a warning + with self.assertDeprecationWarning( + "connect() must be used as a context manager: " + "with connect(...) as websocket: ...; alternatively, use " + "websocket = connect(..., legacy=True) to connect directly" + ): + client.ping() + # Later calls don't trigger a warning + client.pong() From d21a39ba50821effd7f1aae12d130f5dfaebbb5c Mon Sep 17 00:00:00 2001 From: Aymeric Augustin Date: Wed, 19 Aug 2026 09:01:42 +0200 Subject: [PATCH 6/7] Follow redirects in the sync implementation. --- Makefile | 2 +- docs/project/changelog.rst | 3 + docs/reference/features.rst | 2 +- src/websockets/asyncio/client.py | 14 +- src/websockets/sync/client.py | 544 +++++++++++++++++++++--------- src/websockets/sync/connection.py | 3 - src/websockets/sync/server.py | 2 +- src/websockets/trio/client.py | 19 +- tests/asyncio/test_client.py | 9 +- tests/sync/test_client.py | 262 +++++++++++++- tests/sync/test_router.py | 28 +- tests/trio/test_client.py | 28 ++ 12 files changed, 700 insertions(+), 216 deletions(-) diff --git a/Makefile b/Makefile index 36346977a..2ab9be262 100644 --- a/Makefile +++ b/Makefile @@ -24,7 +24,7 @@ types: mypy --strict src tests: - python -m unittest + python -m unittest tests/sync/test_client.py -v coverage: coverage run --source src/websockets,tests -m unittest diff --git a/docs/project/changelog.rst b/docs/project/changelog.rst index ab375b481..d4c821723 100644 --- a/docs/project/changelog.rst +++ b/docs/project/changelog.rst @@ -35,6 +35,9 @@ notice. New features ............ +* :func:`~sync.client.connect` now follows redirects in the :mod:`threading` + implementation. + * :func:`~sync.client.connect` can connect to another host and port than those specified in the URI in the :mod:`threading` implementation. diff --git a/docs/reference/features.rst b/docs/reference/features.rst index 2bc505b5e..0a6c2f614 100644 --- a/docs/reference/features.rst +++ b/docs/reference/features.rst @@ -161,7 +161,7 @@ Client +------------------------------------+--------+--------+--------+--------+--------+ | Connect to non-ASCII IRIs | ✅ | ✅ | ✅ | ✅ | ✅ | +------------------------------------+--------+--------+--------+--------+--------+ - | Follow HTTP redirects | ✅ | ❌ | ✅ | — | ✅ | + | Follow HTTP redirects | ✅ | ✅ | ✅ | — | ✅ | +------------------------------------+--------+--------+--------+--------+--------+ | Perform HTTP Basic Authentication | ✅ | ✅ | ✅ | ✅ | ✅ | +------------------------------------+--------+--------+--------+--------+--------+ diff --git a/src/websockets/asyncio/client.py b/src/websockets/asyncio/client.py index f49af72da..7f6d8c9d8 100644 --- a/src/websockets/asyncio/client.py +++ b/src/websockets/asyncio/client.py @@ -554,13 +554,7 @@ def process_redirect(self, exc: Exception) -> Exception | str: return new_uri - # ... = await connect(...) - - def __await__(self) -> Generator[Any, None, ClientConnection]: - # Create a suitable iterator by calling __await__ on a coroutine. - return self.__await_impl__().__await__() - - async def __await_impl__(self) -> ClientConnection: + async def connect(self) -> ClientConnection: try: async with asyncio.timeout(self.open_timeout): for _ in range(MAX_REDIRECTS): @@ -605,6 +599,12 @@ async def __await_impl__(self) -> ClientConnection: # Re-raise exception with an informative error message. raise TimeoutError("timed out during opening handshake") from exc + # ... = await connect(...) + + def __await__(self) -> Generator[Any, None, ClientConnection]: + # Create a suitable iterator by calling __await__ on a coroutine. + return self.connect().__await__() + # async with connect(...) as ...: ... async def __aenter__(self) -> ClientConnection: diff --git a/src/websockets/sync/client.py b/src/websockets/sync/client.py index 6aebf027a..a158e9d30 100644 --- a/src/websockets/sync/client.py +++ b/src/websockets/sync/client.py @@ -1,16 +1,26 @@ from __future__ import annotations import logging +import os import socket import ssl as ssl_module import threading +import urllib.parse import warnings from collections.abc import Sequence +from types import TracebackType from typing import Any, Callable, Literal, TypeVar, cast +from ..asyncio.client import process_exception from ..client import ClientProtocol -from ..datastructures import HeadersLike -from ..exceptions import InvalidProxyMessage, InvalidProxyStatus, ProxyError +from ..datastructures import Headers, HeadersLike +from ..exceptions import ( + InvalidProxyMessage, + InvalidProxyStatus, + InvalidStatus, + ProxyError, + SecurityError, +) from ..extensions.base import ClientExtensionFactory from ..extensions.permessage_deflate import enable_client_permessage_deflate from ..headers import validate_subprotocols @@ -26,6 +36,8 @@ __all__ = ["connect", "unix_connect", "ClientConnection"] +MAX_REDIRECTS = int(os.environ.get("WEBSOCKETS_MAX_REDIRECTS", "10")) + class ClientConnection(Connection): """ @@ -144,6 +156,335 @@ def recv_events(self) -> None: self.response_rcvd.set() +class _connect: + def __init__( + self, + uri: str, + *, + # TCP/TLS + sock: socket.socket | None = None, + ssl: ssl_module.SSLContext | None = None, + server_hostname: str | None = None, + # WebSocket + origin: Origin | None = None, + extensions: Sequence[ClientExtensionFactory] | None = None, + subprotocols: Sequence[Subprotocol] | None = None, + compression: str | None = "deflate", + # HTTP + additional_headers: HeadersLike | None = None, + user_agent_header: str | None = USER_AGENT, + proxy: str | Literal[True] | None = True, + proxy_ssl: ssl_module.SSLContext | None = None, + proxy_server_hostname: str | None = None, + # Timeouts + open_timeout: float | None = 10, + ping_interval: float | None = 20, + ping_timeout: float | None = 20, + close_timeout: float | None = 10, + # Limits + max_size: int | None | tuple[int | None, int | None] = 2**20, + max_queue: int | None | tuple[int | None, int | None] = 16, + # Logging + logger: LoggerLike | None = None, + # Escape hatch for advanced customization + create_connection: type[ClientConnection] | None = None, + # Other keyword arguments are passed to socket.create_connection + **kwargs: Any, + ) -> None: + # Backwards compatibility: ssl used to be called ssl_context. + if ssl is None and "ssl_context" in kwargs: + ssl = kwargs.pop("ssl_context") + warnings.warn( # deprecated in 13.0 - 2024-08-20 + "ssl_context was renamed to ssl", + DeprecationWarning, + ) + + self.uri = uri + self.ws_uri = parse_uri(uri) + if not self.ws_uri.secure and ssl is not None: + raise ValueError("ssl argument is incompatible with a ws:// URI") + + if subprotocols is not None: + validate_subprotocols(subprotocols) + + if compression == "deflate": + extensions = enable_client_permessage_deflate(extensions) + elif compression is not None: + raise ValueError(f"unsupported compression: {compression}") + + if logger is None: + logger = logging.getLogger("websockets.client") + + if create_connection is None: + create_connection = ClientConnection + + self.sock = sock + self.ssl = ssl + self.server_hostname = server_hostname + self.additional_headers = additional_headers + self.user_agent_header = user_agent_header + self.proxy = proxy + self.proxy_ssl = proxy_ssl + self.proxy_server_hostname = proxy_server_hostname + self.process_exception = process_exception + self.open_timeout = open_timeout + self.logger = logger + self.create_connection = create_connection + self.open_socket_kwargs = kwargs + self.protocol_kwargs = dict( + origin=origin, + extensions=extensions, + subprotocols=subprotocols, + max_size=max_size, + logger=logger, + ) + self.connection_kwargs = dict( + ping_interval=ping_interval, + ping_timeout=ping_timeout, + close_timeout=close_timeout, + max_queue=max_queue, + ) + + def open_socket(self, deadline: Deadline) -> socket.socket: + """Open a TCP or Unix connection to the server, possibly through a proxy.""" + kwargs = self.open_socket_kwargs.copy() + unix = kwargs.pop("unix", False) + + proxy = self.proxy + if unix: + proxy = None + if proxy is True: + proxy = get_proxy(self.ws_uri) + + if unix: + sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + try: + sock.settimeout(deadline.timeout()) + sock.connect(kwargs.pop("path")) + except Exception: + sock.close() + raise + + elif proxy is not None: + proxy_parsed = parse_proxy(proxy) + + if proxy_parsed.scheme[:5] == "socks": + sock = connect_socks_proxy( + proxy_parsed, + self.ws_uri, + deadline, + # websockets is consistent with the socket module while + # python_socks is consistent across implementations. + local_addr=kwargs.pop("source_address", None), + ) + + elif proxy_parsed.scheme[:4] == "http": + if proxy_parsed.scheme != "https" and self.proxy_ssl is not None: + raise ValueError( + "proxy_ssl argument is incompatible with an http:// proxy" + ) + sock = connect_http_proxy( + proxy_parsed, + self.ws_uri, + deadline, + user_agent_header=self.user_agent_header, + ssl=self.proxy_ssl, + server_hostname=self.proxy_server_hostname, + **kwargs, + ) + + else: + raise AssertionError("parse_proxy returned unsupported proxy") + + else: # proxy is None + kwargs.setdefault("address", (self.ws_uri.host, self.ws_uri.port)) + kwargs.setdefault("timeout", deadline.timeout()) + sock = socket.create_connection(**kwargs) + + sock.settimeout(None) + return sock + + def enable_tls(self, sock: socket.socket, deadline: Deadline) -> socket.socket: + """Enable TLS on the connection.""" + if self.ssl is None: + ssl = ssl_module.create_default_context() + else: + ssl = self.ssl + if self.server_hostname is None: + server_hostname = self.ws_uri.host + else: + server_hostname = self.server_hostname + sock.settimeout(deadline.timeout()) + if self.proxy_ssl is None: + sock = ssl.wrap_socket(sock, server_hostname=server_hostname) + else: + sock_2 = SSLSSLSocket(sock, ssl, server_hostname=server_hostname) + # Let's pretend that sock is a socket, even though it isn't. + sock = cast(socket.socket, sock_2) + sock.settimeout(None) + return sock + + def open_connection(self, deadline: Deadline) -> ClientConnection: + """Create a WebSocket connection.""" + if self.sock is None: + sock = self.open_socket(deadline) + else: + sock = self.sock + + try: + if sock.family in {socket.AF_INET, socket.AF_INET6}: + sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, True) + + if self.ws_uri.secure: + sock = self.enable_tls(sock, deadline) + + protocol = ClientProtocol( + self.ws_uri, + **self.protocol_kwargs, # type: ignore + ) + + # self.create_connection defaults to ClientConnection. + connection = self.create_connection( + sock, + protocol, + **self.connection_kwargs, # type: ignore + ) + + except Exception: + sock.close() + raise + + try: + connection.handshake( + self.additional_headers, + self.user_agent_header, + deadline.timeout(), + ) + except Exception: + connection.close_socket() + connection.recv_events_thread.join() + raise + + return connection + + def process_redirect(self, exc: Exception) -> Exception | str: + """ + Determine whether a connection error is a redirect that can be followed. + + Return the new URI if it's a valid redirect. Else, return an exception. + + """ + if not ( + isinstance(exc, InvalidStatus) + and exc.response.status_code + in [ + 300, # Multiple Choices + 301, # Moved Permanently + 302, # Found + 303, # See Other + 307, # Temporary Redirect + 308, # Permanent Redirect + ] + and "Location" in exc.response.headers + ): + return exc + + old_ws_uri = self.ws_uri + new_uri = urllib.parse.urljoin(self.uri, exc.response.headers["Location"]) + new_ws_uri = parse_uri(new_uri) + + # If connect() received a socket, it is closed and cannot be reused. + if self.sock is not None: + return ValueError( + f"cannot follow redirect to {new_uri} with a preexisting socket" + ) + + # TLS downgrade is forbidden. + if old_ws_uri.secure and not new_ws_uri.secure: + return SecurityError(f"cannot follow redirect to non-secure URI {new_uri}") + + # Apply restrictions to cross-origin redirects. + if ( + old_ws_uri.secure != new_ws_uri.secure + or old_ws_uri.host != new_ws_uri.host + or old_ws_uri.port != new_ws_uri.port + ): + # Cross-origin redirects on Unix sockets don't quite make sense. + if self.open_socket_kwargs.get("unix", False): + return ValueError( + f"cannot follow cross-origin redirect to {new_uri} " + f"with a Unix socket" + ) + # Cross-origin redirects when host and port are overridden are ill-defined. + if self.open_socket_kwargs.get("address") is not None: + return ValueError( + f"cannot follow cross-origin redirect to {new_uri} " + f"with an explicit host and port" + ) + + # Strip credentials to avoid leaking them to a different origin. + if self.additional_headers is not None: + self.additional_headers = Headers( + ( + (key, value) + for key, value in Headers(self.additional_headers).raw_items() + if key.lower() + not in ["authorization", "cookie", "proxy-authorization"] + ) + ) + + return new_uri + + def connect(self) -> ClientConnection: + """Connect to a WebSocket server, following redirects.""" + # Calculate timeouts on the TCP, TLS, and WebSocket handshakes. + # The TCP and TLS timeouts must be set on the socket, then removed + # to avoid conflicting with the WebSocket timeout in handshake(). + deadline = Deadline(self.open_timeout) + + for _ in range(MAX_REDIRECTS): + try: + connection = self.open_connection(deadline) + except Exception as exc: + exc_or_uri = self.process_redirect(exc) + if isinstance(exc_or_uri, Exception): + # Response isn't a valid redirect; raise the exception. + if exc_or_uri is exc: + raise + else: + raise exc_or_uri from exc + else: + # Response is a valid redirect; follow it. + self.uri = exc_or_uri + self.ws_uri = parse_uri(exc_or_uri) + continue + + else: + connection.start_keepalive() + return connection + else: + raise SecurityError(f"more than {MAX_REDIRECTS} redirects") + + # with connect(...) as ...: ... + + def __enter__(self) -> ClientConnection: + if hasattr(self, "connection"): + raise RuntimeError("connect() isn't reentrant") + self.connection = self.connect() + return self.connection + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + exc_traceback: TracebackType | None, + ) -> None: + try: + self.connection.close() + finally: + del self.connection + + def connect( uri: str, *, @@ -263,171 +604,41 @@ def connect( TimeoutError: If the opening handshake times out. """ + # Backwards compatibility: connect() can return a ClientConnection. + legacy: bool | None = kwargs.pop("legacy", None) + + connecter = _connect( + uri, + sock=sock, + ssl=ssl, + server_hostname=server_hostname, + origin=origin, + extensions=extensions, + subprotocols=subprotocols, + compression=compression, + additional_headers=additional_headers, + user_agent_header=user_agent_header, + proxy=proxy, + proxy_ssl=proxy_ssl, + proxy_server_hostname=proxy_server_hostname, + open_timeout=open_timeout, + ping_interval=ping_interval, + ping_timeout=ping_timeout, + close_timeout=close_timeout, + max_size=max_size, + max_queue=max_queue, + logger=logger, + create_connection=create_connection, + **kwargs, + ) - # Process parameters - - # Backwards compatibility: ssl used to be called ssl_context. - if ssl is None and "ssl_context" in kwargs: - ssl = kwargs.pop("ssl_context") - warnings.warn( # deprecated in 13.0 - 2024-08-20 - "ssl_context was renamed to ssl", - DeprecationWarning, - ) - - # Backwards compatibility: connect can return a ClientConnection. - legacy = kwargs.pop("legacy", False) - - ws_uri = parse_uri(uri) - if not ws_uri.secure and ssl is not None: - raise ValueError("ssl argument is incompatible with a ws:// URI") - - if subprotocols is not None: - validate_subprotocols(subprotocols) - - if compression == "deflate": - extensions = enable_client_permessage_deflate(extensions) - elif compression is not None: - raise ValueError(f"unsupported compression: {compression}") - - if logger is None: - logger = logging.getLogger("websockets.client") - - if create_connection is None: - create_connection = ClientConnection - - # Private APIs for unix_connect() - unix: bool = kwargs.pop("unix", False) - path: str | None = kwargs.pop("path", None) - - if unix: - if path is None and sock is None: - raise ValueError("missing path argument") - elif path is not None and sock is not None: - raise ValueError("path is incompatible with sock") - - if unix: - proxy = None - if sock is not None: - proxy = None - if proxy is True: - proxy = get_proxy(ws_uri) - - # Calculate timeouts on the TCP, TLS, and WebSocket handshakes. - # The TCP and TLS timeouts must be set on the socket, then removed - # to avoid conflicting with the WebSocket timeout in handshake(). - deadline = Deadline(open_timeout) - - try: - # Connect socket - - if sock is None: - if unix: - sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) - sock.settimeout(deadline.timeout()) - assert path is not None # mypy cannot figure this out - sock.connect(path) - - elif proxy is not None: - proxy_parsed = parse_proxy(proxy) - - if proxy_parsed.scheme[:5] == "socks": - sock = connect_socks_proxy( - proxy_parsed, - ws_uri, - deadline, - # websockets is consistent with the socket module while - # python_socks is consistent across implementations. - local_addr=kwargs.pop("source_address", None), - ) - - elif proxy_parsed.scheme[:4] == "http": - if proxy_parsed.scheme != "https" and proxy_ssl is not None: - raise ValueError( - "proxy_ssl argument is incompatible with an http:// proxy" - ) - sock = connect_http_proxy( - proxy_parsed, - ws_uri, - deadline, - user_agent_header=user_agent_header, - ssl=proxy_ssl, - server_hostname=proxy_server_hostname, - **kwargs, - ) - - else: - raise AssertionError("parse_proxy returned unsupported proxy") - - else: # proxy is None - kwargs.setdefault("address", (ws_uri.host, ws_uri.port)) - kwargs.setdefault("timeout", deadline.timeout()) - sock = socket.create_connection(**kwargs) - - sock.settimeout(None) - - # Disable Nagle algorithm - - if not unix: - sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, True) - - # Initialize TLS wrapper and perform TLS handshake - - if ws_uri.secure: - if ssl is None: - ssl = ssl_module.create_default_context() - if server_hostname is None: - server_hostname = ws_uri.host - sock.settimeout(deadline.timeout()) - if proxy_ssl is None: - sock = ssl.wrap_socket(sock, server_hostname=server_hostname) - else: - sock_2 = SSLSSLSocket(sock, ssl, server_hostname=server_hostname) - # Let's pretend that sock is a socket, even though it isn't. - sock = cast(socket.socket, sock_2) - sock.settimeout(None) - - # Initialize WebSocket protocol - - protocol = ClientProtocol( - ws_uri, - origin=origin, - extensions=extensions, - subprotocols=subprotocols, - max_size=max_size, - logger=logger, - ) - - # Initialize WebSocket connection - - # create_connection defaults to ClientConnection. - connection = create_connection( - sock, - protocol, - ping_interval=ping_interval, - ping_timeout=ping_timeout, - close_timeout=close_timeout, - max_queue=max_queue, - ) - except Exception: - if sock is not None: - sock.close() - raise + # Forwards compatibility: this will be the default behavior in the future. + if legacy is False: + return connecter # type: ignore + connection = connecter.connect() if legacy: connection.pending_legacy_warning = False - - try: - connection.handshake( - additional_headers, - user_agent_header, - deadline.timeout(), - ) - except Exception: - connection.close_socket() - connection.recv_events_thread.join() - raise - - connection.start_keepalive() return connection @@ -452,12 +663,19 @@ def unix_connect( ``wss://localhost/``. """ + sock = kwargs.get("sock") + if path is None and sock is None: + raise ValueError("missing path argument") + elif path is not None and sock is not None: + raise ValueError("path is incompatible with sock") + if uri is None: # Backwards compatibility: ssl used to be called ssl_context. if kwargs.get("ssl") is None and kwargs.get("ssl_context") is None: uri = "ws://localhost/" else: uri = "wss://localhost/" + return connect(uri=uri, unix=True, path=path, **kwargs) diff --git a/src/websockets/sync/connection.py b/src/websockets/sync/connection.py index b2b7415cb..8af3812a0 100644 --- a/src/websockets/sync/connection.py +++ b/src/websockets/sync/connection.py @@ -594,7 +594,6 @@ def close( reason: WebSocket close reason. """ - self.maybe_raise_legacy_warning() try: # The context manager takes care of waiting for the TCP connection # to terminate after calling a method that sends a close frame. @@ -651,7 +650,6 @@ def ping( the corresponding pong wasn't received yet. """ - self.maybe_raise_legacy_warning() if isinstance(data, BytesLike): data = bytes(data) elif isinstance(data, str): @@ -689,7 +687,6 @@ def pong(self, data: DataLike = b"") -> None: ConnectionClosed: When the connection is closed. """ - self.maybe_raise_legacy_warning() if isinstance(data, BytesLike): data = bytes(data) elif isinstance(data, str): diff --git a/src/websockets/sync/server.py b/src/websockets/sync/server.py index 0d1c19fbe..97a8d9d46 100644 --- a/src/websockets/sync/server.py +++ b/src/websockets/sync/server.py @@ -686,7 +686,7 @@ def sock_handler(sock: socket.socket, addr: Any) -> None: try: # Disable Nagle algorithm - if not unix: + if sock.family in {socket.AF_INET, socket.AF_INET6}: sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, True) # Perform TLS handshake diff --git a/src/websockets/trio/client.py b/src/websockets/trio/client.py index 18a348dc9..8b72355b1 100644 --- a/src/websockets/trio/client.py +++ b/src/websockets/trio/client.py @@ -326,15 +326,16 @@ def __init__( async def open_tcp_stream(self) -> trio.abc.Stream: """Open a TCP or Unix connection to the server, possibly through a proxy.""" kwargs = self.open_tcp_stream_kwargs.copy() + unix = kwargs.pop("unix", False) proxy = self.proxy - if kwargs.get("unix", False): + if unix: proxy = None if proxy is True: proxy = get_proxy(self.ws_uri) - if kwargs.pop("unix", False): - return await trio.open_unix_socket(kwargs["path"]) + if unix: + return await trio.open_unix_socket(kwargs.pop("path")) elif proxy is not None: proxy_parsed = parse_proxy(proxy) @@ -391,7 +392,6 @@ async def enable_tls(self, stream: trio.abc.Stream) -> trio.abc.Stream: async def open_connection(self, nursery: trio.Nursery) -> ClientConnection: """Create a WebSocket connection.""" - # TCP connection is already established. if self.stream is None: stream = await self.open_tcp_stream() else: @@ -419,8 +419,6 @@ async def open_connection(self, nursery: trio.Nursery) -> ClientConnection: self.user_agent_header, ) - return connection - except trio.Cancelled: await trio.aclose_forcefully(stream) # The nursery running this coroutine was canceled. @@ -434,6 +432,8 @@ async def open_connection(self, nursery: trio.Nursery) -> ClientConnection: await trio.aclose_forcefully(stream) raise + return connection + def process_redirect(self, exc: Exception) -> Exception | str: """ Determine whether a connection error is a redirect that can be followed. @@ -660,11 +660,18 @@ def unix_connect( ``wss://localhost/``. """ + stream = kwargs.get("stream") + if path is None and stream is None: + raise ValueError("missing path argument") + elif path is not None and stream is not None: + raise ValueError("path is incompatible with stream") + if uri is None: if kwargs.get("ssl") is None: uri = "ws://localhost/" else: uri = "wss://localhost/" + return connect(uri=uri, unix=True, path=path, **kwargs) diff --git a/tests/asyncio/test_client.py b/tests/asyncio/test_client.py index b011a1fd8..66fa7f815 100644 --- a/tests/asyncio/test_client.py +++ b/tests/asyncio/test_client.py @@ -993,6 +993,13 @@ async def test_set_server_hostname(self): ssl_object = client.transport.get_extra_info("ssl_object") self.assertEqual(ssl_object.server_hostname, "overridden") + async def test_non_existing_path(self): + """Client attempts to connect to a non-existing Unix socket path.""" + with temp_unix_socket_path() as path: + with self.assertRaises(FileNotFoundError): + async with unix_connect(path + ".doesnotexist"): + self.fail("did not raise") + class ClientUsageErrorsTests(unittest.IsolatedAsyncioTestCase): async def test_ssl_without_secure_uri(self): @@ -1026,7 +1033,7 @@ async def test_proxy_ssl_without_https_proxy(self): "proxy_ssl argument is incompatible with an http:// proxy", ) - async def test_https_proxy_without_ssl(self): + async def test_https_proxy_without_proxy_ssl(self): """Client rejects proxy_ssl=None when proxy is HTTPS.""" with self.assertRaises(ValueError) as raised: await connect( diff --git a/tests/sync/test_client.py b/tests/sync/test_client.py index 214b13b42..f28d0846c 100644 --- a/tests/sync/test_client.py +++ b/tests/sync/test_client.py @@ -1,3 +1,4 @@ +import contextlib import http import logging import os @@ -18,6 +19,7 @@ InvalidStatus, InvalidURI, ProxyError, + SecurityError, ) from websockets.extensions.permessage_deflate import PerMessageDeflate from websockets.sync.client import * @@ -33,6 +35,18 @@ from .server import get_host_port, get_uri, run_server, run_unix_server +@contextlib.contextmanager +def few_redirects(): + from websockets.sync import client + + max_redirects = client.MAX_REDIRECTS + client.MAX_REDIRECTS = 2 + try: + yield + finally: + client.MAX_REDIRECTS = max_redirects + + class ClientTests(unittest.TestCase): def test_context_manager(self): """Client connects to server and disconnects automatically.""" @@ -142,6 +156,153 @@ def create_connection(*args, **kwargs): ) as client: self.assertTrue(client.create_connection_ran) + def test_redirect(self): + """Client follows redirect.""" + + def redirect(connection, request): + if request.path == "/redirect": + response = connection.respond(http.HTTPStatus.FOUND, "") + response.headers["Location"] = "/" + return response + + with run_server(process_request=redirect) as server: + with connect(get_uri(server) + "/redirect") as client: + self.assertEqual(client.protocol.uri.path, "/") + + def test_cross_origin_redirect(self): + """Client follows redirect to a secure URI on a different origin.""" + + def redirect(connection, request): + response = connection.respond(http.HTTPStatus.FOUND, "") + response.headers["Location"] = get_uri(other_server) + return response + + with run_server(process_request=redirect) as server: + with run_server() as other_server: + with connect(get_uri(server)): + self.assertFalse(server.connections) + self.assertTrue(other_server.connections) + + @few_redirects() + def test_redirect_limit(self): + """Client stops following redirects after limit is reached.""" + + def redirect(connection, request): + response = connection.respond(http.HTTPStatus.FOUND, "") + response.headers["Location"] = request.path + return response + + with run_server(process_request=redirect) as server: + with self.assertRaises(SecurityError) as raised: + with connect(get_uri(server)): + self.fail("did not raise") + + self.assertEqual( + str(raised.exception), + "more than 2 redirects", + ) + + def test_redirect_with_explicit_host_port(self): + """Client follows redirect with an explicit host / port.""" + + def redirect(connection, request): + if request.path == "/redirect": + response = connection.respond(http.HTTPStatus.FOUND, "") + response.headers["Location"] = "/" + return response + + with run_server(process_request=redirect) as server: + address = get_host_port(server) + with connect("ws://overridden/redirect", address=address) as client: + self.assertEqual(client.protocol.uri.path, "/") + + def test_cross_origin_redirect_with_explicit_host_port(self): + """Client doesn't follow cross-origin redirect with an explicit host / port.""" + + def redirect(connection, request): + response = connection.respond(http.HTTPStatus.FOUND, "") + response.headers["Location"] = "ws://other/" + return response + + with run_server(process_request=redirect) as server: + address = get_host_port(server) + with self.assertRaises(ValueError) as raised: + with connect("ws://overridden/", address=address): + self.fail("did not raise") + + self.assertEqual( + str(raised.exception), + "cannot follow cross-origin redirect to ws://other/ " + "with an explicit host and port", + ) + + def test_redirect_with_existing_socket(self): + """Client doesn't follow redirect when using a pre-existing socket.""" + + def redirect(connection, request): + response = connection.respond(http.HTTPStatus.FOUND, "") + response.headers["Location"] = "/" + return response + + with run_server(process_request=redirect) as server: + with socket.create_connection(get_host_port(server)) as sock: + with self.assertRaises(ValueError) as raised: + # Use a non-existing domain to ensure we connect via sock. + with connect("ws://invalid/redirect", sock=sock): + self.fail("did not raise") + + self.assertEqual( + str(raised.exception), + "cannot follow redirect to ws://invalid/ with a preexisting socket", + ) + + def test_cross_origin_redirect_strips_credentials(self): + """Client strips credentials when following a cross-origin redirect.""" + + def redirect(connection, request): + response = connection.respond(http.HTTPStatus.FOUND, "") + response.headers["Location"] = get_uri(other_server) + return response + + with run_server(process_request=redirect) as server: + with run_server() as other_server: + with connect( + get_uri(server), + additional_headers={ + "Authorization": "Bearer secret", + "Cookie": "session=secret", + "Proxy-Authorization": "Basic secret", + "X-Custom": "keep", + }, + ) as client: + self.assertNotIn("Authorization", client.request.headers) + self.assertNotIn("Cookie", client.request.headers) + self.assertNotIn("Proxy-Authorization", client.request.headers) + self.assertIn("X-Custom", client.request.headers) + + def test_same_origin_redirect_preserves_credentials(self): + """Client preserves credentials when following a same-origin redirect.""" + + def redirect(connection, request): + if request.path == "/redirect": + response = connection.respond(http.HTTPStatus.FOUND, "") + response.headers["Location"] = "/" + return response + + with run_server(process_request=redirect) as server: + with connect( + get_uri(server) + "/redirect", + additional_headers={ + "Authorization": "Bearer secret", + "Cookie": "session=secret", + "Proxy-Authorization": "Basic secret", + "X-Custom": "keep", + }, + ) as client: + self.assertIn("Authorization", client.request.headers) + self.assertIn("Cookie", client.request.headers) + self.assertIn("Proxy-Authorization", client.request.headers) + def test_invalid_uri(self): """Client receives an invalid URI.""" with self.assertRaises(InvalidURI): @@ -370,6 +531,40 @@ def test_reject_invalid_server_hostname(self): str(raised.exception), ) + def test_cross_origin_redirect(self): + """Client follows redirect to a secure URI on a different origin.""" + + def redirect(connection, request): + response = connection.respond(http.HTTPStatus.FOUND, "") + response.headers["Location"] = get_uri(other_server) + return response + + with run_server(ssl=SERVER_CONTEXT, process_request=redirect) as server: + with run_server(ssl=SERVER_CONTEXT) as other_server: + with connect(get_uri(server), ssl=CLIENT_CONTEXT): + self.assertFalse(server.connections) + self.assertTrue(other_server.connections) + + def test_redirect_to_insecure_uri(self): + """Client doesn't follow redirect from secure URI to non-secure URI.""" + + def redirect(connection, request): + response = connection.respond(http.HTTPStatus.FOUND, "") + response.headers["Location"] = insecure_uri + return response + + with run_server(ssl=SERVER_CONTEXT, process_request=redirect) as server: + with self.assertRaises(SecurityError) as raised: + secure_uri = get_uri(server) + insecure_uri = secure_uri.replace("wss://", "ws://") + with connect(secure_uri, ssl=CLIENT_CONTEXT): + self.fail("did not raise") + + self.assertEqual( + str(raised.exception), + f"cannot follow redirect to non-secure URI {insecure_uri}", + ) + @unittest.skipUnless("mitmproxy" in sys.modules, "mitmproxy not installed") class SocksProxyClientTests(ProxyMixin, unittest.TestCase): @@ -685,6 +880,25 @@ def test_set_host_header(self): with unix_connect(path, uri="ws://overridden/") as client: self.assertEqual(client.request.headers["Host"], "overridden") + def test_cross_origin_redirect(self): + """Client doesn't follows redirect to a URI on a different origin.""" + + def redirect(connection, request): + response = connection.respond(http.HTTPStatus.FOUND, "") + response.headers["Location"] = "ws://other/" + return response + + with temp_unix_socket_path() as path: + with run_unix_server(path, process_request=redirect): + with self.assertRaises(ValueError) as raised: + with unix_connect(path): + self.fail("did not raise") + + self.assertEqual( + str(raised.exception), + "cannot follow cross-origin redirect to ws://other/ with a Unix socket", + ) + def test_secure_context_manager(self): """Client connects to Unix server securely and disconnects automatically.""" with temp_unix_socket_path() as path: @@ -698,7 +912,7 @@ def test_secure_direct_connection(self): """Client connects to Unix server securely and directly.""" with temp_unix_socket_path() as path: with run_unix_server(path, ssl=SERVER_CONTEXT): - client = unix_connect(path, ssl=CLIENT_CONTEXT) + client = unix_connect(path, ssl=CLIENT_CONTEXT, legacy=True) self.addCleanup(client.close) self.assertEqual(client.protocol.state.name, "OPEN") self.assertEqual(client.socket.version()[:3], "TLS") @@ -713,6 +927,13 @@ def test_set_server_hostname(self): ) as client: self.assertEqual(client.socket.server_hostname, "overridden") + def test_non_existing_path(self): + """Client attempts to connect to a non-existing Unix socket path.""" + with temp_unix_socket_path() as path: + with self.assertRaises(FileNotFoundError): + with unix_connect(path + ".doesnotexist"): + self.fail("did not raise") + class ClientUsageErrorsTests(unittest.TestCase): def test_ssl_without_secure_uri(self): @@ -737,15 +958,6 @@ def test_proxy_ssl_without_https_proxy(self): "proxy_ssl argument is incompatible with an http:// proxy", ) - def test_unix_without_path_or_sock(self): - """Unix client requires path when sock isn't provided.""" - with self.assertRaises(ValueError) as raised: - unix_connect() - self.assertEqual( - str(raised.exception), - "missing path argument", - ) - def test_unsupported_proxy(self): """Client rejects unsupported proxy.""" with self.assertRaises(InvalidProxy) as raised: @@ -756,6 +968,15 @@ def test_unsupported_proxy(self): "other://localhost:58080 isn't a valid proxy: scheme other isn't supported", ) + def test_unix_without_path_or_sock(self): + """Unix client requires path when sock isn't provided.""" + with self.assertRaises(ValueError) as raised: + unix_connect() + self.assertEqual( + str(raised.exception), + "missing path argument", + ) + def test_unix_with_path_and_sock(self): """Unix client rejects path when sock is provided.""" sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) @@ -785,6 +1006,19 @@ def test_unsupported_compression(self): "unsupported compression: False", ) + def test_reentrancy(self): + """Client isn't reentrant.""" + with run_server() as server: + connecter = connect(get_uri(server), legacy=False) + with connecter: + with self.assertRaises(RuntimeError) as raised: + with connecter: + self.fail("did not raise") + self.assertEqual( + str(raised.exception), + "connect() isn't reentrant", + ) + class BackwardsCompatibilityTests(DeprecationTestCase): def test_ssl_context_argument(self): @@ -806,9 +1040,9 @@ def test_direct_connection_without_legacy_flag(self): "with connect(...) as websocket: ...; alternatively, use " "websocket = connect(..., legacy=True) to connect directly" ): - client.ping() + client.send("2 + 2") # Later calls don't trigger a warning - client.pong() + self.assertEqual(client.recv(), "4") def test_direct_unix_connection_without_legacy_flag(self): """Client connects to Unix server without legacy=True.""" @@ -823,6 +1057,6 @@ def test_direct_unix_connection_without_legacy_flag(self): "with connect(...) as websocket: ...; alternatively, use " "websocket = connect(..., legacy=True) to connect directly" ): - client.ping() + client.send("2 + 2") # Later calls don't trigger a warning - client.pong() + self.assertEqual(client.recv(), "4") diff --git a/tests/sync/test_router.py b/tests/sync/test_router.py index 5c8b4de2d..1d16b7238 100644 --- a/tests/sync/test_router.py +++ b/tests/sync/test_router.py @@ -63,27 +63,17 @@ def test_route_with_query_string(self): def test_redirect(self): """Router redirects connections according to redirect_to.""" - with run_router(url_map, server_name="localhost") as server: - with self.assertRaises(InvalidStatus) as raised: - with connect(get_uri(server) + "/r"): - self.fail("did not raise") - self.assertEqual( - raised.exception.response.headers["Location"], - "ws://localhost/", - ) + with run_router(url_map) as server: + with connect(get_uri(server) + "/r") as client: + self.assertEval(client, "ws.request.path", "/") def test_secure_redirect(self): - """Router redirects connections to a wss:// URI when TLS is enabled.""" - with run_router(url_map, server_name="localhost", ssl=SERVER_CONTEXT) as server: - with self.assertRaises(InvalidStatus) as raised: - with connect(get_uri(server) + "/r", ssl=CLIENT_CONTEXT): - self.fail("did not raise") - self.assertEqual( - raised.exception.response.headers["Location"], - "wss://localhost/", - ) + """Router redirects connections according to redirect_to when TLS is enabled.""" + with run_router(url_map, ssl=SERVER_CONTEXT) as server: + with connect(get_uri(server) + "/r", ssl=CLIENT_CONTEXT) as client: + self.assertEval(client, "ws.request.path", "/") - @patch("websockets.asyncio.client.connect.process_redirect", lambda _, exc: exc) + @patch("websockets.sync.client._connect.process_redirect", lambda _, exc: exc) def test_force_secure_redirect(self): """Router redirects ws:// connections to a wss:// URI when ssl=True.""" with run_router(url_map, ssl=True) as server: @@ -96,7 +86,7 @@ def test_force_secure_redirect(self): redirect_uri + "/", ) - @patch("websockets.asyncio.client.connect.process_redirect", lambda _, exc: exc) + @patch("websockets.sync.client._connect.process_redirect", lambda _, exc: exc) def test_force_redirect_server_name(self): """Router redirects connections to the host declared in server_name.""" with run_router(url_map, server_name="other") as server: diff --git a/tests/trio/test_client.py b/tests/trio/test_client.py index cedd078cf..89269e048 100644 --- a/tests/trio/test_client.py +++ b/tests/trio/test_client.py @@ -978,6 +978,13 @@ async def test_set_server_hostname(self): ssl_object = client.stream._ssl_object self.assertEqual(ssl_object.server_hostname, "overridden") + async def test_non_existing_path(self): + """Client attempts to connect to a non-existing Unix socket path.""" + with temp_unix_socket_path() as path: + with self.assertRaises(FileNotFoundError): + async with unix_connect(path + ".doesnotexist"): + self.fail("did not raise") + class ClientUsageErrorsTests(IsolatedTrioTestCase): async def test_ssl_without_secure_uri(self): @@ -1014,6 +1021,27 @@ async def test_unsupported_proxy(self): "other://localhost:51080 isn't a valid proxy: scheme other isn't supported", ) + async def test_unix_without_path_or_sock(self): + """Unix client requires path when sock isn't provided.""" + with self.assertRaises(ValueError) as raised: + async with unix_connect(): + self.fail("did not raise") + self.assertEqual( + str(raised.exception), + "missing path argument", + ) + + async def test_unix_with_path_and_stream(self): + """Unix client rejects path when stream is provided.""" + stream, _ = trio.testing.memory_stream_pair() + with self.assertRaises(ValueError) as raised: + async with unix_connect(path="/", stream=stream): + self.fail("did not raise") + self.assertEqual( + str(raised.exception), + "path is incompatible with stream", + ) + async def test_invalid_subprotocol(self): """Client rejects single value of subprotocols.""" with self.assertRaises(TypeError) as raised: From 4c020d3c6905cbe3e90f0947f92c6296a5fb50f2 Mon Sep 17 00:00:00 2001 From: Aymeric Augustin Date: Tue, 11 Aug 2026 08:44:45 +0200 Subject: [PATCH 7/7] Support reconnecting in the threading implementation. --- docs/project/changelog.rst | 4 + docs/reference/features.rst | 2 +- docs/reference/sync/client.rst | 6 + src/websockets/sync/client.py | 274 +++++++++++++++++++++++++++++---- tests/sync/test_client.py | 141 ++++++++++++++++- tests/sync/test_router.py | 4 +- 6 files changed, 394 insertions(+), 37 deletions(-) diff --git a/docs/project/changelog.rst b/docs/project/changelog.rst index d4c821723..a7e73ebaf 100644 --- a/docs/project/changelog.rst +++ b/docs/project/changelog.rst @@ -35,6 +35,10 @@ notice. New features ............ +* Added support for reconnecting automatically by using + :func:`~sync.client.reconnect` as an iterator to the :mod:`threading` + implementation. + * :func:`~sync.client.connect` now follows redirects in the :mod:`threading` implementation. diff --git a/docs/reference/features.rst b/docs/reference/features.rst index 0a6c2f614..01909bdb1 100644 --- a/docs/reference/features.rst +++ b/docs/reference/features.rst @@ -149,7 +149,7 @@ Client +------------------------------------+--------+--------+--------+--------+--------+ | Close connection on context exit | ✅ | ✅ | ✅ | — | ✅ | +------------------------------------+--------+--------+--------+--------+--------+ - | Reconnect automatically | ✅ | ❌ | ✅ | — | ✅ | + | Reconnect automatically | ✅ | ✅ | ✅ | — | ✅ | +------------------------------------+--------+--------+--------+--------+--------+ | Configure ``Origin`` header | ✅ | ✅ | ✅ | ✅ | ✅ | +------------------------------------+--------+--------+--------+--------+--------+ diff --git a/docs/reference/sync/client.rst b/docs/reference/sync/client.rst index fdc772ed2..d98f04c58 100644 --- a/docs/reference/sync/client.rst +++ b/docs/reference/sync/client.rst @@ -8,8 +8,14 @@ Opening a connection .. autofunction:: connect +.. autofunction:: reconnect + .. autofunction:: unix_connect +.. autofunction:: unix_reconnect + +.. autofunction:: process_exception + Using a connection ------------------ diff --git a/src/websockets/sync/client.py b/src/websockets/sync/client.py index a158e9d30..01ea78cf1 100644 --- a/src/websockets/sync/client.py +++ b/src/websockets/sync/client.py @@ -5,14 +5,16 @@ import socket import ssl as ssl_module import threading +import time +import traceback import urllib.parse import warnings -from collections.abc import Sequence +from collections.abc import Generator, Iterator, Sequence from types import TracebackType -from typing import Any, Callable, Literal, TypeVar, cast +from typing import Any, Callable, Literal, TypeVar, cast, overload from ..asyncio.client import process_exception -from ..client import ClientProtocol +from ..client import ClientProtocol, backoff from ..datastructures import Headers, HeadersLike from ..exceptions import ( InvalidProxyMessage, @@ -34,7 +36,7 @@ from .utils import Deadline -__all__ = ["connect", "unix_connect", "ClientConnection"] +__all__ = ["connect", "unix_connect", "reconnect", "unix_reconnect", "ClientConnection"] MAX_REDIRECTS = int(os.environ.get("WEBSOCKETS_MAX_REDIRECTS", "10")) @@ -156,7 +158,43 @@ def recv_events(self) -> None: self.response_rcvd.set() -class _connect: +class reconnect: + """ + Similar to :func:`connect`, with support for automatic reconnection. + + :func:`reconnect` can also be treated as an infinite asynchronous iterator + to reconnect automatically on errors:: + + for websocket in reconnect(...): + try: + ... + except websockets.exceptions.ConnectionClosed: + continue + + If the connection fails with a transient error, it is retried with + exponential backoff. If it fails with a fatal error, the exception is + raised, breaking out of the loop. + + The connection is closed automatically after each iteration of the loop. + + :func:`reconnect` accepts the same arguments as :func:`connect`, except the + ``legacy`` flag, and raises the same exceptions. + + .. admonition:: Why is :func:`reconnect` a separate API from :func:`connect`? + :class: tip + + A new API was necessary to maintain backwards compatibility with this + historical behavior of :func:`connect`:: + + websocket = connect(...) + for message in websocket: + ... + + Once the deprecation period elapses, :func:`connect` will be changed to + behave like :func:`reconnect` by default. + + """ + def __init__( self, uri: str, @@ -176,11 +214,13 @@ def __init__( proxy: str | Literal[True] | None = True, proxy_ssl: ssl_module.SSLContext | None = None, proxy_server_hostname: str | None = None, + process_exception: Callable[[Exception], Exception | None] = process_exception, # Timeouts open_timeout: float | None = 10, ping_interval: float | None = 20, ping_timeout: float | None = 20, close_timeout: float | None = 10, + reconnect_delays: Callable[[], Generator[float]] = backoff, # Limits max_size: int | None | tuple[int | None, int | None] = 2**20, max_queue: int | None | tuple[int | None, int | None] = 16, @@ -228,6 +268,7 @@ def __init__( self.proxy_server_hostname = proxy_server_hostname self.process_exception = process_exception self.open_timeout = open_timeout + self.reconnect_delays = reconnect_delays self.logger = logger self.create_connection = create_connection self.open_socket_kwargs = kwargs @@ -349,7 +390,6 @@ def open_connection(self, deadline: Deadline) -> ClientConnection: protocol, **self.connection_kwargs, # type: ignore ) - except Exception: sock.close() raise @@ -458,7 +498,6 @@ def connect(self) -> ClientConnection: self.uri = exc_or_uri self.ws_uri = parse_uri(exc_or_uri) continue - else: connection.start_keepalive() return connection @@ -484,6 +523,123 @@ def __exit__( finally: del self.connection + # for ... in reconnect(...): ... + + def __iter__(self) -> Iterator[ClientConnection]: + delays: Generator[float] | None = None + while True: + try: + with self as connection: + yield connection + except Exception as exc: + # Determine whether the exception is retryable or fatal. + # The API of process_exception is "return an exception or None"; + # "raise an exception" is also supported because it's a frequent + # mistake. It isn't documented in order to keep the API simple. + try: + new_exc = self.process_exception(exc) + except Exception as raised_exc: + new_exc = raised_exc + + # The connection failed with a fatal error. + # Raise the exception and exit the loop. + if new_exc is exc: + raise + if new_exc is not None: + raise new_exc from exc + + # The connection failed with a retryable error. + # Start or continue backoff and reconnect. + if delays is None: + delays = self.reconnect_delays() + delay = next(delays) + self.logger.info( + "connect failed; reconnecting in %.1f seconds: %s", + delay, + traceback.format_exception_only(exc)[0].strip(), + ) + time.sleep(delay) + + else: + # The connection succeeded. Reset backoff. + delays = None + + +@overload +def connect( + uri: str, + *, + # TCP/TLS + sock: socket.socket | None = ..., + ssl: ssl_module.SSLContext | None = ..., + server_hostname: str | None = ..., + # WebSocket + origin: Origin | None = ..., + extensions: Sequence[ClientExtensionFactory] | None = ..., + subprotocols: Sequence[Subprotocol] | None = ..., + compression: str | None = ..., + # HTTP + additional_headers: HeadersLike | None = ..., + user_agent_header: str | None = ..., + proxy: str | Literal[True] | None = ..., + proxy_ssl: ssl_module.SSLContext | None = ..., + proxy_server_hostname: str | None = ..., + # Timeouts + open_timeout: float | None = ..., + ping_interval: float | None = ..., + ping_timeout: float | None = ..., + close_timeout: float | None = ..., + # Limits + max_size: int | None | tuple[int | None, int | None] = ..., + max_queue: int | None | tuple[int | None, int | None] = ..., + # Logging + logger: LoggerLike | None = ..., + # Escape hatch for advanced customization + create_connection: type[ClientConnection] | None = ..., + # Backwards and forwards compatibility + legacy: Literal[True] | None = ..., + # Other keyword arguments are passed to socket.create_connection + **kwargs: Any, +) -> ClientConnection: ... + + +@overload +def connect( + uri: str, + *, + # TCP/TLS + sock: socket.socket | None = ..., + ssl: ssl_module.SSLContext | None = ..., + server_hostname: str | None = ..., + # WebSocket + origin: Origin | None = ..., + extensions: Sequence[ClientExtensionFactory] | None = ..., + subprotocols: Sequence[Subprotocol] | None = ..., + compression: str | None = ..., + # HTTP + additional_headers: HeadersLike | None = ..., + user_agent_header: str | None = ..., + proxy: str | Literal[True] | None = ..., + proxy_ssl: ssl_module.SSLContext | None = ..., + proxy_server_hostname: str | None = ..., + # Timeouts + open_timeout: float | None = ..., + ping_interval: float | None = ..., + ping_timeout: float | None = ..., + close_timeout: float | None = ..., + # Limits + max_size: int | None | tuple[int | None, int | None] = ..., + max_queue: int | None | tuple[int | None, int | None] = ..., + # Logging + logger: LoggerLike | None = ..., + # Escape hatch for advanced customization + create_connection: type[ClientConnection] | None = ..., + # Backwards and forwards compatibility + legacy: Literal[False], + # Other keyword arguments are passed to socket.create_connection + **kwargs: Any, +) -> reconnect: ... + def connect( uri: str, @@ -515,9 +671,11 @@ def connect( logger: LoggerLike | None = None, # Escape hatch for advanced customization create_connection: type[ClientConnection] | None = None, + # Backwards and forwards compatibility + legacy: bool | None = None, # Other keyword arguments are passed to socket.create_connection **kwargs: Any, -) -> ClientConnection: +) -> ClientConnection | reconnect: """ Connect to the WebSocket server at ``uri``. @@ -531,6 +689,8 @@ def connect( The connection is closed automatically when exiting the context. + Use :func:`reconnect` to reconnect automatically on errors. + For backwards compatibility, :func:`connect` may be called directly:: websocket = await connect(..., legacy=True) @@ -538,6 +698,14 @@ def connect( In that case, you're responsible for closing the connection with :meth:`ClientConnection.close` when no longer needed. + When the ``legacy`` flag is enabled, :func:`connect` returns directly a + :class:`ClientConnection` and iterating that connection yields messages. + Currently, this is the default behavior when ``legacy`` isn't specified. + + When the ``legacy`` flag is explicitly disabled, :func:`connect` behaves + like :func:`reconnect`: using it as an iterator returns a new connection + at each iteration, making it easy to reconnect automatically on errors. + Args: uri: URI of the WebSocket server. sock: Preexisting TCP socket. ``sock`` overrides the host and port @@ -586,6 +754,8 @@ def connect( logger: Logger for this client. It defaults to ``logging.getLogger("websockets.client")``. See the :doc:`logging guide <../../topics/logging>` for details. + legacy: Set to :obj:`True` to opt into the historical behavior of + returning a :class:`ClientConnection`, without deprecation warning. create_connection: Factory for the :class:`ClientConnection` managing the connection. Set it to a wrapper or a subclass to customize connection handling. @@ -604,10 +774,7 @@ def connect( TimeoutError: If the opening handshake times out. """ - # Backwards compatibility: connect() can return a ClientConnection. - legacy: bool | None = kwargs.pop("legacy", None) - - connecter = _connect( + connecter = reconnect( uri, sock=sock, ssl=ssl, @@ -631,36 +798,29 @@ def connect( create_connection=create_connection, **kwargs, ) - - # Forwards compatibility: this will be the default behavior in the future. + # For backwards compatibility, connect defaults to the historical behavior. + # For forwards compatibility, the future behavior can be chosen explicitly. if legacy is False: - return connecter # type: ignore - + return connecter connection = connecter.connect() + # Users can opt in to the historical behavior to remain unaffected when the + # future behavior becomes the default. if legacy: connection.pending_legacy_warning = False return connection -def unix_connect( +def unix_reconnect( path: str | None = None, uri: str | None = None, + *, + legacy: bool | None = None, **kwargs: Any, -) -> ClientConnection: +) -> reconnect: """ - Connect to a WebSocket server listening on a Unix socket. - - This function accepts the same keyword arguments as :func:`connect`. + Similar to :func:`unix_connect`, with support for automatic reconnection. - It's only available on Unix. - - It's mainly useful for debugging servers listening on Unix sockets. - - Args: - path: File system path to the Unix socket. - uri: URI of the WebSocket server. ``uri`` defaults to - ``ws://localhost/`` or, when a ``ssl`` is provided, to - ``wss://localhost/``. + Refer to the documentation of :func:`reconnect` for details on its behavior. """ sock = kwargs.get("sock") @@ -676,7 +836,59 @@ def unix_connect( else: uri = "wss://localhost/" - return connect(uri=uri, unix=True, path=path, **kwargs) + return reconnect(uri=uri, unix=True, path=path, **kwargs) + + +@overload +def unix_connect( + path: str | None = ..., + uri: str | None = ..., + *, + legacy: Literal[True] | None = ..., + **kwargs: Any, +) -> ClientConnection: ... + + +@overload +def unix_connect( + path: str | None = ..., + uri: str | None = ..., + *, + legacy: Literal[False], + **kwargs: Any, +) -> reconnect: ... + + +def unix_connect( + *args: Any, legacy: bool | None = None, **kwargs: Any +) -> ClientConnection | reconnect: + """ + Connect to a WebSocket server listening on a Unix socket. + + This function accepts the same keyword arguments as :func:`connect`. + + It's only available on Unix. + + It's mainly useful for debugging servers listening on Unix sockets. + + Args: + path: File system path to the Unix socket. + uri: URI of the WebSocket server. ``uri`` defaults to + ``ws://localhost/`` or, when a ``ssl`` is provided, to + ``wss://localhost/``. + + """ + connecter = unix_reconnect(*args, **kwargs) + # For backwards compatibility, connect defaults to the historical behavior. + # For forwards compatibility, the future behavior can be chosen explicitly. + if legacy is False: + return connecter + connection = connecter.connect() + # Users can opt in to the historical behavior to remain unaffected when the + # future behavior becomes the default. + if legacy: + connection.pending_legacy_warning = False + return connection try: diff --git a/tests/sync/test_client.py b/tests/sync/test_client.py index f28d0846c..88d54ab65 100644 --- a/tests/sync/test_client.py +++ b/tests/sync/test_client.py @@ -11,6 +11,7 @@ import unittest from unittest.mock import patch +from websockets.client import backoff from websockets.exceptions import ( InvalidHandshake, InvalidMessage, @@ -35,6 +36,16 @@ from .server import get_host_port, get_uri, run_server, run_unix_server +def short_backoff(): + defaults = backoff.__defaults__ + yield from backoff( + defaults[0] * MS, + defaults[1] * MS, + defaults[2] * MS, + defaults[3], + ) + + @contextlib.contextmanager def few_redirects(): from websockets.sync import client @@ -156,6 +167,101 @@ def create_connection(*args, **kwargs): ) as client: self.assertTrue(client.create_connection_ran) + def test_reconnect(self): + """Client reconnects to server.""" + iterations = 0 + successful = 0 + + def process_request(connection, request): + nonlocal iterations + iterations += 1 + # Retriable errors + if iterations == 1: + time.sleep(3 * MS) + elif iterations == 2: + connection.socket.close() + elif iterations == 3: + return connection.respond(http.HTTPStatus.SERVICE_UNAVAILABLE, "🚒") + # Fatal error + elif iterations == 6: + return connection.respond(http.HTTPStatus.PAYMENT_REQUIRED, "💸") + + with run_server(process_request=process_request) as server: + with self.assertRaises(InvalidStatus) as raised: + for client in reconnect( + get_uri(server), + open_timeout=3 * MS, + reconnect_delays=short_backoff, + ): + self.assertEqual(client.protocol.state.name, "OPEN") + successful += 1 + + self.assertEqual( + str(raised.exception), + "server rejected WebSocket connection: HTTP 402", + ) + self.assertEqual(iterations, 6) + self.assertEqual(successful, 2) + + def test_reconnect_with_custom_process_exception(self): + """Client runs process_exception to tell if errors are retryable or fatal.""" + iteration = 0 + + def process_request(connection, request): + nonlocal iteration + iteration += 1 + if iteration == 1: + return connection.respond(http.HTTPStatus.SERVICE_UNAVAILABLE, "🚒") + return connection.respond(http.HTTPStatus.IM_A_TEAPOT, "🫖") + + def process_exception(exc): + if isinstance(exc, InvalidStatus): + if 500 <= exc.response.status_code < 600: + return None + if exc.response.status_code == 418: + return Exception("🫖 💔 ☕️") + self.fail("unexpected exception") + + with run_server(process_request=process_request) as server: + with self.assertRaises(Exception) as raised: + for _ in reconnect( + get_uri(server), + process_exception=process_exception, + reconnect_delays=short_backoff, + ): + self.fail("did not raise") + + self.assertEqual(iteration, 2) + self.assertEqual( + str(raised.exception), + "🫖 💔 ☕️", + ) + + def test_reconnect_with_custom_process_exception_raising_exception(self): + """Client supports raising an exception in process_exception.""" + + def process_request(connection, request): + return connection.respond(http.HTTPStatus.IM_A_TEAPOT, "🫖") + + def process_exception(exc): + if isinstance(exc, InvalidStatus) and exc.response.status_code == 418: + raise Exception("🫖 💔 ☕️") + self.fail("unexpected exception") + + with run_server(process_request=process_request) as server: + with self.assertRaises(Exception) as raised: + for _ in reconnect( + get_uri(server), + process_exception=process_exception, + reconnect_delays=short_backoff, + ): + self.fail("did not raise") + + self.assertEqual( + str(raised.exception), + "🫖 💔 ☕️", + ) + def test_redirect(self): """Client follows redirect.""" @@ -948,11 +1054,12 @@ def test_ssl_without_secure_uri(self): def test_proxy_ssl_without_https_proxy(self): """Client rejects proxy_ssl when proxy isn't HTTPS.""" with self.assertRaises(ValueError) as raised: - connect( + with connect( "ws://localhost/", proxy="http://localhost:8080", proxy_ssl=CLIENT_CONTEXT, - ) + ): + self.fail("did not raise") self.assertEqual( str(raised.exception), "proxy_ssl argument is incompatible with an http:// proxy", @@ -1009,7 +1116,7 @@ def test_unsupported_compression(self): def test_reentrancy(self): """Client isn't reentrant.""" with run_server() as server: - connecter = connect(get_uri(server), legacy=False) + connecter = reconnect(get_uri(server)) with connecter: with self.assertRaises(RuntimeError) as raised: with connecter: @@ -1028,6 +1135,34 @@ def test_ssl_context_argument(self): with connect(get_uri(server), ssl_context=CLIENT_CONTEXT): pass + def test_set_legacy_flag_explicitly(self): + """Client connects to server with legacy=True.""" + with run_server() as server: + client = connect(get_uri(server), legacy=True) + self.addCleanup(client.close) + self.assertIsInstance(client, ClientConnection) + + def test_unset_legacy_flag_explicitly(self): + """Client connects to server with legacy=False.""" + with run_server() as server: + client = connect(get_uri(server), legacy=False) + self.assertIsInstance(client, reconnect) + + def test_unix_set_legacy_flag_explicitly(self): + """Client connects to server with legacy=True.""" + with temp_unix_socket_path() as path: + with run_unix_server(path): + client = unix_connect(path, legacy=True) + self.addCleanup(client.close) + self.assertIsInstance(client, ClientConnection) + + def test_unix_unset_legacy_flag_explicitly(self): + """Client connects to server with legacy=False.""" + with temp_unix_socket_path() as path: + with run_unix_server(path): + client = unix_connect(path, legacy=False) + self.assertIsInstance(client, reconnect) + def test_direct_connection_without_legacy_flag(self): """Client connects to server without legacy=True.""" with run_server() as server: diff --git a/tests/sync/test_router.py b/tests/sync/test_router.py index 1d16b7238..4d2b2ad99 100644 --- a/tests/sync/test_router.py +++ b/tests/sync/test_router.py @@ -73,7 +73,7 @@ def test_secure_redirect(self): with connect(get_uri(server) + "/r", ssl=CLIENT_CONTEXT) as client: self.assertEval(client, "ws.request.path", "/") - @patch("websockets.sync.client._connect.process_redirect", lambda _, exc: exc) + @patch("websockets.sync.client.reconnect.process_redirect", lambda _, exc: exc) def test_force_secure_redirect(self): """Router redirects ws:// connections to a wss:// URI when ssl=True.""" with run_router(url_map, ssl=True) as server: @@ -86,7 +86,7 @@ def test_force_secure_redirect(self): redirect_uri + "/", ) - @patch("websockets.sync.client._connect.process_redirect", lambda _, exc: exc) + @patch("websockets.sync.client.reconnect.process_redirect", lambda _, exc: exc) def test_force_redirect_server_name(self): """Router redirects connections to the host declared in server_name.""" with run_router(url_map, server_name="other") as server: