Add sslmode=allow support and fix =prefer retry

fantix · fantix · commit 77d47424ec29 · 2021-03-21T23:25:59.000-04:00
We didn't really retry the connection without SSL if the first SSL connection fails, that led to an issue when the server has SSL support but explicitly denies SSL connection through pg_hba.conf. This commit adds a retry in a new connection, which makes it easy to implement the sslmode=allow retry. Fixes #716
diff --git a/asyncpg/connect_utils.py b/asyncpg/connect_utils.py
@@ -35,7 +35,7 @@
 'password',
 'database',
 'ssl',
- 'ssl_is_advisory',
+ 'alt_retry_ssl_first',
 'connect_timeout',
 'server_settings',
 ])
@@ -402,8 +402,13 @@ def _parse_connect_dsn_and_args(*, dsn, host, port, user,
 if ssl is None and have_tcp_addrs:
 ssl = 'prefer'
 
- # ssl_is_advisory is only allowed to come from the sslmode parameter.
- ssl_is_advisory = None
+ # alt_retry_ssl_first is particularly for "allow" and "prefer"
+ # to alternatively try SSL/non-SSL connections (once each if supported):
+ # False - allow (try non-SSL first)
+ # True - prefer (try SSL first)
+ # None - other (don't retry, stick with the "ssl" parameter)
+ alt_retry_ssl_first = None
+
 if isinstance(ssl, str):
 SSLMODES ={
 'disable': 0,
@@ -420,26 +425,21 @@ def _parse_connect_dsn_and_args(*, dsn, host, port, user,
 raise exceptions.InterfaceError(
 '`sslmode` parameter must be one of:{}'.format(modes))
 
- # sslmode 'allow' is currently handled as 'prefer' because we're
- # missing the "retry with SSL" behavior for 'allow', but do have the
- # "retry without SSL" behavior for 'prefer'.
- # Not changing 'allow' to 'prefer' here would be effectively the same
- # as changing 'allow' to 'disable'.
 if sslmode == SSLMODES['allow']:
- sslmode = SSLMODES['prefer']
+ alt_retry_ssl_first = False
+ elif sslmode == SSLMODES['prefer']:
+ alt_retry_ssl_first = True
 
 # docs at https://www.postgresql.org/docs/10/static/libpq-connect.html
 # Not implemented: sslcert & sslkey & sslrootcert & sslcrl params.
- if sslmode <= SSLMODES['allow']:
+ if sslmode < SSLMODES['allow']:
 ssl = False
- ssl_is_advisory = sslmode >= SSLMODES['allow']
 else:
 ssl = ssl_module.create_default_context()
 ssl.check_hostname = sslmode >= SSLMODES['verify-full']
 ssl.verify_mode = ssl_module.CERT_REQUIRED
 if sslmode <= SSLMODES['require']:
 ssl.verify_mode = ssl_module.CERT_NONE
- ssl_is_advisory = sslmode <= SSLMODES['prefer']
 elif ssl is True:
 ssl = ssl_module.create_default_context()
 
@@ -453,7 +453,8 @@ def _parse_connect_dsn_and_args(*, dsn, host, port, user,
 
 params = _ConnectionParameters(
 user=user, password=password, database=database, ssl=ssl,
- ssl_is_advisory=ssl_is_advisory, connect_timeout=connect_timeout,
+ alt_retry_ssl_first=alt_retry_ssl_first,
+ connect_timeout=connect_timeout,
 server_settings=server_settings)
 
 return addrs, params
@@ -520,9 +521,8 @@ def data_received(self, data):
 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.
+ # sslmode=prefer. 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(
@@ -566,6 +566,7 @@ async def _create_ssl_connection(protocol_factory, host, port, *,
 new_tr = tr
 
 pg_proto = protocol_factory()
+ pg_proto.is_ssl = do_ssl_upgrade
 pg_proto.connection_made(new_tr)
 new_tr.set_protocol(pg_proto)
 
@@ -584,7 +585,9 @@ async def _create_ssl_connection(protocol_factory, host, port, *,
 tr.close()
 
 try:
- return await conn_factory(sock=sock)
+ new_tr, pg_proto = await conn_factory(sock=sock)
+ pg_proto.is_ssl = do_ssl_upgrade
+ return new_tr, pg_proto
 except (Exception, asyncio.CancelledError):
 sock.close()
 raise
@@ -605,8 +608,6 @@ async def _connect_addr(
 if timeout <= 0:
 raise asyncio.TimeoutError
 
- connected = _create_future(loop)
-
 params_input = params
 if callable(params.password):
 if inspect.iscoroutinefunction(params.password):
@@ -615,6 +616,44 @@ async def _connect_addr(
 password = params.password()
 
 params = params._replace(password=password)
+ args = (addr, loop, config, connection_class, record_class, params_input)
+
+ # skip retry if alt_retry is not enabled
+ if params.alt_retry_ssl_first is None:
+ return await __connect_addr(params, timeout, *args)
+
+ # prepare the params (which attempt has ssl) for the 2 attempts
+ params_retry = params._replace(ssl=None)
+ if not params.alt_retry_ssl_first:
+ params, params_retry = params_retry, params
+
+ # first attempt
+ before = time.monotonic()
+ try:
+ return await __connect_addr(params, timeout, *args)
+ except ConnectionError:
+ pass
+
+ # the second attempt with alt_retry_ssl_first=None
+ timeout -= time.monotonic() - before
+ if timeout <= 0:
+ raise asyncio.TimeoutError
+ else:
+ params_retry = params_retry._replace(alt_retry_ssl_first=None)
+ return await __connect_addr(params_retry, timeout, *args)
+
+
+async def __connect_addr(
+ params,
+ timeout,
+ addr,
+ loop,
+ config,
+ connection_class,
+ record_class,
+ params_input,
+):
+ connected = _create_future(loop)
 
 proto_factory = lambda: protocol.Protocol(
 addr, connected, params, record_class, loop)
@@ -625,7 +664,7 @@ async def _connect_addr(
 elif params.ssl:
 connector = _create_ssl_connection(
 proto_factory, *addr, loop=loop, ssl_context=params.ssl,
- ssl_is_advisory=params.ssl_is_advisory)
+ ssl_is_advisory=params.alt_retry_ssl_first)
 else:
 connector = loop.create_connection(proto_factory, *addr)
 
@@ -638,6 +677,23 @@ async def _connect_addr(
 if timeout <= 0:
 raise asyncio.TimeoutError
 await compat.wait_for(connected, timeout=timeout)
+ except exceptions.InvalidAuthorizationSpecificationError:
+ tr.close()
+
+ # pr.is_ssl is a bool, so this equal test implies
+ # alt_retry_ssl_first is not None (should do alt_retry)
+ if params.alt_retry_ssl_first == pr.is_ssl:
+ # Elevate the error to ConnectionError to trigger retry
+ raise ConnectionError("Connection rejected trying{} SSL".format(
+ 'with' if pr.is_ssl else 'without'))
+
+ else:
+ # Don't retry if alt_retry_ssl_first is None, or we don't need to
+ # (alt_retry_ssl_first=True and pr.is_ssl=False means the server
+ # doesn't support SSL, and we've already tried to Startup without
+ # SSL but failed; The opposite case doesn't exist).
+ raise
+
 except (Exception, asyncio.CancelledError):
 tr.close()
 raise
@@ -684,6 +740,7 @@ class CancelProto(asyncio.Protocol):
 
 def __init__(self):
 self.on_disconnect = _create_future(loop)
+ self.is_ssl = False
 
 def connection_lost(self, exc):
 if not self.on_disconnect.done():
@@ -692,17 +749,30 @@ def connection_lost(self, exc):
 if isinstance(addr, str):
 tr, pr = await loop.create_unix_connection(CancelProto, addr)
 else:
- if params.ssl:
- tr, pr = await _create_ssl_connection(
- CancelProto,
- *addr,
- loop=loop,
- ssl_context=params.ssl,
- ssl_is_advisory=params.ssl_is_advisory)
+ async def _connect(params_in, ssl_is_advisory):
+ if params_in.ssl:
+ return await _create_ssl_connection(
+ CancelProto,
+ *addr,
+ loop=loop,
+ ssl_context=params_in.ssl,
+ ssl_is_advisory=ssl_is_advisory)
+ else:
+ return await loop.create_connection(
+ CancelProto, *addr)
+ _set_nodelay(_get_socket(tr))
+
+ if params.alt_retry_ssl_first is None:
+ tr, pr = await _connect(params, False)
 else:
- tr, pr = await loop.create_connection(
- CancelProto, *addr)
- _set_nodelay(_get_socket(tr))
+ params_retry = params._replace(ssl=None)
+ if not params.alt_retry_ssl_first:
+ params, params_retry = params_retry, params
+ try:
+ tr, pr = await _connect(params, True)
+ except ConnectionError:
+ tr, pr = await _connect(
+ params._replace(alt_retry_ssl_first=None), False)
 
 # Pack a CancelRequest message
 msg = struct.pack('!llll', 16, 80877102, backend_pid, backend_secret)
diff --git a/asyncpg/connection.py b/asyncpg/connection.py
@@ -1879,7 +1879,8 @@ async def connect(dsn=None, *,
 - ``'disable'`` - SSL is disabled (equivalent to ``False``)
 - ``'prefer'`` - try SSL first, fallback to non-SSL connection
 if SSL connection fails
- - ``'allow'`` - currently equivalent to ``'prefer'``
+ - ``'allow'`` - try without SSL first, then retry with SSL if the first
+ attempt fails.
 - ``'require'`` - only try an SSL connection. Certificate
 verification errors are ignored
 - ``'verify-ca'`` - only try an SSL connection, and verify
diff --git a/asyncpg/protocol/protocol.pxd b/asyncpg/protocol/protocol.pxd
@@ -52,6 +52,8 @@ cdef class BaseProtocol(CoreProtocol):
 
 readonly uint64_t queries_count
 
+ bint _is_ssl
+
 PreparedStatementState statement
 
 cdef get_connection(self)
diff --git a/asyncpg/protocol/protocol.pyx b/asyncpg/protocol/protocol.pyx
@@ -103,6 +103,8 @@ cdef class BaseProtocol(CoreProtocol):
 
 self.queries_count = 0
 
+ self._is_ssl = False
+
 try:
 self.create_future = loop.create_future
 except AttributeError:
@@ -943,6 +945,14 @@ cdef class BaseProtocol(CoreProtocol):
 def resume_writing(self):
 self.writing_allowed.set()
 
+ @property
+ def is_ssl(self):
+ return self._is_ssl
+
+ @is_ssl.setter
+ def is_ssl(self, value):
+ self._is_ssl = value
+
 
 class Timer:
 def __init__(self, budget):
diff --git a/tests/test_connect.py b/tests/test_connect.py

-Original file line number
+Diff line change
 self.queries_count =0
 +self._is_ssl =False
++
 try:
 self.create_future = loop.create_future
 exceptAttributeError:
 defresume_writing(self):
 self.writing_allowed.set()
 +@property
 +defis_ssl(self):
 +returnself._is_ssl
++
 +@is_ssl.setter
 +defis_ssl(self, value):
 +self._is_ssl = value
++
 classTimer:
 def__init__(self, budget):