Fix: connection_acquisition_timeout now covers TLS handshake

Author: robsdedude

src/neo4j/_async/io/_bolt_socket.py

@@ -327,7 +327,11 @@ async def connect(
             s = None
             try:
                 s = await cls._connect_secure(
-                    resolved_address, tcp_timeout, keep_alive, ssl_context
+                    resolved_address,
+                    tcp_timeout,
+                    deadline,
+                    keep_alive,
+                    ssl_context,
                 )
                 agreed_version = await s._handshake(resolved_address, deadline)
                 return s, agreed_version

src/neo4j/_async_compat/network/_bolt_socket.py

@@ -205,13 +205,19 @@ def kill(self):
 
     @classmethod
     async def _connect_secure(
-        cls, resolved_address, timeout, keep_alive, ssl_context
+        cls,
+        resolved_address: Address,
+        timeout: float | None,
+        deadline: Deadline,
+        keep_alive: bool,
+        ssl_context: SSLContext | None,
     ) -> t.Self:
         """
         Connect to the address and return the socket.
 
         :param resolved_address:
         :param timeout: seconds
+        :param deadline: deadline for the whole operation
         :param keep_alive: True or False
         :param ssl_context: SSLContext or None
 
@@ -223,24 +229,35 @@ async def _connect_secure(
 
         # TODO: tomorrow me: fix this mess
         try:
-            if len(resolved_address) == 2:
-                s = socket(AF_INET)
-            elif len(resolved_address) == 4:
-                s = socket(AF_INET6)
-            else:
-                raise ValueError(f"Unsupported address {resolved_address!r}")
-            s.setblocking(False)  # asyncio + blocking = no-no!
-            log.debug("[#0000]  C: <OPEN> %s", resolved_address)
-            await wait_for(loop.sock_connect(s, resolved_address), timeout)
-            local_port = s.getsockname()[1]
+            try:
+                if len(resolved_address) == 2:
+                    s = socket(AF_INET)
+                elif len(resolved_address) == 4:
+                    s = socket(AF_INET6)
+                else:
+                    raise ValueError(
+                        f"Unsupported address {resolved_address!r}"
+                    )
+                s.setblocking(False)  # asyncio + blocking = no-no!
+                log.debug("[#0000]  C: <OPEN> %s", resolved_address)
+                await wait_for(loop.sock_connect(s, resolved_address), timeout)
+                local_port = s.getsockname()[1]
 
-            keep_alive = 1 if keep_alive else 0
-            s.setsockopt(SOL_SOCKET, SO_KEEPALIVE, keep_alive)
+                s.setsockopt(SOL_SOCKET, SO_KEEPALIVE, 1 if keep_alive else 0)
+            except asyncio.TimeoutError:
+                log.debug("[#0000]  S: <TIMEOUT> %s", resolved_address)
+                raise ServiceUnavailable(
+                    "Timed out trying to establish connection to "
+                    f"{resolved_address!r}"
+                ) from None
+            except asyncio.CancelledError:
+                log.debug("[#0000]  S: <CANCELLED> %s", resolved_address)
+                raise
 
             ssl_kwargs: dict[str, t.Any] = {}
 
+            hostname = resolved_address._host_name or None
             if ssl_context is not None:
-                hostname = resolved_address._host_name or None
                 sni_host = hostname if HAS_SNI and hostname else None
                 ssl_kwargs.update(ssl=ssl_context, server_hostname=sni_host)
                 log.debug("[#%04X]  C: <SECURE> %s", local_port, hostname)
@@ -250,9 +267,28 @@ async def _connect_secure(
                 loop=loop,
             )
             protocol = asyncio.StreamReaderProtocol(reader, loop=loop)
-            transport, _ = await loop.create_connection(
-                lambda: protocol, sock=s, **ssl_kwargs
-            )
+
+            try:
+                transport, _ = await wait_for(
+                    loop.create_connection(
+                        lambda: protocol, sock=s, **ssl_kwargs
+                    ),
+                    deadline.to_timeout(),
+                )
+            except (OSError, SSLError, CertificateError) as error:
+                log.debug(
+                    "[#0000]  S: <SECURE FAILURE> %s: %r",
+                    resolved_address,
+                    error,
+                )
+                raise BoltSecurityError(
+                    message="Failed to establish encrypted connection.",
+                    address=(hostname, local_port),
+                ) from error
+            except asyncio.CancelledError:
+                log.debug("[#0000]  S: <CANCELLED> %s", resolved_address)
+                raise
+
             writer = asyncio.StreamWriter(transport, protocol, reader, loop)
 
             if ssl_context is not None:
@@ -265,39 +301,8 @@ async def _connect_secure(
                     raise BoltProtocolError(
                         "When using an encrypted socket, the server should "
                         "always provide a certificate",
-                        address=(resolved_address._host_name, local_port),
+                        address=(hostname, local_port),
                     )
-
-            return cls(reader, protocol, writer)
-
-        except asyncio.TimeoutError:
-            log.debug("[#0000]  S: <TIMEOUT> %s", resolved_address)
-            log.debug("[#0000]  C: <CLOSE> %s", resolved_address)
-            if s:
-                cls._kill_raw_socket(s)
-            raise ServiceUnavailable(
-                "Timed out trying to establish connection to "
-                f"{resolved_address!r}"
-            ) from None
-        except asyncio.CancelledError:
-            log.debug("[#0000]  S: <CANCELLED> %s", resolved_address)
-            log.debug("[#0000]  C: <CLOSE> %s", resolved_address)
-            if s:
-                cls._kill_raw_socket(s)
-            raise
-        except (SSLError, CertificateError) as error:
-            log.debug(
-                "[#0000]  S: <SECURE FAILURE> %s: %s",
-                resolved_address,
-                error,
-            )
-            log.debug("[#0000]  C: <CLOSE> %s", resolved_address)
-            if s:
-                cls._kill_raw_socket(s)
-            raise BoltSecurityError(
-                message="Failed to establish encrypted connection.",
-                address=(resolved_address._host_name, local_port),
-            ) from error
         except Exception as error:
             log.debug(
                 "[#0000]  S: <ERROR> %s %s",
@@ -314,6 +319,8 @@ async def _connect_secure(
                 ) from error
             raise
 
+        return cls(reader, protocol, writer)
+
     @abc.abstractmethod
     async def _handshake(
         self,
@@ -463,13 +470,19 @@ def kill(self):
 
     @classmethod
     def _connect_secure(
-        cls, resolved_address, timeout, keep_alive, ssl_context
+        cls,
+        resolved_address: Address,
+        timeout: float | None,
+        deadline: Deadline,
+        keep_alive: bool,
+        ssl_context: SSLContext | None,
     ):
         """
         Connect to the address and return the socket.
 
         :param resolved_address:
         :param timeout: seconds
+        :param deadline: deadline for the whole operation
         :param keep_alive: True or False
         :returns: socket object
         """
@@ -497,26 +510,14 @@ def _connect_secure(
                 log.debug("[#0000]  C: <OPEN> %s", resolved_address)
                 s.connect(resolved_address)
                 s.settimeout(t)
-                keep_alive = 1 if keep_alive else 0
-                s.setsockopt(SOL_SOCKET, SO_KEEPALIVE, keep_alive)
+                s.setsockopt(SOL_SOCKET, SO_KEEPALIVE, 1 if keep_alive else 0)
             except TimeoutError:
                 log.debug("[#0000]  S: <TIMEOUT> %s", resolved_address)
-                log.debug("[#0000]  C: <CLOSE> %s", resolved_address)
-                if s:
-                    cls._kill_raw_socket(s)
                 raise ServiceUnavailable(
                     "Timed out trying to establish connection to "
                     f"{resolved_address!r}"
                 ) from None
             except Exception as error:
-                log.debug(
-                    "[#0000]  S: <ERROR> %s %s",
-                    type(error).__name__,
-                    " ".join(map(repr, error.args)),
-                )
-                log.debug("[#0000]  C: <CLOSE> %s", resolved_address)
-                if s:
-                    cls._kill_raw_socket(s)
                 if isinstance(error, OSError):
                     raise ServiceUnavailable(
                         "Failed to establish connection to "
@@ -531,16 +532,17 @@ def _connect_secure(
                 sni_host = hostname if HAS_SNI and hostname else None
                 log.debug("[#%04X]  C: <SECURE> %s", local_port, hostname)
                 try:
+                    t = s.gettimeout()
+                    if timeout:
+                        s.settimeout(deadline.to_timeout())
                     s = ssl_context.wrap_socket(s, server_hostname=sni_host)
+                    s.settimeout(t)
                 except (OSError, SSLError, CertificateError) as cause:
                     log.debug(
-                        "[#0000]  S: <SECURE FAILURE> %s: %s",
+                        "[#0000]  S: <SECURE FAILURE> %s: %r",
                         resolved_address,
                         cause,
                     )
-                    log.debug("[#0000]  C: <CLOSE> %s", resolved_address)
-                    if s:
-                        cls._kill_raw_socket(s)
                     raise BoltSecurityError(
                         message="Failed to establish encrypted connection.",
                         address=(hostname, local_port),
@@ -554,15 +556,17 @@ def _connect_secure(
                         "[#0000]  S: <SECURE FAILURE> %s: no certificate",
                         resolved_address,
                     )
-                    log.debug("[#0000]  C: <CLOSE> %s", resolved_address)
-                    if s:
-                        cls._kill_raw_socket(s)
                     raise BoltProtocolError(
                         "When using an encrypted socket, the server should"
                         "always provide a certificate",
                         address=(hostname, local_port),
                     )
-        except Exception:
+        except Exception as error:
+            log.debug(
+                "[#0000]  S: <ERROR> %s %s",
+                type(error).__name__,
+                " ".join(map(repr, error.args)),
+            )
             if s is not None:
                 log.debug("[#0000]  C: <CLOSE> %s", resolved_address)
                 cls._kill_raw_socket(s)