Per-frame cap and opt-in heartbeat trust

max_continuation_frames: the per-query row cap is already enforced,
but a hostile server sending 1-row frames can still pin a client CPU
with ~max_total_rows iterations of Python-level decode work. Add a
complementary ``max_continuation_frames`` knob (default 100_000)
enforced inside ``_drain_continuations``. An earlier design pass
originally called for both knobs; only ``max_total_rows`` was
delivered then. Plumbed through :class:`DqliteProtocol` and
:class:`DqliteConnection`; callers get the guard with no code
change.

trust_server_heartbeat: previously the handshake unconditionally
widened the client's per-read deadline to the server-advertised
heartbeat (capped at 300 s) — a 30× amplification from the common
10 s default. A hostile server could therefore override the
operator's configured timeout. Invert the default:
``trust_server_heartbeat=False`` means the server value is recorded
for diagnostics only and ``timeout`` is authoritative. Opt-in
preserves the previous behavior for callers who need adaptive
behavior.

Validation: extract a shared ``_validate_positive_int_or_none``
helper used for both ``max_total_rows`` and
``max_continuation_frames``.

Tests: extend test_max_total_rows_cap.py with
TestMaxContinuationFramesEnforcement (exceeds-cap, within-cap,
None-disables) and TestTrustServerHeartbeat (default no
amplification, opt-in honored).

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

Author: antoineleclair

src/dqliteclient/connection.py

@@ -14,7 +14,11 @@
     OperationalError,
     ProtocolError,
 )
-from dqliteclient.protocol import DqliteProtocol, _validate_max_total_rows
+from dqliteclient.protocol import (
+    DqliteProtocol,
+    _validate_max_total_rows,
+    _validate_positive_int_or_none,
+)
 from dqlitewire.exceptions import EncodeError as _WireEncodeError
 
 # dqlite error codes that indicate a leader change (SQLite extended error codes)
@@ -76,6 +80,8 @@ def __init__(
         database: str = "default",
         timeout: float = 10.0,
         max_total_rows: int | None = 10_000_000,
+        max_continuation_frames: int | None = 100_000,
+        trust_server_heartbeat: bool = False,
     ) -> None:
         """Initialize connection (does not connect yet).
 
@@ -87,13 +93,27 @@ def __init__(
                 frames for a single query. Prevents a slow-drip server
                 from keeping the client alive indefinitely within the
                 per-operation deadline. Set to ``None`` to disable.
+            max_continuation_frames: Maximum number of continuation
+                frames in a single query result. Caps the per-query
+                Python-side decode work a hostile server can inflict
+                by sending many 1-row frames (ISSUE-98). Set to
+                ``None`` to disable.
+            trust_server_heartbeat: When True, widen the per-read
+                deadline to the server-advertised heartbeat (subject
+                to a 300 s hard cap). When False (default), ``timeout``
+                is authoritative — the server value cannot amplify it
+                (ISSUE-101).
         """
         if not math.isfinite(timeout) or timeout <= 0:
             raise ValueError(f"timeout must be a positive finite number, got {timeout}")
         self._address = address
         self._database = database
         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
         self._protocol: DqliteProtocol | None = None
         self._db_id: int | None = None
         self._in_transaction = False
@@ -148,6 +168,8 @@ async def connect(self) -> None:
                 writer,
                 timeout=self._timeout,
                 max_total_rows=self._max_total_rows,
+                max_continuation_frames=self._max_continuation_frames,
+                trust_server_heartbeat=self._trust_server_heartbeat,
             )
 
             try:

src/dqliteclient/protocol.py

@@ -32,6 +32,23 @@
 _READ_CHUNK_SIZE = 4096
 
 
