Add sslmode=allow support and fix =prefer retry

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 MagicStack#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,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)

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