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..b6545ad7 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,11 @@ 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) - # Authentication must not report an unobserved failure after - # cancellation, even if the transport closes afterwards. + # 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()