Use loop.start_tls() to upgrade connections to SSL

The old way of TLS upgrade (openining a connection, asking postgres to do TLS and then duping the underlying socket) seems not to work anymore on Windows with Python 3.8.

Author: 1st1

asyncpg/connect_utils.py

@@ -504,6 +504,95 @@ def _parse_connect_arguments(*, dsn, host, port, user, password, passfile,
 returnaddrs, params, config
 
 
+classTLSUpgradeProto(asyncio.Protocol):
+def__init__(self, loop, host, port, ssl_context, ssl_is_advisory):
+self.on_data=_create_future(loop)
+self.host=host
+self.port=port
+self.ssl_context=ssl_context
+self.ssl_is_advisory=ssl_is_advisory
+
+defdata_received(self, data):
+ifdata==b'S':
+self.on_data.set_result(True)
+elif (self.ssl_is_advisoryand
+self.ssl_context.verify_mode==ssl_module.CERT_NONEand
+data==b'N'):
+# ssl_is_advisory will imply that ssl.verify_mode == CERT_NONE,
+# since the only way to get ssl_is_advisory is from
+# sslmode=prefer (or sslmode=allow). But be extra sure to
+# disallow insecure connections when the ssl context asks for
+# real security.
+self.on_data.set_result(False)
+else:
+self.on_data.set_exception(
+ConnectionError(
+'PostgreSQL server at "{host}:{port}" '
+'rejected SSL upgrade'.format(
+host=self.host, port=self.port)))
+
+defconnection_lost(self, exc):
+ifnotself.on_data.done():
+ifexcisNone:
+exc=ConnectionError('unexpected connection_lost() call')
+self.on_data.set_exception(exc)
+
+
+asyncdef_create_ssl_connection(protocol_factory, host, port, *,
+loop, ssl_context, ssl_is_advisory=False):
+
+ifssl_contextisTrue:
+ssl_context=ssl_module.create_default_context()
+
+tr, pr=awaitloop.create_connection(
+lambda: TLSUpgradeProto(loop, host, port,
+ssl_context, ssl_is_advisory),
+host, port)
+
+tr.write(struct.pack('!ll', 8, 80877103)) # SSLRequest message.
+
+try:
+do_ssl_upgrade=awaitpr.on_data
+except (Exception, asyncio.CancelledError):
+tr.close()
+raise
+
+ifhasattr(loop, 'start_tls'):
+ifdo_ssl_upgrade:
+try:
+new_tr=awaitloop.start_tls(
+tr, pr, ssl_context, server_hostname=host)
+except (Exception, asyncio.CancelledError):
+tr.close()
+raise
+else:
+new_tr=tr
+
+pg_proto=protocol_factory()
+pg_proto.connection_made(new_tr)
+new_tr.set_protocol(pg_proto)
+
+returnnew_tr, pg_proto
+else:
+conn_factory=functools.partial(
+loop.create_connection, protocol_factory)
+
+ifdo_ssl_upgrade:
+conn_factory=functools.partial(
+conn_factory, ssl=ssl_context, server_hostname=host)
+
+sock=_get_socket(tr)
+sock=sock.dup()
+_set_nodelay(sock)
+tr.close()
+
+try:
+returnawaitconn_factory(sock=sock)
+except (Exception, asyncio.CancelledError):
+sock.close()
+raise
+
+
 asyncdef_connect_addr(*, addr, loop, timeout, params, config,
 connection_class):
 assertloopisnotNone
@@ -526,8 +615,6 @@ async def _connect_addr(*, addr, loop, timeout, params, config,
 else:
 connector=loop.create_connection(proto_factory, *addr)
 
-connector=asyncio.ensure_future(connector)
-
 before=time.monotonic()
 try:
 tr, pr=awaitasyncio.wait_for(
@@ -575,79 +662,41 @@ async def _connect(*, loop, timeout, connection_class, **kwargs):
 raiselast_error
 
 
-asyncdef_negotiate_ssl_connection(host, port, conn_factory, *, loop, ssl,
-server_hostname, ssl_is_advisory=False):
-# Note: ssl_is_advisory only affects behavior when the server does not
-# accept SSLRequests. If the SSLRequest is accepted but either the SSL
-# negotiation fails or the PostgreSQL user isn't permitted to use SSL,
-# there's nothing that would attempt to reconnect with a non-SSL socket.
-reader, writer=awaitasyncio.open_connection(host, port)
-
-tr=writer.transport
-try:
-sock=_get_socket(tr)
-_set_nodelay(sock)
-
-writer.write(struct.pack('!ll', 8, 80877103)) # SSLRequest message.
-awaitwriter.drain()
-resp=awaitreader.readexactly(1)
-
-ifresp==b'S':
-conn_factory=functools.partial(
-conn_factory, ssl=ssl, server_hostname=server_hostname)
-elif (ssl_is_advisoryand
-ssl.verify_mode==ssl_module.CERT_NONEand
-resp==b'N'):
-# ssl_is_advisory will imply that ssl.verify_mode == CERT_NONE,
-# since the only way to get ssl_is_advisory is from sslmode=prefer
-# (or sslmode=allow). But be extra sure to disallow insecure
-# connections when the ssl context asks for real security.
-pass
-else:
-raiseConnectionError(
-'PostgreSQL server at "{}:{}" rejected SSL upgrade'.format(
-host, port))
-
-sock=sock.dup() # Must come before tr.close()
-finally:
-writer.close()
-awaitcompat.wait_closed(writer)
-
-try:
-returnawaitconn_factory(sock=sock) # Must come after tr.close()
-except (Exception, asyncio.CancelledError):
-sock.close()
-raise
+asyncdef_cancel(*, loop, addr, params: _ConnectionParameters,
+backend_pid, backend_secret):
 
+classCancelProto(asyncio.Protocol):
 
-asyncdef_create_ssl_connection(protocol_factory, host, port, *,
-loop, ssl_context, ssl_is_advisory=False):
-returnawait_negotiate_ssl_connection(
-host, port,
-functools.partial(loop.create_connection, protocol_factory),
-loop=loop,
-ssl=ssl_context,
-server_hostname=host,
-ssl_is_advisory=ssl_is_advisory)
+def__init__(self):
+self.on_disconnect=_create_future(loop)
 
+defconnection_lost(self, exc):
+ifnotself.on_disconnect.done():
+self.on_disconnect.set_result(True)
 
-asyncdef_open_connection(*, loop, addr, params: _ConnectionParameters):
 ifisinstance(addr, str):
-r, w=awaitasyncio.open_unix_connection(addr)
+tr, pr=awaitloop.create_unix_connection(CancelProto, addr)
 else:
 ifparams.ssl:
-r, w=await_negotiate_ssl_connection(
+tr, pr=await_create_ssl_connection(
+CancelProto,
 *addr,
-asyncio.open_connection,
 loop=loop,
-ssl=params.ssl,
-server_hostname=addr[0],
+ssl_context=params.ssl,
 ssl_is_advisory=params.ssl_is_advisory)
 else:
-r, w=awaitasyncio.open_connection(*addr)
-_set_nodelay(_get_socket(w.transport))
+tr, pr=awaitloop.create_connection(
+CancelProto, *addr)
+_set_nodelay(_get_socket(tr))
+
+# Pack a CancelRequest message
+msg=struct.pack('!llll', 16, 80877102, backend_pid, backend_secret)
 
-returnr, w
+try:
+tr.write(msg)
+awaitpr.on_disconnect
+finally:
+tr.close()
 
 
 def_get_socket(transport):

asyncpg/connection.py

@@ -10,7 +10,6 @@
 importcollections
 importcollections.abc
 importitertools
-importstruct
 importsys
 importtime
 importtraceback
@@ -1186,24 +1185,16 @@ async def _cleanup_stmts(self):
 awaitself._protocol.close_statement(stmt, protocol.NO_TIMEOUT)
 
 asyncdef_cancel(self, waiter):
-r=w=None
-
 try:
 # Open new connection to the server
-r, w=awaitconnect_utils._open_connection(
-loop=self._loop, addr=self._addr, params=self._params)
-
-# Pack CancelRequest message
-msg=struct.pack('!llll', 16, 80877102,
-self._protocol.backend_pid,
-self._protocol.backend_secret)
-
-w.write(msg)
-awaitr.read() # Wait until EOF
+awaitconnect_utils._cancel(
+loop=self._loop, addr=self._addr, params=self._params,
+backend_pid=self._protocol.backend_pid,
+backend_secret=self._protocol.backend_secret)
 exceptConnectionResetErrorasex:
 # On some systems Postgres will reset the connection
 # after processing the cancellation command.
-ifrisNoneandnotwaiter.done():
+ifnotwaiter.done():
 waiter.set_exception(ex)
 exceptasyncio.CancelledError:
 # There are two scenarios in which the cancellation
@@ -1221,9 +1212,6 @@ async def _cancel(self, waiter):
 compat.current_asyncio_task(self._loop))
 ifnotwaiter.done():
 waiter.set_result(None)
-ifwisnotNone:
-w.close()
-awaitcompat.wait_closed(w)
 
 def_cancel_current_command(self, waiter):
 self._cancellations.add(self._loop.create_task(self._cancel(waiter)))