+def _validate_positive_int_or_none(
+    value: int | None, name: str
+) -> int | None:
+    """Shared validation for positive-int-or-None parameters.
+
+    Used for both ``max_total_rows`` and ``max_continuation_frames``
+    (ISSUE-98). None disables the cap; any int value must be > 0.
+    """
+    if value is None:
+        return None
+    if not isinstance(value, int) or isinstance(value, bool):
+        raise TypeError(f"{name} must be int or None, got {type(value).__name__}")
+    if value <= 0:
+        raise ValueError(f"{name} must be > 0 or None, got {value}")
+    return value
+
+
 def _validate_max_total_rows(value: int | None) -> int | None:
     """Validate the ``max_total_rows`` constructor argument.
 
@@ -56,6 +73,8 @@ def __init__(
         writer: asyncio.StreamWriter,
         timeout: float = 15.0,
         max_total_rows: int | None = 10_000_000,
+        max_continuation_frames: int | None = 100_000,
+        trust_server_heartbeat: bool = False,
     ) -> None:
         self._reader = reader
         self._writer = writer
@@ -69,6 +88,21 @@ def __init__(
         # could legitimately allocate hundreds of millions of rows over
         # the full deadline. None disables the cap.
         self._max_total_rows = _validate_max_total_rows(max_total_rows)
+        # Per-query frame cap. Complements max_total_rows: a server
+        # sending 10M 1-row frames to reach the row cap would still
+        # burn 10M × decode-cost of Python work; the frame cap bounds
+        # that at ~100k iterations (ISSUE-98).
+        self._max_continuation_frames = _validate_positive_int_or_none(
+            max_continuation_frames, "max_continuation_frames"
+        )
+        # When True, the client honors the server-advertised heartbeat
+        # timeout to adjust its per-read deadline (subject to the 300 s
+        # hard cap). When False (default), the server value is recorded
+        # for diagnostics only and the operator-configured ``timeout``
+        # is authoritative. Opt-in protects operators whose timeout is
+        # a latency-SLO boundary from server-induced amplification
+        # (ISSUE-101).
+        self._trust_server_heartbeat = trust_server_heartbeat
 
     async def handshake(self, client_id: int | None = None) -> int:
         """Perform protocol handshake.
@@ -96,8 +130,14 @@ async def handshake(self, client_id: int | None = None) -> int:
 
         self._client_id = client_id
         self._heartbeat_timeout = response.heartbeat_timeout
-        # Use heartbeat timeout for subsequent reads if larger than default
-        if response.heartbeat_timeout > 0:
+        # Use the server-advertised heartbeat only when explicitly
+        # trusted. Previously we always widened ``self._timeout`` up
+        # to 300 s based on the server value, which let a hostile
+        # server amplify the operator's configured timeout up to 30×
+        # (ISSUE-101). Default is now opt-out: the server value is
+        # recorded for diagnostics but does not change the per-read
+        # deadline.
+        if self._trust_server_heartbeat and response.heartbeat_timeout > 0:
             heartbeat_seconds = response.heartbeat_timeout / 1000.0
             # Cap to prevent a malicious/buggy server from disabling timeouts
             self._timeout = max(self._timeout, min(heartbeat_seconds, 300.0))
