fix: reorder decode_handshake commit to survive async exceptions

decode_handshake used to commit via three sequential statements:

    self._buffer.read_bytes(8)     # (1) consume bytes
    self._version = version         # (2) record version
    self._handshake_done = True     # (3) mark done

A KeyboardInterrupt (or any other PyErr_SetAsyncExc delivery) landing
between (1) and (3) left the buffer with the 8 handshake bytes
consumed but _handshake_done still False. On retry, decode_handshake
re-peeked 8 bytes — but those were now the first 8 bytes of the
NEXT message, which almost always fail the _SUPPORTED_VERSIONS
check. The caller then saw a misleading "Unsupported protocol
version" error pointing at what was actually valid message data.

Reorder: commit _version/_handshake_done FIRST, then consume. If
read_bytes(8) raises for any BaseException subclass, revert the
state commit so the buffer is still coherent (8 bytes still
available) and a retry repeats the peek/validate/commit cycle
cleanly.

Does NOT fix hazard 1 from issue 041 (two-thread TOCTOU) — that
remains issue 021's single-owner contract territory.

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

Author: antoineleclair

src/dqlitewire/codec.py

@@ -358,6 +358,19 @@ def decode_handshake(self) -> int | None:
         On an unsupported version, the 8 handshake bytes are left in the
         buffer untouched so that a retry is deterministic (same bytes, same
         error) rather than silently consuming the next 8 bytes of real data.
+
+        Signal-safety (issue 041): the commit order is
+        ``_version``/``_handshake_done`` FIRST, then ``read_bytes(8)``.
+        If an async exception (``KeyboardInterrupt``) lands between the
+        state commit and the buffer consume, the except block reverts
+        ``_handshake_done`` and ``_version`` so the buffer is still
+        coherent: the 8 handshake bytes are still there and a retry
+        repeats the peek/validate/commit cycle. This replaces the
+        previous "consume then commit" order, which allowed a signal
+        to leave the bytes consumed but the state not yet marked —
+        retry would then re-peek 8 bytes of real message data as a
+        handshake and almost always raise a misleading "Unsupported
+        protocol version" error.
         """
         if self._handshake_done:
             raise ProtocolError("Handshake already completed")
@@ -370,10 +383,17 @@ def decode_handshake(self) -> int | None:
         version = int.from_bytes(peek, "little")
         if version not in _SUPPORTED_VERSIONS:
             raise ProtocolError(f"Unsupported protocol version: {version:#x}")
-        # Valid — commit by advancing past the handshake bytes.
-        self._buffer.read_bytes(8)
+        # Commit state BEFORE consuming bytes. If the consume is
+        # interrupted by an async exception, revert so the peek/commit
+        # pair becomes atomic from the caller's perspective.
         self._version = version
         self._handshake_done = True
+        try:
+            self._buffer.read_bytes(8)
+        except BaseException:
+            self._handshake_done = False
+            self._version = None
+            raise
         return version
 
 

tests/test_codec_signal_safety.py

@@ -1,4 +1,4 @@
-"""Signal-safety tests for MessageDecoder (issue 045).
+"""Signal-safety tests for MessageDecoder (issues 041 and 045).
 
 ``MessageDecoder.decode()`` and ``decode_continuation()`` both consume
 bytes from the buffer via ``read_message()`` and then parse them via
@@ -15,9 +15,18 @@
 caller catches it at the top level and retries; the retry reads the
 *next* message boundary, silently losing exactly one message.
 
+``MessageDecoder.decode_handshake()`` has an analogous hazard: the
+commit used to be (1) consume bytes, (2) set ``_version``, (3) set
+``_handshake_done``. A signal between (1) and (3) left the bytes
+consumed but state not marked; the retry re-peeked the first 8 bytes
+of the *next* message as a handshake and raised a misleading
+"Unsupported protocol version" error.
+
 These tests use ``sys.settrace`` to inject a ``KeyboardInterrupt`` at
-a specific source line inside ``decode_bytes`` and assert that the
-decoder is left poisoned.
+a specific source line inside a decoder method — the most reliable
+way to exercise CPython's bytecode-boundary async-exception model
+without relying on wallclock timing — and assert the decoder ends up
+in a state the caller can detect and recover from, not silently torn.
 """
 
 from __future__ import annotations
@@ -31,7 +40,7 @@
 import pytest
 
 from dqlitewire.codec import MessageDecoder
-from dqlitewire.constants import ResponseType
+from dqlitewire.constants import PROTOCOL_VERSION, ResponseType
 from dqlitewire.exceptions import DecodeError, ProtocolError
 from dqlitewire.messages import DbResponse, LeaderResponse
 
@@ -65,10 +74,12 @@ def tracer(frame: FrameType, event: str, arg: object) -> Any:
 
 
 class TestDecodeSignalSafety:
