Plumb max_continuation_frames + trust_server_heartbeat through pool and cluster

Post-review follow-up. The security review flagged that these knobs
existed only on DqliteProtocol/DqliteConnection — operators using
ConnectionPool or ClusterClient.connect (which is the common path)
had no way to tune them. Add both parameters to:

- ConnectionPool.__init__
- ClusterClient.connect

and forward them through _create_connection / try_connect to the
underlying DqliteConnection.

Also strengthens the trust_server_heartbeat tests:
- opt-in with handshake actually amplifies timeout to the server value
- default ignores handshake heartbeat entirely (no amplification)
- opt-in respects the 300 s hard cap

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: antoineleclair

src/dqliteclient/cluster.py

@@ -160,14 +160,17 @@ async def connect(
         database: str = "default",
         *,
         max_total_rows: int | None = 10_000_000,
+        max_continuation_frames: int | None = 100_000,
+        trust_server_heartbeat: bool = False,
         max_attempts: int | None = None,
     ) -> DqliteConnection:
         """Connect to the cluster leader.
 
-        Returns a connection to the current leader. ``max_total_rows``
-        is forwarded to the underlying :class:`DqliteConnection` so
-        callers (including :class:`ConnectionPool`) can tune the
-        cumulative row cap from one place.
+        Returns a connection to the current leader. ``max_total_rows``,
+        ``max_continuation_frames``, and ``trust_server_heartbeat`` are
+        forwarded to the underlying :class:`DqliteConnection` so callers
+        (including :class:`ConnectionPool`) can tune security/DoS
+        governors from one place.
 
         ``max_attempts`` overrides the default
         :data:`_DEFAULT_CONNECT_MAX_ATTEMPTS` (ISSUE-109).
@@ -194,6 +197,8 @@ async def try_connect() -> DqliteConnection:
                     database=database,
                     timeout=self._timeout,
                     max_total_rows=max_total_rows,
+                    max_continuation_frames=max_continuation_frames,
+                    trust_server_heartbeat=trust_server_heartbeat,
                 )
                 await conn.connect()
                 return conn

src/dqliteclient/pool.py

@@ -11,7 +11,7 @@
 from dqliteclient.connection import DqliteConnection
 from dqliteclient.exceptions import DqliteConnectionError
 from dqliteclient.node_store import NodeStore
-from dqliteclient.protocol import _validate_max_total_rows
+from dqliteclient.protocol import _validate_max_total_rows, _validate_positive_int_or_none
 
 logger = logging.getLogger(__name__)
 
@@ -62,6 +62,8 @@ def __init__(
         cluster: ClusterClient | None = None,
         node_store: NodeStore | None = None,
         max_total_rows: int | None = 10_000_000,
+        max_continuation_frames: int | None = 100_000,
+        trust_server_heartbeat: bool = False,
     ) -> None:
         """Initialize connection pool.
 
@@ -99,6 +101,16 @@ def __init__(
                 connection inherits the same governor. ``None`` disables
                 the cap entirely (not recommended in production —
                 bounds memory against slow-drip attacks).
+            max_continuation_frames: Per-query continuation-frame cap
+                (ISSUE-98). Complements ``max_total_rows``: a server
+                sending 1-row-per-frame can inflict O(n) Python decode
+                work where n is the row cap; the frame cap bounds that.
+                Forwarded to every :class:`DqliteConnection`.
+            trust_server_heartbeat: When True, the per-read deadline on
+                every connection widens to the server-advertised
+                heartbeat (up to a 300 s hard cap). Default False —
+                operator-configured ``timeout`` is authoritative and
+                the server cannot amplify it (ISSUE-101).
         """
         if min_size < 0:
             raise ValueError(f"min_size must be non-negative, got {min_size}")
@@ -119,6 +131,10 @@ def __init__(
         self._max_size = max_size
         self._timeout = timeout
         self._max_total_rows = _validate_max_total_rows(max_total_rows)
+        self._max_continuation_frames = _validate_positive_int_or_none(
+            max_continuation_frames, "max_continuation_frames"
+        )
+        self._trust_server_heartbeat = trust_server_heartbeat
 
         if cluster is not None:
             self._cluster = cluster
@@ -172,7 +188,10 @@ async def _create_connection(self) -> DqliteConnection:
         reservation.
         """
         return await self._cluster.connect(
-            database=self._database, max_total_rows=self._max_total_rows
+            database=self._database,
+            max_total_rows=self._max_total_rows,
+            max_continuation_frames=self._max_continuation_frames,
+            trust_server_heartbeat=self._trust_server_heartbeat,
         )
 
     async def _release_reservation(self) -> None:

tests/test_max_total_rows_cap.py

@@ -158,3 +158,60 @@ def test_opt_in_respected(self) -> None:
         writer = MagicMock()
         p = DqliteProtocol(reader, writer, timeout=5.0, trust_server_heartbeat=True)
         assert p._trust_server_heartbeat is True
+
+    def test_opt_in_amplifies_timeout_via_handshake(self) -> None:
+        """With trust_server_heartbeat=True, a handshake reply with a
+        larger-than-local heartbeat widens the per-read deadline (up to
+        the 300 s hard cap). Exercises the actual handshake code path
+        rather than just checking the flag is set (ISSUE-101 review)."""
+        from dqlitewire.messages import WelcomeResponse
+
+        reader = AsyncMock()
+        writer = MagicMock()
+        writer.drain = AsyncMock()
+        writer.close = MagicMock()
+        # heartbeat=60_000 ms = 60 s, well above timeout=5 s, under 300 s cap
+        reader.read.side_effect = [
+            WelcomeResponse(heartbeat_timeout=60_000).encode(),
+        ]
+        p = DqliteProtocol(reader, writer, timeout=5.0, trust_server_heartbeat=True)
+        asyncio.run(p.handshake())
+        assert p._timeout == 60.0, (
+            f"trust_server_heartbeat=True should widen timeout to 60s, got {p._timeout}"
+        )
+
+    def test_default_ignores_handshake_heartbeat(self) -> None:
+        """With trust_server_heartbeat=False (default), a handshake
+        reply with a larger heartbeat is recorded but does NOT widen
+        the per-read deadline."""
+        from dqlitewire.messages import WelcomeResponse
+
+        reader = AsyncMock()
+        writer = MagicMock()
+        writer.drain = AsyncMock()
+        writer.close = MagicMock()
+        reader.read.side_effect = [
+            WelcomeResponse(heartbeat_timeout=300_000).encode(),
+        ]
+        p = DqliteProtocol(reader, writer, timeout=5.0)  # default = False
+        asyncio.run(p.handshake())
+        # Server value recorded for diagnostics, but timeout unchanged.
+        assert p._heartbeat_timeout == 300_000
+        assert p._timeout == 5.0, f"default should not widen timeout, got {p._timeout}"
+
+    def test_opt_in_respects_hard_300s_cap(self) -> None:
+        """Even with trust_server_heartbeat=True, a server sending an
+        absurdly large heartbeat is clamped at 300 s."""
+        from dqlitewire.messages import WelcomeResponse
+
+        reader = AsyncMock()
+        writer = MagicMock()
+        writer.drain = AsyncMock()
+        writer.close = MagicMock()
+        # heartbeat=3_600_000 ms = 1 h; clamp to 300 s.
+        reader.read.side_effect = [
+            WelcomeResponse(heartbeat_timeout=3_600_000).encode(),
+        ]
+        p = DqliteProtocol(reader, writer, timeout=5.0, trust_server_heartbeat=True)
+        asyncio.run(p.handshake())
+        assert p._timeout == 300.0