Fix: connection_acquisition_timeout now covers TLS handshake

Author: robsdedude

src/neo4j/_async/io/_bolt_socket.py

@@ -326,7 +326,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, handshake, response = await s._handshake(
                     resolved_address, deadline

src/neo4j/_async_compat/network/_bolt_socket.py

@@ -170,13 +170,14 @@ def kill(self):
 
     @classmethod
     async def _connect_secure(
-        cls, resolved_address, timeout, keep_alive, ssl_context
+        cls, resolved_address, timeout, deadline, keep_alive, ssl_context
     ) -> te.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
 
@@ -215,8 +216,9 @@ async def _connect_secure(
                 loop=loop,
             )
             protocol = asyncio.StreamReaderProtocol(reader, loop=loop)
-            transport, _ = await loop.create_connection(
-                lambda: protocol, sock=s, **ssl_kwargs
+            transport, _ = await wait_for(
+                loop.create_connection(lambda: protocol, sock=s, **ssl_kwargs),
+                deadline.to_timeout(),
             )
             writer = asyncio.StreamWriter(transport, protocol, reader, loop)
 
@@ -374,13 +376,14 @@ def kill(self):
 
     @classmethod
     def _connect_secure(
-        cls, resolved_address, timeout, keep_alive, ssl_context
+        cls, resolved_address, timeout, deadline, keep_alive, ssl_context
     ):
         """
         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
         """
@@ -436,7 +439,11 @@ 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:
                     raise BoltSecurityError(
                         message="Failed to establish encrypted connection.",