-    def test_keyboard_interrupt_inside_parse_leaves_decoder_poisoned(self) -> None:
-        """Regression for issue 045.
+    """Regression tests for issue 045 (consumed-but-unpoisoned window
+    in ``decode()`` / ``decode_continuation()``).
+    """
 
-        A ``KeyboardInterrupt`` delivered while ``decode_bytes`` is
+    def test_keyboard_interrupt_inside_parse_leaves_decoder_poisoned(self) -> None:
+        """A ``KeyboardInterrupt`` delivered while ``decode_bytes`` is
         parsing a consumed message used to propagate without
         poisoning the decoder, because ``except Exception`` does not
         catch ``BaseException``. The retry read the next message
@@ -177,3 +188,116 @@ def test_db_response_decodes_cleanly_without_interrupt() -> None:
     msg = dec.decode()
     assert isinstance(msg, DbResponse)
     assert msg.db_id == 42
+
+
+class TestDecodeHandshakeSignalSafety:
+    """Regression tests for issue 041 (hazard 2: single-threaded signal
+    split inside ``decode_handshake``).
+
+    The pre-fix ``decode_handshake`` committed via three separate
+    statements:
+
+        self._buffer.read_bytes(8)       # consume bytes
+        self._version = version           # record version
+        self._handshake_done = True       # mark done
+
+    A ``KeyboardInterrupt`` delivered between ``read_bytes(8)`` and
+    the final store left the buffer with the 8 handshake bytes
+    consumed but ``_handshake_done`` still ``False``. On retry,
+    ``decode_handshake`` re-peeked ``peek_bytes(8)`` — but those were
+    now the first 8 bytes of the **next** message, which almost
+    always fail the ``_SUPPORTED_VERSIONS`` check. The caller saw a
+    misleading "Unsupported protocol version" error pointing at what
+    was actually legitimate message data.
+
+    The fix marks ``_handshake_done`` BEFORE consuming the bytes and
+    reverts on failure. An interrupt between the mark and the consume
+    leaves the buffer unconsumed and the state "already completed";
+    the retry raises a deterministic error instead of silently
+    misreading real data.
+    """
+
+    def test_keyboard_interrupt_post_consume_is_not_silently_misleading(
+        self,
+    ) -> None:
+        dec = MessageDecoder(is_request=True)
+        # 8 valid handshake bytes followed by 8 bytes that WOULD fail
+        # the version check if the retry misread them as a handshake.
+        dec.feed(PROTOCOL_VERSION.to_bytes(8, "little") + b"\xff" * 8)
+
+        state = {"raised": False}
+
+        def tracer(frame: FrameType, event: str, arg: object) -> Any:
+            if event != "line":
+                return tracer
+            if frame.f_code.co_name != "decode_handshake":
+                return tracer
+            if state["raised"]:
+                return tracer
+            try:
+                with open(frame.f_code.co_filename) as f:
+                    src_line = f.readlines()[frame.f_lineno - 1]
+            except OSError:
+                return tracer
+            # Inject at the `self._version = version` line — it is
+            # present in both the pre-fix and post-fix layouts, and
+            # in the pre-fix source it runs AFTER read_bytes(8).
+            if "self._version = version" in src_line:
+                state["raised"] = True
+                raise KeyboardInterrupt("injected mid-commit")
+            return tracer
+
+        sys.settrace(tracer)
+        try:
+            with contextlib.suppress(KeyboardInterrupt):
+                dec.decode_handshake()
+        finally:
+            sys.settrace(None)
+
+        # After the torn window, a retry must NOT report
+        # "Unsupported protocol version" for what were actually real
+        # bytes of the next message (0xff*8). Acceptable outcomes:
+        #   1. The fix reverted state — retry succeeds with the
+        #      original 8 handshake bytes still in the buffer.
+        #   2. The fix set _handshake_done before the consume and
+        #      did not revert — retry raises "already completed".
+        retry_err: Exception | None = None
+        retry_result: int | None = None
+        try:
+            retry_result = dec.decode_handshake()
+        except ProtocolError as e:
+            retry_err = e
+
+        if retry_err is not None:
+            assert "Unsupported" not in str(retry_err), (
+                f"retry reported misleading error: {retry_err}"
+            )
+        else:
+            assert retry_result == PROTOCOL_VERSION
+            assert dec._handshake_done is True
+
+    def test_happy_path_handshake_still_works(self) -> None:
+        """Sanity: without any interrupt, decode_handshake completes
+        normally and sets all three state bits.
+        """
+        dec = MessageDecoder(is_request=True)
+        dec.feed(PROTOCOL_VERSION.to_bytes(8, "little"))
+        version = dec.decode_handshake()
+        assert version == PROTOCOL_VERSION
+        assert dec._handshake_done is True
+        assert dec._version == PROTOCOL_VERSION
+
+    def test_unsupported_version_does_not_consume_bytes(self) -> None:
+        """Counter-test for issues 027 and 041: an unsupported version
+        must leave the 8 peeked bytes in the buffer so that retry
+        semantics remain deterministic.
+        """
+        dec = MessageDecoder(is_request=True)
+        bogus = (0xDEADBEEF).to_bytes(8, "little")
+        dec.feed(bogus)
+
+        with pytest.raises(ProtocolError, match="Unsupported"):
+            dec.decode_handshake()
+
+        assert dec._handshake_done is False
+        assert dec._buffer.available() == 8