From 08db166aaccea01ff7445061d1b18df02fafa731 Mon Sep 17 00:00:00 2001
From: qcha0 <agnewee@gmail.com>
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 <agnewee@gmail.com>
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 <qinchao@uyunsoft.cn>
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 <qinchao@uyunsoft.cn>
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 <qcha0@outlook.com>
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 <qcha0@outlook.com>
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 <qcha0@outlook.com>
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