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