Add sslmode=allow support and fix =prefer retry

We didn't really retry the connection without SSL if the first SSL connection fails under sslmode=prefer, 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

Author: fantix

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,31 @@ 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:
+ rv = await loop.create_connection(
+ CancelProto, *addr)
+ _set_nodelay(_get_socket(rv[0]))
+ return rv
+
+ 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)

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

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)

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):