From 08db166aaccea01ff7445061d1b18df02fafa731 Mon Sep 17 00:00:00 2001 From: qcha0 Date: Sat, 13 Nov 2021 12:49:50 +0800 Subject: [PATCH 1/7] avoid self reference --- tornado/websocket.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tornado/websocket.py b/tornado/websocket.py index fbfd700887..14239b8b50 100644 --- a/tornado/websocket.py +++ b/tornado/websocket.py @@ -1439,6 +1439,10 @@ def close(self, code: Optional[int] = None, reason: Optional[str] = None) -> Non def on_connection_close(self) -> None: if not self.connect_future.done(): self.connect_future.set_exception(StreamClosedError()) + else: + # avoid self reference + self.connect_future = Future() + self.connect_future.set_result(None) self._on_message(None) self.tcp_client.close() super().on_connection_close() From b57a1ba52f75c302490f594f1bc0e90a2d4c48f8 Mon Sep 17 00:00:00 2001 From: qcha0 Date: Sat, 13 Nov 2021 12:51:20 +0800 Subject: [PATCH 2/7] add test case --- tornado/test/websocket_test.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/tornado/test/websocket_test.py b/tornado/test/websocket_test.py index 4d39f37046..21e8079d40 100644 --- a/tornado/test/websocket_test.py +++ b/tornado/test/websocket_test.py @@ -483,6 +483,14 @@ def test_write_after_close(self): with self.assertRaises(WebSocketClosedError): ws.write_message("hello") + @gen_test + def test_reference_self_after_been_closed(self): + ws = yield websocket_connect( + "ws://127.0.0.1:%d/close_reason" % self.get_http_port()) + msg = yield ws.read_message() + self.assertIs(msg, None) + self.assertIs(ws.connect_future.result(), None) + @gen_test def test_async_prepare(self): # Previously, an async prepare method triggered a bug that would From 468ce2c6cf920ad6ac96a422b4ffbf39a0e392b9 Mon Sep 17 00:00:00 2001 From: qinchao Date: Tue, 10 Oct 2023 17:11:33 +0800 Subject: [PATCH 3/7] set connect_future to None after connection close --- tornado/websocket.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/tornado/websocket.py b/tornado/websocket.py index 14239b8b50..2e1eed4728 100644 --- a/tornado/websocket.py +++ b/tornado/websocket.py @@ -1439,13 +1439,10 @@ def close(self, code: Optional[int] = None, reason: Optional[str] = None) -> Non def on_connection_close(self) -> None: if not self.connect_future.done(): self.connect_future.set_exception(StreamClosedError()) - else: - # avoid self reference - self.connect_future = Future() - self.connect_future.set_result(None) self._on_message(None) self.tcp_client.close() super().on_connection_close() + self.connect_future = None def on_ws_connection_close( self, close_code: Optional[int] = None, close_reason: Optional[str] = None From 84cc2a1be9470d3f134affcfd7a31ae257aad81f Mon Sep 17 00:00:00 2001 From: qinchao Date: Tue, 10 Oct 2023 17:11:45 +0800 Subject: [PATCH 4/7] update test case --- tornado/test/websocket_test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tornado/test/websocket_test.py b/tornado/test/websocket_test.py index 21e8079d40..8cdaf6a317 100644 --- a/tornado/test/websocket_test.py +++ b/tornado/test/websocket_test.py @@ -489,7 +489,7 @@ def test_reference_self_after_been_closed(self): "ws://127.0.0.1:%d/close_reason" % self.get_http_port()) msg = yield ws.read_message() self.assertIs(msg, None) - self.assertIs(ws.connect_future.result(), None) + self.assertIs(ws.connect_future, None) @gen_test def test_async_prepare(self): From e4926339a85b23db9ecf31a5b2b19096b478d339 Mon Sep 17 00:00:00 2001 From: qcha0 Date: Tue, 10 Oct 2023 17:40:49 +0800 Subject: [PATCH 5/7] reformat websocket_test.py --- tornado/test/websocket_test.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tornado/test/websocket_test.py b/tornado/test/websocket_test.py index 8cdaf6a317..7427c70a9d 100644 --- a/tornado/test/websocket_test.py +++ b/tornado/test/websocket_test.py @@ -486,7 +486,8 @@ def test_write_after_close(self): @gen_test def test_reference_self_after_been_closed(self): ws = yield websocket_connect( - "ws://127.0.0.1:%d/close_reason" % self.get_http_port()) + "ws://127.0.0.1:%d/close_reason" % self.get_http_port() + ) msg = yield ws.read_message() self.assertIs(msg, None) self.assertIs(ws.connect_future, None) From 5296bd8752b4131e0f85391a58d359dc60a07c6d Mon Sep 17 00:00:00 2001 From: qcha0 Date: Tue, 10 Oct 2023 18:13:47 +0800 Subject: [PATCH 6/7] fix connect_future types in assignment --- tornado/websocket.py | 19 ++++++++++++------- 1 file changed, 12 insertions(+), 7 deletions(-) diff --git a/tornado/websocket.py b/tornado/websocket.py index 2e1eed4728..573258a434 100644 --- a/tornado/websocket.py +++ b/tornado/websocket.py @@ -1360,7 +1360,9 @@ def __init__( subprotocols: Optional[List[str]] = None, resolver: Optional[Resolver] = None, ) -> None: - self.connect_future = Future() # type: Future[WebSocketClientConnection] + self.connect_future = ( + Future() + ) # type: Union[Future[WebSocketClientConnection], None] self.read_queue = Queue(1) # type: Queue[Union[None, str, bytes]] self.key = base64.b64encode(os.urandom(16)) self._on_message_callback = on_message_callback @@ -1437,7 +1439,7 @@ def close(self, code: Optional[int] = None, reason: Optional[str] = None) -> Non self.protocol = None # type: ignore def on_connection_close(self) -> None: - if not self.connect_future.done(): + if self.connect_future and not self.connect_future.done(): self.connect_future.set_exception(StreamClosedError()) self._on_message(None) self.tcp_client.close() @@ -1452,7 +1454,7 @@ def on_ws_connection_close( self.on_connection_close() def _on_http_response(self, response: httpclient.HTTPResponse) -> None: - if not self.connect_future.done(): + if self.connect_future and not self.connect_future.done(): if response.error: self.connect_future.set_exception(response.error) else: @@ -1488,7 +1490,8 @@ async def headers_received( # ability to see exceptions. self.final_callback = None # type: ignore - future_set_result_unless_cancelled(self.connect_future, self) + if self.connect_future: + future_set_result_unless_cancelled(self.connect_future, self) def write_message( self, message: Union[str, bytes, Dict[str, Any]], binary: bool = False @@ -1665,6 +1668,8 @@ def websocket_connect( subprotocols=subprotocols, resolver=resolver, ) - if callback is not None: - IOLoop.current().add_future(conn.connect_future, callback) - return conn.connect_future + if conn.connect_future: + if callback is not None: + IOLoop.current().add_future(conn.connect_future, callback) + return conn.connect_future + raise WebSocketError("Initialize websocket client") From 3b4bab6f0258ebf513f8e2a3d2e19286fae50bc0 Mon Sep 17 00:00:00 2001 From: qcha0 Date: Tue, 10 Oct 2023 18:43:19 +0800 Subject: [PATCH 7/7] fix test case --- tornado/test/websocket_test.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/tornado/test/websocket_test.py b/tornado/test/websocket_test.py index 7427c70a9d..47314d6d07 100644 --- a/tornado/test/websocket_test.py +++ b/tornado/test/websocket_test.py @@ -485,11 +485,7 @@ def test_write_after_close(self): @gen_test def test_reference_self_after_been_closed(self): - ws = yield websocket_connect( - "ws://127.0.0.1:%d/close_reason" % self.get_http_port() - ) - msg = yield ws.read_message() - self.assertIs(msg, None) + ws = yield self.ws_connect("/close_reason") self.assertIs(ws.connect_future, None) @gen_test