Skip to content

Commit 98ed271

Browse files
aaugustinaaugustin
committed
Chain causes to ConnectionClosed consistently.
1 parent d3c1ce1 commit 98ed271

8 files changed

Lines changed: 172 additions & 86 deletions

File tree

‎docs/project/changelog.rst‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,12 @@ Security fix
4848

4949
:meth:`~asyncio.connection.Connection.recv_streaming` isn't affected.
5050

51+
Improvements
52+
............
53+
54+
* Chained consistently to :exc:`~exceptions.ConnectionClosed` exceptions that
55+
close the connection.
56+
5157
.. _17.3:
5258

5359
17.3

‎src/websockets/asyncio/connection.py‎

Lines changed: 30 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -116,9 +116,9 @@ def __init__(
116116
# Task that sends keepalive pings. None when ping_interval is None.
117117
self.keepalive_task: asyncio.Task[None] | None = None
118118

119-
# Exception raised while reading from the connection, to be chained to
120-
# ConnectionClosed in order to show why the TCP connection dropped.
121-
self.recv_exc: BaseException | None = None
119+
# Exception that caused the connection to drop, to chain to
120+
# ConnectionClosed.
121+
self.connection_exc: BaseException | None = None
122122

123123
# Completed when the TCP connection is closed and the WebSocket
124124
# connection state becomes CLOSED.
@@ -305,10 +305,12 @@ async def recv(self, decode: bool | None = None) -> Data:
305305
"cannot call recv while another coroutine "
306306
"is already running recv or recv_streaming"
307307
) from None
308-
except TooManyFragments:
308+
except TooManyFragments as exc:
309+
self.set_connection_exc(exc)
309310
async with self.send_context():
310311
self.protocol.fail(CloseCode.MESSAGE_TOO_BIG, "too many fragments")
311312
except UnicodeDecodeError as exc:
313+
self.set_connection_exc(exc)
312314
async with self.send_context():
313315
self.protocol.fail(
314316
CloseCode.INVALID_DATA,
@@ -318,7 +320,7 @@ async def recv(self, decode: bool | None = None) -> Data:
318320

319321
# Wait for the protocol state to be CLOSED before accessing close_exc.
320322
await asyncio.shield(self.connection_lost_waiter)
321-
raise self.protocol.close_exc from self.recv_exc
323+
raise self.close_exc
322324

323325
@overload
324326
def recv_streaming(self, decode: Literal[True]) -> AsyncIterator[str]: ...
@@ -386,6 +388,7 @@ async def recv_streaming(self, decode: bool | None = None) -> AsyncIterator[Data
386388
"is already running recv or recv_streaming"
387389
) from None
388390
except UnicodeDecodeError as exc:
391+
self.set_connection_exc(exc)
389392
async with self.send_context():
390393
self.protocol.fail(
391394
CloseCode.INVALID_DATA,
@@ -395,7 +398,7 @@ async def recv_streaming(self, decode: bool | None = None) -> AsyncIterator[Data
395398

396399
# Wait for the protocol state to be CLOSED before accessing close_exc.
397400
await asyncio.shield(self.connection_lost_waiter)
398-
raise self.protocol.close_exc from self.recv_exc
401+
raise self.close_exc
399402

400403
async def send(
401404
self,
@@ -885,9 +888,6 @@ async def send_context(
885888
wait_for_close = False
886889
# Should we close the transport and raise ConnectionClosed?
887890
raise_close_exc = False
888-
# What exception should we chain ConnectionClosed to?
889-
original_exc: BaseException | None = None
890-
891891
if self.protocol.state is expected_state:
892892
# Let the caller interact with the protocol.
893893
try:
@@ -902,7 +902,7 @@ async def send_context(
902902
# connection and raise the exception to the caller.
903903
wait_for_close = False
904904
raise_close_exc = True
905-
original_exc = exc
905+
self.set_connection_exc(exc)
906906
else:
907907
# Check if the connection is expected to close soon.
908908
if self.protocol.close_expected():
@@ -927,7 +927,7 @@ async def send_context(
927927
# other exceptions would be treated identically.
928928
wait_for_close = False
929929
raise_close_exc = True
930-
original_exc = exc
930+
self.set_connection_exc(exc)
931931

932932
else: # self.protocol.state is not expected_state
933933
# Minor layering violation: we assume that the connection
@@ -946,22 +946,20 @@ async def send_context(
946946
async with asyncio.timeout_at(self.close_deadline):
947947
await asyncio.shield(self.connection_lost_waiter)
948948
except TimeoutError:
949-
# There's no risk of overwriting another error because
950-
# original_exc is never set when wait_for_close is True.
951-
assert original_exc is None
952-
original_exc = TimeoutError("timed out while closing connection")
953-
# Set recv_exc before closing the transport in order to get
954-
# proper exception reporting.
949+
# Set connection_exc before closing the transport in order to
950+
# get proper exception reporting.
955951
raise_close_exc = True
956-
self.set_recv_exc(original_exc)
952+
self.set_connection_exc(
953+
TimeoutError("timed out while closing connection")
954+
)
957955

958956
# If an error occurred, close the transport to terminate the connection and
959957
# raise an exception.
960958
if raise_close_exc:
961959
self.transport.abort()
962960
# Wait for the protocol state to be CLOSED before accessing close_exc.
963961
await asyncio.shield(self.connection_lost_waiter)
964-
raise self.protocol.close_exc from original_exc
962+
raise self.close_exc
965963

966964
def send_data(self) -> None:
967965
"""
@@ -988,15 +986,20 @@ def send_data(self) -> None:
988986
self.logger.debug("x closing TCP connection")
989987
self.transport.close()
990988

991-
def set_recv_exc(self, exc: BaseException | None) -> None:
992-
"""
993-
Set recv_exc, if not set yet.
989+
@property
990+
def close_exc(self) -> ConnectionClosed:
991+
exc = self.protocol.close_exc
992+
if self.connection_exc is not None:
993+
exc.__cause__ = self.connection_exc
994+
return exc
994995

995-
This method must be called only from connection callbacks.
996+
def set_connection_exc(self, exc: BaseException | None) -> None:
997+
"""
998+
Set connection_exc, if not set yet.
996999
9971000
"""
998-
if self.recv_exc is None:
999-
self.recv_exc = exc
1001+
if self.connection_exc is None:
1002+
self.connection_exc = exc
10001003

10011004
# asyncio.Protocol methods
10021005

@@ -1022,7 +1025,7 @@ def connection_lost(self, exc: Exception | None) -> None:
10221025
self.protocol.receive_eof()
10231026
assert self.protocol.state is CLOSED
10241027

1025-
self.set_recv_exc(exc)
1028+
self.set_connection_exc(exc)
10261029

10271030
# Abort recv() and pending pings with a ConnectionClosed exception.
10281031
self.recv_messages.close()
@@ -1096,7 +1099,7 @@ def data_received(self, data: bytes) -> None:
10961099
except Exception as exc:
10971100
if self.debug:
10981101
self.logger.debug("! error while sending data", exc_info=True)
1099-
self.set_recv_exc(exc)
1102+
self.set_connection_exc(exc)
11001103

11011104
# If needed, set the close deadline based on the close timeout.
11021105
if self.protocol.close_expected():

‎src/websockets/protocol.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -155,6 +155,8 @@ def __init__(
155155
self.writes: list[bytes] = []
156156
self.parser = self.parse()
157157
next(self.parser) # start coroutine
158+
159+
# Exception that caused the parser to stop, chained to ConnectionClosed.
158160
self.parser_exc: Exception | None = None
159161

160162
@property

‎src/websockets/sync/connection.py‎

Lines changed: 36 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -127,9 +127,9 @@ def __init__(
127127
# Thread that sends keepalive pings. None when ping_interval is None.
128128
self.keepalive_thread: threading.Thread | None = None
129129

130-
# Exception raised while reading from the connection, to be chained to
131-
# ConnectionClosed in order to show why the TCP connection dropped.
132-
self.recv_exc: BaseException | None = None
130+
# Exception that caused the connection to drop, to chain to
131+
# ConnectionClosed.
132+
self.connection_exc: BaseException | None = None
133133

134134
# Notification to wake up recv_events() when the connection is closed.
135135
self.shutdown_watcher, self.shutdown_notifier = socket.socketpair()
@@ -334,10 +334,14 @@ def recv(self, timeout: float | None = None, decode: bool | None = None) -> Data
334334
"cannot call recv while another thread "
335335
"is already running recv or recv_streaming"
336336
) from None
337-
except TooManyFragments:
337+
except TooManyFragments as exc:
338+
with self.protocol_mutex:
339+
self.set_connection_exc(exc)
338340
with self.send_context():
339341
self.protocol.fail(CloseCode.MESSAGE_TOO_BIG, "too many fragments")
340342
except UnicodeDecodeError as exc:
343+
with self.protocol_mutex:
344+
self.set_connection_exc(exc)
341345
with self.send_context():
342346
self.protocol.fail(
343347
CloseCode.INVALID_DATA,
@@ -347,7 +351,7 @@ def recv(self, timeout: float | None = None, decode: bool | None = None) -> Data
347351

348352
# Wait for the protocol state to be CLOSED before accessing close_exc.
349353
self.recv_events_thread.join()
350-
raise self.protocol.close_exc from self.recv_exc
354+
raise self.close_exc
351355

352356
@overload
353357
def recv_streaming(self, decode: Literal[True]) -> Iterator[str]: ...
@@ -410,6 +414,8 @@ def recv_streaming(self, decode: bool | None = None) -> Iterator[Data]:
410414
"is already running recv or recv_streaming"
411415
) from None
412416
except UnicodeDecodeError as exc:
417+
with self.protocol_mutex:
418+
self.set_connection_exc(exc)
413419
with self.send_context():
414420
self.protocol.fail(
415421
CloseCode.INVALID_DATA,
@@ -419,7 +425,7 @@ def recv_streaming(self, decode: bool | None = None) -> Iterator[Data]:
419425

420426
# Wait for the protocol state to be CLOSED before accessing close_exc.
421427
self.recv_events_thread.join()
422-
raise self.protocol.close_exc from self.recv_exc
428+
raise self.close_exc
423429

424430
def send(
425431
self,
@@ -890,10 +896,10 @@ def recv_events(self) -> None:
890896
)
891897
# When the closing handshake is initiated by our side,
892898
# select() may block until send_context() closes the socket.
893-
# In that case, send_context() already set recv_exc.
894-
# Calling set_recv_exc() avoids overwriting it.
899+
# In that case, send_context() already set connection_exc.
900+
# Calling set_connection_exc() avoids overwriting it.
895901
with self.protocol_mutex:
896-
self.set_recv_exc(exc)
902+
self.set_connection_exc(exc)
897903
break
898904

899905
if data == b"":
@@ -920,7 +926,7 @@ def recv_events(self) -> None:
920926
# set by send_context(), in case of a race condition
921927
# i.e. send_context() closes the socket after select()
922928
# returns above but before send_data() calls send().
923-
self.set_recv_exc(exc)
929+
self.set_connection_exc(exc)
924930
break
925931

926932
# If needed, set the close deadline based on the close timeout.
@@ -966,7 +972,7 @@ def recv_events(self) -> None:
966972
# This branch should never run. It's a safety net in case of bugs.
967973
self.logger.error("unexpected internal error", exc_info=True)
968974
with self.protocol_mutex:
969-
self.set_recv_exc(exc)
975+
self.set_connection_exc(exc)
970976
finally:
971977
# This isn't expected to raise an exception.
972978
self.close_socket()
@@ -999,9 +1005,6 @@ def send_context(
9991005
wait_for_close = False
10001006
# Should we close the socket and raise ConnectionClosed?
10011007
raise_close_exc = False
1002-
# What exception should we chain ConnectionClosed to?
1003-
original_exc: BaseException | None = None
1004-
10051008
# Acquire the protocol lock.
10061009
with self.protocol_mutex:
10071010
if self.protocol.state is expected_state:
@@ -1018,7 +1021,7 @@ def send_context(
10181021
# connection and raise the exception to the caller.
10191022
wait_for_close = False
10201023
raise_close_exc = True
1021-
original_exc = exc
1024+
self.set_connection_exc(exc)
10221025
else:
10231026
# Check if the connection is expected to close soon.
10241027
if self.protocol.close_expected():
@@ -1042,7 +1045,7 @@ def send_context(
10421045
# other exceptions would be treated identically.
10431046
wait_for_close = False
10441047
raise_close_exc = True
1045-
original_exc = exc
1048+
self.set_connection_exc(exc)
10461049

10471050
else: # self.protocol.state is not expected_state
10481051
# Minor layering violation: we assume that the connection
@@ -1064,23 +1067,21 @@ def send_context(
10641067
timeout = self.close_deadline.timeout(raise_if_elapsed=False)
10651068
self.recv_events_thread.join(timeout)
10661069
if self.recv_events_thread.is_alive():
1067-
# There's no risk of overwriting another error because
1068-
# original_exc is never set when wait_for_close is True.
1069-
assert original_exc is None
1070-
original_exc = TimeoutError("timed out while closing connection")
1071-
# Set recv_exc before closing the socket in order to get
1070+
# Set connection_exc before closing the socket in order to get
10721071
# proper exception reporting.
10731072
raise_close_exc = True
10741073
with self.protocol_mutex:
1075-
self.set_recv_exc(original_exc)
1074+
self.set_connection_exc(
1075+
TimeoutError("timed out while closing connection")
1076+
)
10761077

10771078
# If an error occurred, close the socket to terminate the connection and
10781079
# raise an exception.
10791080
if raise_close_exc:
10801081
self.close_socket()
10811082
# Wait for the protocol state to be CLOSED before accessing close_exc.
10821083
self.recv_events_thread.join()
1083-
raise self.protocol.close_exc from original_exc
1084+
raise self.close_exc
10841085

10851086
def send_data(self) -> None:
10861087
"""
@@ -1101,17 +1102,23 @@ def send_data(self) -> None:
11011102
except OSError: # socket already closed
11021103
pass
11031104

1104-
def set_recv_exc(self, exc: BaseException | None) -> None:
1105+
@property
1106+
def close_exc(self) -> ConnectionClosed:
1107+
exc = self.protocol.close_exc
1108+
if self.connection_exc is not None:
1109+
exc.__cause__ = self.connection_exc
1110+
return exc
1111+
1112+
def set_connection_exc(self, exc: BaseException | None) -> None:
11051113
"""
1106-
Set recv_exc, if not set yet.
1114+
Set connection_exc, if not set yet.
11071115
1108-
This method requires holding protocol_mutex and must be called only from
1109-
the thread running recv_events().
1116+
This method requires holding protocol_mutex.
11101117
11111118
"""
11121119
assert self.protocol_mutex.locked()
1113-
if self.recv_exc is None:
1114-
self.recv_exc = exc
1120+
if self.connection_exc is None:
1121+
self.connection_exc = exc
11151122

11161123
def close_socket(self) -> None:
11171124
"""

0 commit comments

Comments
 (0)