diff --git a/cassandra/connection.py b/cassandra/connection.py index 155b95d510..a0e8fa08b2 100644 --- a/cassandra/connection.py +++ b/cassandra/connection.py @@ -1786,7 +1786,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 5f49c24554..91607b5727 100644 --- a/tests/unit/test_connection.py +++ b/tests/unit/test_connection.py @@ -27,7 +27,7 @@ from cassandra.connection import (Connection, HEADER_DIRECTION_TO_CLIENT, ProtocolError, locally_supported_compressions, ConnectionHeartbeat, HeartbeatFuture, _Frame, Timer, TimerManager, ConnectionException, ConnectionSendError, ConnectionShutdown, ConnectionBusy, DefaultEndPoint, ShardAwarePortGenerator, - DRIVER_NAME, DRIVER_VERSION) + DRIVER_NAME, DRIVER_VERSION, segment_codec_no_compression) from cassandra.driver_config import DRIVER_CONFIG_OPTION, SESSION_ID_OPTION from cassandra.ssl_session_cache import SSLSessionCache from cassandra.marshal import uint8_pack, uint32_pack, int32_pack @@ -165,6 +165,34 @@ def test_negative_body_length(self, *args): args, kwargs = c.defunct.call_args assert isinstance(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, [])}