diff --git a/cassandra/connection.py b/cassandra/connection.py index d9bca62718..04834b413d 100644 --- a/cassandra/connection.py +++ b/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: diff --git a/tests/unit/test_connection.py b/tests/unit/test_connection.py index 3bca654c55..4c2b02508a 100644 --- a/tests/unit/test_connection.py +++ b/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, [])}