diff --git a/asyncpg/connect_utils.py b/asyncpg/connect_utils.py index e215391b..5cb90655 100644 --- a/asyncpg/connect_utils.py +++ b/asyncpg/connect_utils.py @@ -993,10 +993,26 @@ async def _create_ssl_connection( retry: bool = False, ) -> typing.Tuple[asyncio.Transport, _ProctolFactoryR]: - tr, pr = await loop.create_connection( - lambda: TLSUpgradeProto(loop, host, port, - ssl_context, ssl_is_advisory), - host, port) + upgrade_proto: typing.Optional[TLSUpgradeProto] = None + + def upgrade_protocol_factory() -> TLSUpgradeProto: + nonlocal upgrade_proto + upgrade_proto = TLSUpgradeProto( + loop, host, port, ssl_context, ssl_is_advisory) + return upgrade_proto + + try: + tr, pr = await loop.create_connection( + upgrade_protocol_factory, host, port) + except (Exception, asyncio.CancelledError): + # The protocol can exist before create_connection() returns. If + # that operation fails or is cancelled, nobody will await on_data. + if upgrade_proto is not None: + if not upgrade_proto.on_data.done(): + upgrade_proto.on_data.cancel() + elif not upgrade_proto.on_data.cancelled(): + upgrade_proto.on_data.exception() + raise tr.write(struct.pack('!ll', 8, _SSL_REQUEST_CODE)) # SSLRequest message. diff --git a/tests/test_connect.py b/tests/test_connect.py index 1c25d8d6..21a92afa 100644 --- a/tests/test_connect.py +++ b/tests/test_connect.py @@ -1761,8 +1761,12 @@ async def close(): class TestConnection(tb.ConnectedTestCase): async def test_connection_cancelled_during_transport_setup(self): - for connection_lost_first in (False, True): - with self.subTest(connection_lost_first=connection_lost_first): + for sslmode, connection_lost_first in itertools.product( + (False, 'prefer'), (False, True), + ): + with self.subTest( + sslmode=sslmode, connection_lost_first=connection_lost_first, + ): async def interrupted_connector(factory, *args, **kwargs): proto = factory() if connection_lost_first: @@ -1775,10 +1779,41 @@ async def interrupted_connector(factory, *args, **kwargs): self.loop, 'create_connection', interrupted_connector, ): with self.assertRaises(asyncio.CancelledError): - await self.connect(host='127.0.0.1', ssl=False) + await self.connect(host='127.0.0.1', ssl=sslmode) + + # Neither authentication nor the SSL upgrade must report an + # unobserved failure after cancellation, even if the + # transport closes afterwards. + await asyncio.sleep(0) + gc.collect() - # Authentication must not report an unobserved failure after - # cancellation, even if the transport closes afterwards. + async def test_connection_failed_during_transport_setup(self): + for sslmode, protocol_state in itertools.product( + (False, 'prefer'), ('not created', 'lost first', 'lost later'), + ): + with self.subTest(sslmode=sslmode, protocol_state=protocol_state): + error = ConnectionRefusedError('connection refused') + + async def failing_connector(factory, *args, **kwargs): + if protocol_state != 'not created': + proto = factory() + if protocol_state == 'lost first': + proto.connection_lost(None) + else: + self.loop.call_soon(proto.connection_lost, None) + raise error + + with unittest.mock.patch.object( + self.loop, 'create_connection', failing_connector, + ): + with self.assertRaises(ConnectionRefusedError) as cm: + await self.connect(host='127.0.0.1', ssl=sslmode) + self.assertIs(cm.exception, error) + + # The connection error must reach the caller unchanged, and + # neither authentication nor the SSL upgrade must report an + # unobserved failure after it, whether or not the protocol + # was created. await asyncio.sleep(0) gc.collect()