Keep a partial protocol v5 segment header across reads

cassandra/connection.py

@@ -1198,7 +1198,9 @@ def _process_segment_buffer(self):
                 # re-raise an exception that inherits from ConnectionException
                 raise CrcMismatchException(str(exc), self.endpoint)
         else:
+            # not even a full header yet. rewind too, or reset_io_buffer() drops these bytes
             self._io_buffer._segment_consumed = False
+            self._io_buffer.io_buffer.seek(0)
 
     def process_io_buffer(self):
         while True:

tests/unit/test_connection.py

@@ -23,10 +23,10 @@
 from cassandra.cluster import Cluster
 from cassandra.connection import (Connection, HEADER_DIRECTION_TO_CLIENT, ProtocolError,
                                   locally_supported_compressions, ConnectionHeartbeat, _Frame, Timer, TimerManager,
-                                  ConnectionException, DefaultEndPoint)
+                                  ConnectionException, DefaultEndPoint, segment_codec_no_compression)
 from cassandra.marshal import uint8_pack, uint32_pack, int32_pack
 from cassandra.protocol import (write_stringmultimap, write_int, write_string,
-                                SupportedMessage, ProtocolHandler)
+                                SupportedMessage, ProtocolHandler, ReadyMessage)
 
 from tests.util import wait_until
 
@@ -123,6 +123,34 @@ def test_negative_body_length(self, *args):
         args, kwargs = c.defunct.call_args
         self.assertIsInstance(args[0], ProtocolError)
 
+    def test_v5_segment_header_split_across_reads(self):
+        # a read that ends inside the next segment's 6-byte header must not lose those bytes
+        def frame(stream):
+            body = b'x' * 100
+            return self.make_header_prefix(ReadyMessage, version=5, stream_id=stream) + int32_pack(len(body)) + body
+
+        wire = BytesIO()
+        segment_codec_no_compression.encode(wire, frame(1))
+        segment_codec_no_compression.encode(wire, frame(2))
+        data = wire.getvalue()
+        first_segment_len = len(data) // 2
+
+        for extra in range(8):
+            with self.subTest(header_bytes_in_first_read=extra):
+                c = self.make_connection()
+                c.defunct = Mock()
+                c._enable_checksumming()
+                streams = []
+                c.process_msg = lambda header, body: streams.append(header.stream)
+
+                split_at = first_segment_len + extra
+                for chunk in (data[:split_at], data[split_at:]):
+                    c._iobuf.write(chunk)
+                    c.process_io_buffer()
+
+                c.defunct.assert_not_called()
+                assert streams == [1, 2]
+
     def test_unsupported_cql_version(self, *args):
         c = self.make_connection()
         c._requests = {0: (c._handle_options_response, ProtocolHandler.decode_message, [])}