@@ -230,12 +270,27 @@ async def _drain_continuations(
         """
         all_rows = list(initial.rows)
         response = initial
+        frames = 1  # the initial frame counts
         while response.has_more:
             next_response = await self._read_continuation(deadline=deadline)
+            frames += 1
             if not next_response.rows and next_response.has_more:
                 raise ProtocolError(
                     "ROWS continuation made no progress: frame had 0 rows and has_more=True"
                 )
+            if (
+                self._max_continuation_frames is not None
+                and frames > self._max_continuation_frames
+            ):
+                # Per-frame cap complements max_total_rows (ISSUE-98): a
+                # slow-drip server sending 1-row-per-frame would
+                # otherwise pin a client CPU with O(n) iterations of
+                # decode work, where n is max_total_rows.
+                raise ProtocolError(
+                    f"Query exceeded max_continuation_frames cap "
+                    f"({self._max_continuation_frames}); server may be "
+                    f"slow-dripping rows."
+                )
             if self._max_total_rows is not None and (
                 len(all_rows) + len(next_response.rows) > self._max_total_rows
             ):

tests/test_max_total_rows_cap.py

@@ -69,3 +69,92 @@ def test_none_disables_cap(self) -> None:
 
         rows = asyncio.run(p._drain_continuations(initial, deadline=999999.0))
         assert len(rows) == 10_000
+
+
+class TestMaxContinuationFramesEnforcement:
+    """ISSUE-98: per-frame cap complements max_total_rows.
+
+    A slow-drip server sending many 1-row frames could pin a client CPU
+    with ~max_total_rows iterations of Python-level decode work. The
+    frame cap bounds that at a configurable number of iterations.
+    """
+
+    def test_exceeding_frame_cap_raises(self) -> None:
+        reader = MagicMock()
+        writer = MagicMock()
+        # Cap at 3 frames (initial + 2 continuations).
+        p = DqliteProtocol(
+            reader,
+            writer,
+            timeout=5.0,
+            max_total_rows=None,
+            max_continuation_frames=3,
+        )
+
+        initial = _make_rows_response([[1]], has_more=True)
+        # Every continuation delivers one row with has_more=True. The
+        # test mock returns the same response object each time so the
+        # frame counter is what drives termination.
+        one_row_continuation = _make_rows_response([[2]], has_more=True)
+
+        p._read_continuation = AsyncMock(return_value=one_row_continuation)  # type: ignore[method-assign]
+
+        async def run() -> None:
+            await p._drain_continuations(initial, deadline=999999.0)
+
+        with pytest.raises(ProtocolError, match="max_continuation_frames"):
+            asyncio.run(run())
+
+    def test_within_frame_cap_succeeds(self) -> None:
+        reader = MagicMock()
+        writer = MagicMock()
+        p = DqliteProtocol(
+            reader,
+            writer,
+            timeout=5.0,
+            max_continuation_frames=5,
+        )
+        initial = _make_rows_response([[1]], has_more=True)
+        last = _make_rows_response([[2], [3]], has_more=False)
+
+        p._read_continuation = AsyncMock(return_value=last)  # type: ignore[method-assign]
+        rows = asyncio.run(p._drain_continuations(initial, deadline=999999.0))
+        assert len(rows) == 3
+
+    def test_none_disables_frame_cap(self) -> None:
+        reader = MagicMock()
+        writer = MagicMock()
+        p = DqliteProtocol(
+            reader,
+            writer,
+            timeout=5.0,
+            max_continuation_frames=None,
+        )
+        initial = _make_rows_response([[1]], has_more=True)
+        last = _make_rows_response([[2]], has_more=False)
+
+        p._read_continuation = AsyncMock(return_value=last)  # type: ignore[method-assign]
+        rows = asyncio.run(p._drain_continuations(initial, deadline=999999.0))
+        assert len(rows) == 2
+
+
+class TestTrustServerHeartbeat:
+    """ISSUE-101: server heartbeat no longer widens client timeout by default."""
+
+    def test_default_does_not_amplify_timeout(self) -> None:
+        reader = MagicMock()
+        writer = MagicMock()
+        p = DqliteProtocol(reader, writer, timeout=5.0)
+        # Inject the server-advertised value; the field is set in
+        # handshake() but we bypass the socket dance for a pure
+        # attribute test.
+        p._heartbeat_timeout = 300_000  # 300 s in ms, far above 5 s
+        # trust_server_heartbeat defaults to False, so timeout stayed 5.
+        assert p._timeout == 5.0
+        assert p._trust_server_heartbeat is False
+
+    def test_opt_in_respected(self) -> None:
+        reader = MagicMock()
+        writer = MagicMock()
+        p = DqliteProtocol(reader, writer, timeout=5.0, trust_server_heartbeat=True)
+        assert p._trust_server_heartbeat is True