fix: peek-before-consume in decode_handshake()

Previously the handshake path read 8 bytes BEFORE validating the
version, so an unsupported version left the buffer advanced by 8
with _handshake_done still False. A retry consumed the NEXT 8 bytes
as a "version" — almost always the header of the first real message
— and silently desynchronized the stream.

Add ReadBuffer.peek_bytes(n) for non-consuming reads, then rewrite
decode_handshake() to peek the version bytes, validate, and only
call read_bytes(8) on success. An invalid version now leaves the
bytes in place: a retry is deterministic (same bytes, same error)
instead of silently advancing into real message data.

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

Author: antoineleclair

src/dqlitewire/buffer.py

@@ -218,6 +218,18 @@ def read_bytes(self, n: int) -> bytes | None:
         self._maybe_compact()
         return data
 
+    def peek_bytes(self, n: int) -> bytes | None:
+        """Return the next n bytes without advancing the read position.
+
+        Returns None if fewer than n bytes are available. Symmetric with
+        ``read_bytes(n)`` but non-consuming — use this when you need to
+        validate the bytes before deciding whether to consume them.
+        """
+        available = len(self._data) - self._pos
+        if available < n:
+            return None
+        return bytes(self._data[self._pos : self._pos + n])
+
     def _maybe_compact(self) -> None:
         """Compact buffer if we've consumed a lot."""
         if self._pos > 4096:

src/dqlitewire/codec.py

@@ -333,15 +333,24 @@ def decode_handshake(self) -> int | None:
         Must be called before decode() on request decoders.
         Raises ProtocolError if the version is not recognized or if
         the handshake was already completed.
+
+        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.
         """
         if self._handshake_done:
             raise ProtocolError("Handshake already completed")
-        data = self._buffer.read_bytes(8)
-        if data is None:
+        # Peek first so we only commit on a valid version. An invalid version
+        # leaves the bytes in place — a retry is deterministic rather than
+        # silently advancing into real message data.
+        peek = self._buffer.peek_bytes(8)
+        if peek is None:
             return None
-        version = int.from_bytes(data, "little")
+        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)
         self._version = version
         self._handshake_done = True
         return version

tests/test_codec.py

@@ -221,6 +221,110 @@ def test_decode_handshake_rejects_unknown_version(self) -> None:
         with pytest.raises(ProtocolError, match="[Uu]nsupported protocol version"):
             decoder.decode_handshake()
 
+    def test_decode_handshake_failure_does_not_consume_bytes(self) -> None:
+        """On an unsupported version, decode_handshake() must NOT consume the
+        handshake bytes.
+
+        Previously the method read 8 bytes BEFORE validating the version, so
+        a failed handshake left the buffer advanced by 8 with ``_handshake_done``
+        still False. A retry consumed the next 8 bytes as a "version", which
+        was almost always the header of a real message — silently desynchronizing
+        the stream. Peek-before-consume means the bytes stay in the buffer and
+        a retry is deterministic: same bytes, same error.
+        """
+        from dqlitewire.exceptions import ProtocolError
+
+        decoder = MessageDecoder(is_request=True)
+        bogus = b"\x42" * 8
+        decoder.feed(bogus)
+
+        with pytest.raises(ProtocolError, match="[Uu]nsupported protocol version"):
+            decoder.decode_handshake()
+
+        # The 8 handshake bytes must still be in the buffer.
+        assert decoder._buffer.available() == 8
+        assert not decoder._handshake_done
+
+        # A retry on the same bytes gets the same error, deterministically.
+        with pytest.raises(ProtocolError, match="[Uu]nsupported protocol version"):
+            decoder.decode_handshake()
+        assert decoder._buffer.available() == 8
+
+    def test_decode_handshake_partial_data_leaves_bytes_intact(self) -> None:
+        """With fewer than 8 bytes buffered, decode_handshake() returns None
+        and leaves the partial data untouched so a subsequent feed() can
+        complete the handshake.
+        """
+        decoder = MessageDecoder(is_request=True)
+        version_bytes = PROTOCOL_VERSION_LEGACY.to_bytes(8, "little")
+        decoder.feed(version_bytes[:4])
+        assert decoder.decode_handshake() is None
+        assert decoder._buffer.available() == 4
+
+        decoder.feed(version_bytes[4:])
+        assert decoder.decode_handshake() == PROTOCOL_VERSION_LEGACY
+
+    def test_decode_handshake_failure_preserves_following_bytes(self) -> None:
+        """Direct reproducer of the original bug: a buffer containing a bogus
+        version followed by a real valid version. The pre-fix code consumed
+        both 8-byte chunks across two retries; the fix must consume neither
+        on the first failure.
+        """
+        from dqlitewire.exceptions import ProtocolError
+
+        decoder = MessageDecoder(is_request=True)
+        bogus = b"\x42" * 8
+        valid = PROTOCOL_VERSION_LEGACY.to_bytes(8, "little")
+        decoder.feed(bogus + valid)
+
+        with pytest.raises(ProtocolError, match="[Uu]nsupported protocol version"):
+            decoder.decode_handshake()
+        # All 16 bytes still in the buffer — neither chunk has been consumed.
+        assert decoder._buffer.available() == 16
+
+        # Even after a retry, the valid bytes behind are untouched.
+        with pytest.raises(ProtocolError, match="[Uu]nsupported protocol version"):
+            decoder.decode_handshake()
+        assert decoder._buffer.available() == 16
+
+    def test_decode_handshake_recoverable_via_reset(self) -> None:
+        """After a handshake failure, reset() clears the buffer and the
+        decoder accepts a fresh handshake on a reconnect.
+        """
+        from dqlitewire.exceptions import ProtocolError
+
+        decoder = MessageDecoder(is_request=True)
+        decoder.feed(b"\x42" * 8)
+        with pytest.raises(ProtocolError, match="[Uu]nsupported protocol version"):
+            decoder.decode_handshake()
+
+        decoder.reset()
+        assert decoder._buffer.available() == 0
+        assert not decoder._handshake_done
+
+        # Fresh valid handshake works.
+        decoder.feed(PROTOCOL_VERSION.to_bytes(8, "little"))
+        assert decoder.decode_handshake() == PROTOCOL_VERSION
+        assert decoder._handshake_done
+
+    def test_peek_bytes_does_not_advance_position(self) -> None:
+        """Sanity: ReadBuffer.peek_bytes() must return the requested bytes
+        without advancing _pos. This is the primitive decode_handshake()
+        depends on.
+        """
+        from dqlitewire.buffer import ReadBuffer
+
+        buf = ReadBuffer()
+        buf.feed(b"abcdefghij")
+        assert buf.peek_bytes(4) == b"abcd"
+        assert buf.available() == 10
+        # Repeated peeks return the same bytes.
+        assert buf.peek_bytes(4) == b"abcd"
+        assert buf.available() == 10
+        # Asking for more than available returns None.
+        assert buf.peek_bytes(20) is None
+        assert buf.available() == 10
+
     def test_decode_handshake_rejects_zero_version(self) -> None:
         """Version 0 is not a valid protocol version."""
         from dqlitewire.exceptions import ProtocolError