Fix socket timeouts during connection establishment (#1302)

  * Prevent async socket from passing 0 as `ssl_handshake_timeout` leading to
    a `ValueError` being raised when the acquisition deadline has expired.
    Raise `ServiceUnavailable` instead.
  * Align async driver with sync driver:
    * (lazily) raise a `ValueError` when configured with a negative
      `connection_timeout`.
    * Treat `connection_timeout=0` as "no timeout".
  * Fix sync driver setting the socket timeout to `0` (non-blocking) on expired
    `connection_acquisition_timeout`.
  * Harden internal code against against `None` timeouts.

Author: robsdedude

requirements-dev.txt

@@ -12,6 +12,7 @@ isort>=5.11.5  # TODO: 6.0 - bump when support for Python 3.7 is dropped
 mypy>=1.4.1  # TODO: 6.0 - bump when support for Python 3.7 is dropped
 typing-extensions>=4.7.1
 types-pytz>=2023.3.1.1  # TODO: 6.0 - bump when support for Python 3.7 is dropped
+types-mock>=5.1.0.3  # TODO: 6.0 - bump when support for Python 3.7 is dropped
 ruff>=0.8.2
 
 # needed for running tests

src/neo4j/_async/io/_bolt_socket.py

@@ -318,7 +318,7 @@ async def connect(
         )
         async for resolved_address in resolved_addresses:
             deadline_timeout = deadline.to_timeout()
-            if (
+            if tcp_timeout is None or (
                 deadline_timeout is not None
                 and deadline_timeout <= tcp_timeout
             ):

src/neo4j/_async/io/_pool.py

@@ -301,8 +301,7 @@ async def _re_auth_connection(self, connection, auth, force, unprepared):
             )
         except Exception as exc:
             log.debug(
-                "[#%04X]  _: <POOL> check re_auth failed %r auth=%s "
-                "force=%s",
+                "[#%04X]  _: <POOL> check re_auth failed %r auth=%s force=%s",
                 connection.local_port,
                 exc,
                 log_auth,
@@ -879,8 +878,7 @@ async def fetch_routing_table(
         # No readers
         if num_readers == 0:
             log.debug(
-                "[#0000]  _: <POOL> no read servers returned from "
-                "server %s",
+                "[#0000]  _: <POOL> no read servers returned from server %s",
                 address,
             )
             return None

src/neo4j/_async_compat/network/_bolt_socket.py

@@ -63,6 +63,8 @@
     from ..._async.io import AsyncBolt
     from ..._sync.io import Bolt
 
+    _P = te.ParamSpec("_P")
+
 
 log = logging.getLogger("neo4j.io")
 
@@ -76,6 +78,25 @@ def _sanitize_deadline(deadline):
     return deadline
 
 
+def _validate_timeout(timeout):
+    if timeout is not None and timeout < 0:
+        raise ValueError("Timeout value out of range")
+
+
+def _non_expired_timeout(
+    deadline: Deadline | None,
+    operation: str,
+) -> float | None:
+    if deadline is None:
+        return None
+    timeout = deadline.to_timeout()
+    if timeout is None:
+        return None
+    if timeout <= 0:
+        raise SocketDeadlineExceededError(f"{operation} timed out")
+    return timeout
+
+
 class AsyncBoltSocketBase(abc.ABC):
     Bolt: te.Final[type[AsyncBolt]] = None  # type: ignore[assignment]
 
@@ -86,32 +107,40 @@ def __init__(self, reader, protocol, writer) -> None:
         # 0 - non-blocking
         # None infinitely blocking
         # int - seconds to wait for data
-        self._timeout = None
-        self._deadline = None
-
-    async def _wait_for_io(self, io_async_fn, *args, **kwargs):
+        self._timeout: float | None = None
+        self._deadline: Deadline | None = None
+
+    async def _wait_for_io(
+        self,
+        io_async_fn: t.Callable[_P, t.Coroutine],
+        *args: _P.args,
+        **kwargs: _P.kwargs,
+    ) -> None:
         timeout = self._timeout
-        to_raise = SocketTimeout
-        if self._deadline is not None:
-            deadline_timeout = self._deadline.to_timeout()
-            if deadline_timeout <= 0:
-                raise SocketDeadlineExceededError("timed out")
-            if timeout is None or deadline_timeout <= timeout:
-                timeout = deadline_timeout
-                to_raise = SocketDeadlineExceededError
-
-        io_fut = io_async_fn(*args, **kwargs)
+        to_raise: type[Exception] = SocketTimeout
+        deadline_timeout = _non_expired_timeout(self._deadline, "IO operation")
+        if deadline_timeout is not None and (
+            timeout is None or deadline_timeout <= timeout
+        ):
+            timeout = deadline_timeout
+            to_raise = SocketDeadlineExceededError
+
+        io_fut: t.Awaitable
         if timeout is not None and timeout <= 0:
             # give the io-operation time for one loop cycle to do its thing
-            io_fut = asyncio.create_task(io_fut)
+            io_fut = asyncio.create_task(io_async_fn(*args, **kwargs))
             try:
                 await asyncio.sleep(0)
             except asyncio.CancelledError:
                 # This is emulating non-blocking. There is no cancelling this!
                 # Still, we don't want to silently swallow the cancellation.
                 # Hence, we flag this task as cancelled again, so that the next
                 # `await` will raise the CancelledError.
-                asyncio.current_task().cancel()
+                current_task = asyncio.current_task()
+                if current_task is not None:
+                    current_task.cancel()
+        else:
+            io_fut = io_async_fn(*args, **kwargs)
         try:
             return await wait_for(io_fut, timeout)
         except asyncio.TimeoutError as e:
@@ -197,6 +226,9 @@ async def _connect_secure(
                 raise ValueError(f"Unsupported address {resolved_address!r}")
             s.setblocking(False)  # asyncio + blocking = no-no!
             log.debug("[#0000]  C: <OPEN> %s", resolved_address)
+            _validate_timeout(timeout)
+            if timeout == 0:  # socket timeout of 0 => non-blocking
+                timeout = None
             await wait_for(loop.sock_connect(s, resolved_address), timeout)
             local_port = s.getsockname()[1]
 
@@ -208,18 +240,26 @@ async def _connect_secure(
             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,
-                    ssl_handshake_timeout=deadline.to_timeout(),
-                )
+                ssl_kwargs.update(ssl=ssl_context, server_hostname=sni_host)
                 log.debug("[#%04X]  C: <SECURE> %s", local_port, hostname)
 
             reader = asyncio.StreamReader(
                 limit=2**16,  # 64 KiB,
                 loop=loop,
             )
             protocol = asyncio.StreamReaderProtocol(reader, loop=loop)
+            if ssl_context is not None:
+                try:
+                    ssl_timeout = _non_expired_timeout(
+                        deadline, "SSL handshake"
+                    )
+                except SocketDeadlineExceededError as error:
+                    raise BoltSecurityError(
+                        message="Failed to establish encrypted connection.",
+                        address=(hostname, local_port),
+                    ) from error
+                if ssl_timeout is not None:
+                    ssl_kwargs["ssl_handshake_timeout"] = ssl_timeout
             transport, _ = await loop.create_connection(
                 lambda: protocol, sock=s, **ssl_kwargs
             )
@@ -262,6 +302,16 @@ async def _connect_secure(
                 message="Failed to establish encrypted connection.",
                 address=(resolved_address._host_name, local_port),
             ) from error
+        except BoltSecurityError as error:
+            log.debug(
+                "[#0000]  S: <SECURE FAILURE> %s: %r",
+                resolved_address,
+                error,
+            )
+            if s:
+                log.debug("[#0000]  C: <CLOSE> %s", resolved_address)
+                cls._kill_raw_socket(s)
+            raise
         except Exception as error:
             log.debug(
                 "[#0000]  S: <ERROR> %s %s",
@@ -315,14 +365,14 @@ class BoltSocketBase:
 
     def __init__(self, socket_: socket):
         self._socket = socket_
-        self._deadline = None
+        self._deadline: Deadline | None = None
 
     @property
-    def _socket(self):
+    def _socket(self) -> socket | SSLSocket:
         return self.__socket
 
     @_socket.setter
-    def _socket(self, socket_: socket | SSLSocket):
+    def _socket(self, socket_: socket | SSLSocket) -> None:
         self.__socket = socket_
         self.getsockname = socket_.getsockname
         self.getpeername = socket_.getpeername
@@ -339,14 +389,17 @@ def _socket(self, socket_: socket | SSLSocket):
     gettimeout: t.Callable = None  # type: ignore
     settimeout: t.Callable = None  # type: ignore
 
-    def _wait_for_io(self, func, *args, **kwargs):
-        if self._deadline is None:
-            return func(*args, **kwargs)
-        timeout = self._socket.gettimeout()
-        deadline_timeout = self._deadline.to_timeout()
-        if deadline_timeout <= 0:
-            raise SocketDeadlineExceededError("timed out")
-        if timeout is None or deadline_timeout <= timeout:
+    def _wait_for_io(
+        self,
+        func: t.Callable[_P, t.Any],
+        *args: _P.args,
+        **kwargs: _P.kwargs,
+    ) -> None:
+        timeout: float | None = self._socket.gettimeout()
+        deadline_timeout = _non_expired_timeout(self._deadline, "IO operation")
+        if deadline_timeout is not None and (
+            timeout is None or deadline_timeout <= timeout
+        ):
             self._socket.settimeout(deadline_timeout)
             try:
                 return func(*args, **kwargs)
@@ -443,11 +496,19 @@ def _connect_secure(
                 log.debug("[#%04X]  C: <SECURE> %s", local_port, hostname)
                 try:
                     t = s.gettimeout()
-                    if timeout:
-                        s.settimeout(deadline.to_timeout())
+                    ssl_timeout = _non_expired_timeout(
+                        deadline, "SSL handshake"
+                    )
+                    if ssl_timeout is not None:
+                        s.settimeout(ssl_timeout)
                     s = ssl_context.wrap_socket(s, server_hostname=sni_host)
                     s.settimeout(t)
-                except (OSError, SSLError, CertificateError) as cause:
+                except (
+                    OSError,
+                    SSLError,
+                    CertificateError,
+                    SocketDeadlineExceededError,
+                ) as cause:
                     raise BoltSecurityError(
                         message="Failed to establish encrypted connection.",
                         address=(hostname, local_port),

src/neo4j/_deadline.py

@@ -14,76 +14,97 @@
 # limitations under the License.
 
 
+from __future__ import annotations
+
+import typing as t
 from contextlib import contextmanager
 from time import monotonic
 
 
+if t.TYPE_CHECKING:
+    import typing_extensions as te
+
+    from ._async.io import AsyncBolt
+    from ._sync.io import Bolt
+
+
 class Deadline:
-    def __init__(self, timeout):
+    def __init__(self, timeout: float | None) -> None:
         if timeout is None or timeout == float("inf"):
             self._deadline = float("inf")
         else:
             self._deadline = monotonic() + timeout
         self._original_timeout = timeout
 
     @property
-    def original_timeout(self):
+    def original_timeout(self) -> float | None:
         return self._original_timeout
 
-    def expired(self):
+    def expired(self) -> bool:
         return self.to_timeout() == 0
 
-    def to_timeout(self):
+    def to_timeout(self) -> float | None:
         if self._deadline == float("inf"):
             return None
         timeout = self._deadline - monotonic()
         return max(0, timeout)
 
-    def __eq__(self, other):
+    def __eq__(self, other) -> bool:
         if isinstance(other, Deadline):
             return self._deadline == other._deadline
         return NotImplemented
 
-    def __gt__(self, other):
+    def __gt__(self, other) -> bool:
         if isinstance(other, Deadline):
             return self._deadline > other._deadline
         return NotImplemented
 
-    def __ge__(self, other):
+    def __ge__(self, other) -> bool:
         if isinstance(other, Deadline):
             return self._deadline >= other._deadline
         return NotImplemented
 
-    def __lt__(self, other):
+    def __lt__(self, other) -> bool:
         if isinstance(other, Deadline):
             return self._deadline < other._deadline
         return NotImplemented
 
-    def __le__(self, other):
+    def __le__(self, other) -> bool:
         if isinstance(other, Deadline):
             return self._deadline <= other._deadline
         return NotImplemented
 
     @classmethod
-    def from_timeout_or_deadline(cls, timeout):
-        if isinstance(timeout, cls):
-            return timeout
-        return cls(timeout)
-
-    def __str__(self):
+    def from_timeout_or_deadline(
+        cls, timeout: te.Self | float | None
+    ) -> te.Self:
+        if isinstance(timeout, (float, int)) or timeout is None:
+            return cls(timeout)
+        return timeout
+
+    def __str__(self) -> str:
         return f"Deadline(timeout={self._original_timeout})"
 
+    def __bool__(self) -> bool:
+        """Whether a deadline is set (:data:`True`) or not (:data:`False`)."""
+        return self._deadline != float("inf")
+
 
 merge_deadlines = min
 
 
-def merge_deadlines_and_timeouts(*deadline):
+def merge_deadlines_and_timeouts(
+    *deadline: Deadline | None,
+) -> Deadline:
     deadlines = map(Deadline.from_timeout_or_deadline, deadline)
     return merge_deadlines(deadlines)
 
 
 @contextmanager
-def connection_deadline(connection, deadline):
+def connection_deadline(
+    connection: AsyncBolt | Bolt,
+    deadline: Deadline | None,
+) -> t.Generator[None, None, None]:
     original_deadline = connection.socket.get_deadline()
     if deadline is None and original_deadline is not None:
         # nothing to do here