Skip to content

Commit 28854a4

Browse files
Fix ConnectionResetError not being raised when the transport is close… (#7199)
Co-authored-by: J. Nick Koston <nick@koston.org>
1 parent 565cc21 commit 28854a4

6 files changed

Lines changed: 44 additions & 13 deletions

File tree

‎CHANGES/7180.bugfix‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
``ConnectionResetError`` will always be raised when ``StreamWriter.write`` is called after ``connection_lost`` has been called on the ``BaseProtocol``

‎aiohttp/base_protocol.py‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -18,11 +18,15 @@ def __init__(self, loop: asyncio.AbstractEventLoop) -> None:
1818
self._loop: asyncio.AbstractEventLoop = loop
1919
self._paused = False
2020
self._drain_waiter: Optional[asyncio.Future[None]] = None
21-
self._connection_lost = False
2221
self._reading_paused = False
2322

2423
self.transport: Optional[asyncio.Transport] = None
2524

25+
@property
26+
def connected(self) -> bool:
27+
"""Return True if the connection is open."""
28+
return self.transport is not None
29+
2630
def pause_writing(self) -> None:
2731
assert not self._paused
2832
self._paused = True
@@ -59,7 +63,6 @@ def connection_made(self, transport: asyncio.BaseTransport) -> None:
5963
self.transport = tr
6064

6165
def connection_lost(self, exc: Optional[BaseException]) -> None:
62-
self._connection_lost = True
6366
# Wake up the writer if currently paused.
6467
self.transport = None
6568
if not self._paused:
@@ -76,7 +79,7 @@ def connection_lost(self, exc: Optional[BaseException]) -> None:
7679
waiter.set_exception(exc)
7780

7881
async def _drain_helper(self) -> None:
79-
if self._connection_lost:
82+
if not self.connected:
8083
raise ConnectionResetError("Connection lost")
8184
if not self._paused:
8285
return

‎aiohttp/http_writer.py‎

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,6 @@ def __init__(
3535
on_headers_sent: _T_OnHeadersSent = None,
3636
) -> None:
3737
self._protocol = protocol
38-
self._transport = protocol.transport
3938

4039
self.loop = loop
4140
self.length = None
@@ -52,7 +51,7 @@ def __init__(
5251

5352
@property
5453
def transport(self) -> Optional[asyncio.Transport]:
55-
return self._transport
54+
return self._protocol.transport
5655

5756
@property
5857
def protocol(self) -> BaseProtocol:
@@ -71,10 +70,10 @@ def _write(self, chunk: bytes) -> None:
7170
size = len(chunk)
7271
self.buffer_size += size
7372
self.output_size += size
74-
75-
if self._transport is None or self._transport.is_closing():
73+
transport = self.transport
74+
if not self._protocol.connected or transport is None or transport.is_closing():
7675
raise ConnectionResetError("Cannot write to closing transport")
77-
self._transport.write(chunk)
76+
transport.write(chunk)
7877

7978
async def write(
8079
self, chunk: bytes, *, drain: bool = True, LIMIT: int = 0x10000
@@ -159,7 +158,6 @@ async def write_eof(self, chunk: bytes = b"") -> None:
159158
await self.drain()
160159

161160
self._eof = True
162-
self._transport = None
163161

164162
async def drain(self) -> None:
165163
"""Flush the write buffer.

‎tests/test_base_protocol.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -45,22 +45,22 @@ async def test_connection_lost_not_paused() -> None:
4545
pr = BaseProtocol(loop=loop)
4646
tr = mock.Mock()
4747
pr.connection_made(tr)
48-
assert not pr._connection_lost
48+
assert pr.connected
4949
pr.connection_lost(None)
5050
assert pr.transport is None
51-
assert pr._connection_lost
51+
assert not pr.connected
5252

5353

5454
async def test_connection_lost_paused_without_waiter() -> None:
5555
loop = asyncio.get_event_loop()
5656
pr = BaseProtocol(loop=loop)
5757
tr = mock.Mock()
5858
pr.connection_made(tr)
59-
assert not pr._connection_lost
59+
assert pr.connected
6060
pr.pause_writing()
6161
pr.connection_lost(None)
6262
assert pr.transport is None
63-
assert pr._connection_lost
63+
assert not pr.connected
6464

6565

6666
async def test_drain_lost() -> None:

‎tests/test_client_proto.py‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -134,3 +134,17 @@ async def test_eof_received(loop) -> None:
134134
assert proto._read_timeout_handle is not None
135135
proto.eof_received()
136136
assert proto._read_timeout_handle is None
137+
138+
139+
async def test_connection_lost_sets_transport_to_none(loop, mocker) -> None:
140+
"""Ensure that the transport is set to None when the connection is lost.
141+
142+
This ensures the writer knows that the connection is closed.
143+
"""
144+
proto = ResponseHandler(loop=loop)
145+
proto.connection_made(mocker.Mock())
146+
assert proto.transport is not None
147+
148+
proto.connection_lost(OSError())
149+
150+
assert proto.transport is None

‎tests/test_http_writer.py‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -236,6 +236,21 @@ async def test_write_to_closing_transport(protocol, transport, loop) -> None:
236236
await msg.write(b"After closing")
237237

238238

239+
async def test_write_to_closed_transport(protocol, transport, loop) -> None:
240+
"""Test that writing to a closed transport raises ConnectionResetError.
241+
242+
The StreamWriter checks to see if protocol.transport is None before
243+
writing to the transport. If it is None, it raises ConnectionResetError.
244+
"""
245+
msg = http.StreamWriter(protocol, loop)
246+
247+
await msg.write(b"Before transport close")
248+
protocol.transport = None
249+
250+
with pytest.raises(ConnectionResetError, match="Cannot write to closing transport"):
251+
await msg.write(b"After transport closed")
252+
253+
239254
async def test_drain(protocol, transport, loop) -> None:
240255
msg = http.StreamWriter(protocol, loop)
241256
await msg.drain()

0 commit comments

Comments
 (0)