diff --git a/mssql-tds/src/connection/tds_client.rs b/mssql-tds/src/connection/tds_client.rs index 67d4d16f..77ca1a1d 100644 --- a/mssql-tds/src/connection/tds_client.rs +++ b/mssql-tds/src/connection/tds_client.rs @@ -184,6 +184,18 @@ pub enum CursorColumn { RowEnded, } +/// Result of a non-blocking pull-cursor attempt. +#[derive(Debug, PartialEq)] +pub enum CursorPoll { + /// The operation completed entirely from bytes already buffered by the + /// transport. + Ready(T), + /// The existing async cursor method must continue the operation. + /// + /// No transport bytes or cursor state were consumed by the attempt. + Pending, +} + /// Active TDS connection to a SQL Server instance. /// /// Created by [`TdsConnectionProvider::create_client()`](crate::connection_provider::tds_connection_provider::TdsConnectionProvider::create_client). @@ -3690,6 +3702,53 @@ impl TdsClient { } } + /// Attempts to position the cursor from bytes already buffered by the transport. + /// + /// Returns [`CursorPoll::Pending`] without consuming bytes or cursor state + /// when the current row must first be drained, the next token needs async + /// parsing, encryption keys need resolving, or the row header is incomplete. + pub fn try_next_row_cursor(&mut self) -> TdsResult> { + let Some(metadata) = self.current_metadata.as_ref().map(Arc::clone) else { + return Err(UsageError( + "No metadata found while fetching the next row. Have you called the execute method or was the query supposed to return resultset?".to_string(), + )); + }; + if self.current_result_set_has_been_read_till_end { + return Ok(CursorPoll::Ready(false)); + } + if !matches!(self.active_row_read_state, ActiveRowReadState::Idle) + || self + .cancel_handle + .as_ref() + .is_some_and(|handle| handle.cancel_token.is_cancelled()) + { + return Ok(CursorPoll::Pending); + } + + // Resolving a CEK can call an async key-store provider. Keep encrypted + // result sets on the existing async path. + if !metadata.cek_table.is_empty() + || metadata + .columns + .iter() + .any(|column| column.crypto_metadata.is_some()) + { + return Ok(CursorPoll::Pending); + } + + let context = ParserContext::ColumnMetadata(metadata, None); + let start = self.request_timeout_start(); + let pause_state = self.transport.try_receive_row_header(&context)?; + let Some(pause_state) = pause_state else { + return Ok(CursorPoll::Pending); + }; + if let Some(start) = start { + self.update_remaining_timeout(start); + } + self.active_row_read_state = ActiveRowReadState::RowPaused(Box::new(pause_state)); + Ok(CursorPoll::Ready(true)) + } + /// Positions the cursor on the next row without decoding any column /// (ODBC `SQLFetch`). Returns `Ok(true)` when positioned on a row and /// `Ok(false)` when the result set is exhausted. @@ -3697,7 +3756,8 @@ impl TdsClient { /// After this returns `true`, individual columns are pulled with /// [`read_row_column`](Self::read_row_column). Any previously positioned row /// is drained first; its remaining column bytes are read and discarded rather - /// than returned to the caller. + /// than returned to the caller. Use this directly, or after + /// [`Self::try_next_row_cursor`] returns [`CursorPoll::Pending`]. #[instrument(skip(self), level = "info")] pub async fn next_row_cursor(&mut self) -> TdsResult { if self.current_metadata.is_none() { @@ -3751,38 +3811,76 @@ impl TdsClient { } } + /// Attempts to decode the next sequential column from buffered bytes. + /// + /// Returns [`CursorPoll::Pending`] without consuming bytes or cursor state + /// for PLP, encrypted, skipped, unsupported, or incomplete columns. The + /// caller then continues with [`Self::read_row_column`]. + pub fn try_read_row_column(&mut self, target: usize) -> TdsResult> { + let (next_column, column_count) = match &self.active_row_read_state { + ActiveRowReadState::Idle => return Ok(CursorPoll::Ready(CursorColumn::RowEnded)), + ActiveRowReadState::PlpPaused(_) => return Ok(CursorPoll::Pending), + ActiveRowReadState::RowPaused(pause_state) => { + (pause_state.next_column_index, pause_state.columns().len()) + } + }; + if self + .cancel_handle + .as_ref() + .is_some_and(|handle| handle.cancel_token.is_cancelled()) + { + return Ok(CursorPoll::Pending); + } + if target >= column_count { + return Err(UsageError(format!( + "read_row_column target column {target} is out of range (row has {column_count} columns)" + ))); + } + if target < next_column { + return Ok(CursorPoll::Ready(CursorColumn::AlreadyConsumed)); + } + if target != next_column { + // Skipping intervening columns can require arbitrarily shaped + // decoders. The async path remains authoritative for that case. + return Ok(CursorPoll::Pending); + } + + let start = self.request_timeout_start(); + let value = match &self.active_row_read_state { + ActiveRowReadState::RowPaused(pause_state) => self + .transport + .try_read_buffered_column(pause_state, target)?, + ActiveRowReadState::Idle | ActiveRowReadState::PlpPaused(_) => None, + }; + let Some(value) = value else { + return Ok(CursorPoll::Pending); + }; + if let Some(start) = start { + self.update_remaining_timeout(start); + } + + let ActiveRowReadState::RowPaused(mut pause_state) = + std::mem::replace(&mut self.active_row_read_state, ActiveRowReadState::Idle) + else { + return Err(crate::error::Error::ImplementationError( + "buffered column decode lost its row pause state".to_string(), + )); + }; + pause_state.next_column_index = target + 1; + if pause_state.next_column_index < column_count { + self.active_row_read_state = ActiveRowReadState::RowPaused(pause_state); + } + Ok(CursorPoll::Ready(CursorColumn::Value(value))) + } + /// Pulls column `target` (0-based) of the currently positioned row, /// skipping any intervening columns (ODBC `SQLGetData`). Forward-only: /// `target` must be at or after the cursor's next undecoded column. /// - /// Returns: - /// - [`CursorColumn::Value`] with the decoded value for non-PLP columns, - /// - [`CursorColumn::PlpStreaming`] when `target` is a PLP column whose - /// bytes must be pulled via [`read_active_plp_chunk`](Self::read_active_plp_chunk), - /// - [`CursorColumn::AlreadyConsumed`] if `target` was already read/skipped - /// (including after the whole row has been consumed), - /// - [`CursorColumn::RowEnded`] when no row is positioned. - /// - /// # Errors - /// - /// Returns `UsageError` when `target` is out of range **and** a row is - /// still positioned with unread columns (partially read). When no row is - /// positioned — including after the final column has been read — an - /// out-of-range or backward `target` yields [`CursorColumn::RowEnded`] - /// instead of an error. - /// - /// # Notes - /// - /// Reading the final column advances the cursor to idle, so a subsequent - /// backward or out-of-range request returns [`CursorColumn::RowEnded`]. A - /// caller that needs to distinguish "I rewound past the last column" from - /// "no row is positioned" must track the column it last read itself (as the - /// ODBC layer does). - /// - /// A `true` return from [`next_row_cursor`](Self::next_row_cursor) does not - /// guarantee `read_row_column(0)` yields a value: a zero-column row is - /// positioned with a column count of 0, so `read_row_column(0)` is - /// out-of-range and returns `UsageError`. + /// Returns [`CursorColumn::PlpStreaming`] for a PLP target, + /// [`CursorColumn::AlreadyConsumed`] for a backward target, and + /// [`CursorColumn::RowEnded`] when no row is positioned. Use this directly, + /// or after [`Self::try_read_row_column`] returns [`CursorPoll::Pending`]. // Avoid a span for every SQLGetData column; row and token events retain observability. pub async fn read_row_column(&mut self, target: usize) -> TdsResult { match std::mem::replace(&mut self.active_row_read_state, ActiveRowReadState::Idle) { @@ -4783,6 +4881,8 @@ mod tests { /// `read_row_column` down a specific arm (e.g. a `PlpPaused` result that /// makes the cursor emit `CursorColumn::PlpStreaming`). resume_results: VecDeque, + sync_header_available: bool, + sync_columns: VecDeque, } impl TestTransport { @@ -4795,6 +4895,8 @@ mod tests { packet_data: Vec::new(), packet_pos: 0, resume_results: VecDeque::new(), + sync_header_available: false, + sync_columns: VecDeque::new(), } } @@ -4807,6 +4909,8 @@ mod tests { packet_data: Vec::new(), packet_pos: 0, resume_results: VecDeque::new(), + sync_header_available: false, + sync_columns: VecDeque::new(), } } @@ -4819,6 +4923,8 @@ mod tests { packet_data, packet_pos: 0, resume_results: VecDeque::new(), + sync_header_available: false, + sync_columns: VecDeque::new(), } } @@ -4837,6 +4943,33 @@ mod tests { #[async_trait] impl TdsTokenStreamReader for TestTransport { + fn try_receive_row_header( + &mut self, + context: &ParserContext, + ) -> TdsResult> { + if !self.sync_header_available { + return Ok(None); + } + self.sync_header_available = false; + let ParserContext::ColumnMetadata(metadata, decryptor) = context else { + return Ok(None); + }; + Ok(Some(RowPauseState { + next_column_index: 0, + metadata: Arc::clone(metadata), + nbc_null_bitmap: None, + decryptor: decryptor.clone(), + })) + } + + fn try_read_buffered_column( + &mut self, + _pause_state: &RowPauseState, + _target: usize, + ) -> TdsResult> { + Ok(self.sync_columns.pop_front()) + } + async fn receive_token( &mut self, _context: &ParserContext, @@ -5227,6 +5360,28 @@ mod tests { Arc::new(ColMetadataToken::default()) } + fn int_column_metadata(column_count: usize) -> Arc { + let columns = (0..column_count) + .map(|index| crate::query::metadata::ColumnMetadata { + user_type: 0, + flags: 0, + type_info: crate::datatypes::sqldatatypes::TypeInfo::fixed_len( + crate::datatypes::sqldatatypes::TdsDataType::Int4, + ) + .unwrap(), + data_type: crate::datatypes::sqldatatypes::TdsDataType::Int4, + column_name: format!("c{index}"), + multi_part_name: None, + crypto_metadata: None, + }) + .collect(); + Arc::new(ColMetadataToken { + column_count: u16::try_from(column_count).unwrap(), + columns, + cek_table: vec![], + }) + } + fn done_more() -> Tokens { Tokens::Done(DoneToken { status: DoneStatus::MORE, @@ -5395,6 +5550,113 @@ mod tests { assert_eq!(resolved.duration(), None); } + #[test] + fn sync_cursor_attempt_reads_buffered_header_and_columns_in_order() { + let metadata = int_column_metadata(2); + let mut transport = TestTransport::new(); + transport.sync_header_available = true; + transport.sync_columns = VecDeque::from([ColumnValues::Int(10), ColumnValues::Int(20)]); + let mut client = create_test_client_with_transport(transport); + client.current_metadata = Some(metadata); + client.current_result_set_has_been_read_till_end = false; + + assert_eq!( + client.try_next_row_cursor().unwrap(), + CursorPoll::Ready(true) + ); + assert_eq!( + client.try_read_row_column(0).unwrap(), + CursorPoll::Ready(CursorColumn::Value(ColumnValues::Int(10))) + ); + assert_eq!( + client.try_read_row_column(1).unwrap(), + CursorPoll::Ready(CursorColumn::Value(ColumnValues::Int(20))) + ); + assert_eq!( + client.try_read_row_column(1).unwrap(), + CursorPoll::Ready(CursorColumn::RowEnded) + ); + } + + #[test] + fn sync_cursor_pending_leaves_cursor_state_untouched() { + let mut client = create_test_client(); + client.current_metadata = Some(int_column_metadata(1)); + client.current_result_set_has_been_read_till_end = false; + + assert_eq!(client.try_next_row_cursor().unwrap(), CursorPoll::Pending); + assert!(matches!( + client.active_row_read_state, + ActiveRowReadState::Idle + )); + + client.active_row_read_state = ActiveRowReadState::RowPaused(Box::new(RowPauseState { + next_column_index: 0, + metadata: int_column_metadata(1), + nbc_null_bitmap: None, + decryptor: None, + })); + assert_eq!(client.try_read_row_column(0).unwrap(), CursorPoll::Pending); + assert!(matches!( + client.active_row_read_state, + ActiveRowReadState::RowPaused(ref state) if state.next_column_index == 0 + )); + } + + #[test] + fn sync_cursor_ready_preserves_explicit_zero_timeout() { + let metadata = int_column_metadata(1); + let mut transport = TestTransport::new(); + transport.sync_header_available = true; + transport.sync_columns.push_back(ColumnValues::Int(42)); + let mut client = create_test_client_with_transport(transport); + client.current_metadata = Some(metadata); + client.current_result_set_has_been_read_till_end = false; + client.remaining_request_timeout = Some(Duration::ZERO); + + assert_eq!( + client.try_next_row_cursor().unwrap(), + CursorPoll::Ready(true) + ); + assert_eq!( + client.try_read_row_column(0).unwrap(), + CursorPoll::Ready(CursorColumn::Value(ColumnValues::Int(42))) + ); + assert_eq!(client.remaining_request_timeout, Some(Duration::ZERO)); + } + + #[test] + fn sync_cursor_defers_cancelled_operations_without_consuming_state() { + let metadata = int_column_metadata(1); + let mut transport = TestTransport::new(); + transport.sync_header_available = true; + transport.sync_columns.push_back(ColumnValues::Int(42)); + let mut client = create_test_client_with_transport(transport); + client.current_metadata = Some(Arc::clone(&metadata)); + client.current_result_set_has_been_read_till_end = false; + let cancellation = CancelHandle::new(); + client.cancel_handle = Some(cancellation.child_handle()); + cancellation.cancel(); + + assert_eq!(client.try_next_row_cursor().unwrap(), CursorPoll::Pending); + assert!(matches!( + client.active_row_read_state, + ActiveRowReadState::Idle + )); + + client.active_row_read_state = ActiveRowReadState::RowPaused(Box::new(RowPauseState { + next_column_index: 0, + metadata, + nbc_null_bitmap: None, + decryptor: None, + })); + assert_eq!(client.try_read_row_column(0).unwrap(), CursorPoll::Pending); + assert!(matches!( + client.active_row_read_state, + ActiveRowReadState::RowPaused(ref state) if state.next_column_index == 0 + )); + } + // ── PLP streaming lifecycle contract tests ── #[tokio::test] diff --git a/mssql-tds/src/connection/transport/buffers.rs b/mssql-tds/src/connection/transport/buffers.rs index 99ff5eff..6d34766d 100644 --- a/mssql-tds/src/connection/transport/buffers.rs +++ b/mssql-tds/src/connection/transport/buffers.rs @@ -166,6 +166,10 @@ impl TdsReadBuffer { pub(crate) fn get_slice(&self) -> &[u8] { &self.working_buffer[self.buffer_position..] } + + pub(crate) fn get_buffered_slice(&self) -> &[u8] { + &self.working_buffer[self.buffer_position..self.buffer_length] + } } impl Debug for TdsReadBuffer { diff --git a/mssql-tds/src/connection/transport/network_transport.rs b/mssql-tds/src/connection/transport/network_transport.rs index 236b48bb..10494050 100644 --- a/mssql-tds/src/connection/transport/network_transport.rs +++ b/mssql-tds/src/connection/transport/network_transport.rs @@ -11,6 +11,8 @@ use crate::connection_provider::tds_connection_provider::PARSER_REGISTRY; use crate::core::{ CancelHandle, EncryptionOptions, EncryptionSetting, NegotiatedEncryptionSetting, TdsResult, }; +use crate::datatypes::column_values::ColumnValues; +use crate::datatypes::decoder::GenericDecoder; use crate::datatypes::row_writer::RowWriter; use crate::error::Error::{OperationCancelledError, TimeoutError}; use crate::error::TimeoutErrorType; @@ -26,7 +28,7 @@ use crate::io::token_stream::{ use crate::message::attention::AttentionRequest; use crate::message::login_options::TdsVersion; use crate::message::messages::{PacketStatusFlags, Request, ResetConnectionMode}; -use crate::token::tokens::{DoneStatus, Tokens}; +use crate::token::tokens::{DoneStatus, TokenType, Tokens}; use async_trait::async_trait; use byteorder::{BigEndian, ByteOrder, LittleEndian}; use std::cmp::min; @@ -1401,6 +1403,91 @@ impl TdsPacketReader for NetworkTransport { #[async_trait] impl TdsTokenStreamReader for NetworkTransport { + fn try_receive_row_header( + &mut self, + context: &ParserContext, + ) -> TdsResult> { + let ParserContext::ColumnMetadata(metadata, decryptor) = context else { + return Err(crate::error::Error::ProtocolError( + "Expected ColumnMetadata in context for row decoding".to_string(), + )); + }; + let buffered = self.tds_read_buffer.get_buffered_slice(); + let Some(&token) = buffered.first() else { + return Ok(None); + }; + + if token == TokenType::Row as u8 { + self.tds_read_buffer.consume_bytes(1); + return Ok(Some(RowPauseState { + next_column_index: 0, + metadata: Arc::clone(metadata), + nbc_null_bitmap: None, + decryptor: decryptor.clone(), + })); + } + + if token != TokenType::NbcRow as u8 { + return Ok(None); + } + + let bitmap_len = metadata.columns.len().div_ceil(8); + let Some(bitmap_bytes) = buffered.get(1..1 + bitmap_len) else { + return Ok(None); + }; + let bitmap = if let Some(mut cached) = self.nbc_bitmap_scratch.take() + && cached.len() == bitmap_len + && let Some(buffer) = Arc::get_mut(&mut cached) + { + buffer.copy_from_slice(bitmap_bytes); + self.nbc_bitmap_scratch = Some(Arc::clone(&cached)); + cached + } else { + let bitmap: Arc<[u8]> = Arc::from(bitmap_bytes); + self.nbc_bitmap_scratch = Some(Arc::clone(&bitmap)); + bitmap + }; + self.tds_read_buffer.consume_bytes(1 + bitmap_len); + Ok(Some(RowPauseState { + next_column_index: 0, + metadata: Arc::clone(metadata), + nbc_null_bitmap: Some(bitmap), + decryptor: decryptor.clone(), + })) + } + + fn try_read_buffered_column( + &mut self, + pause_state: &RowPauseState, + target: usize, + ) -> TdsResult> { + if target != pause_state.next_column_index { + return Ok(None); + } + let Some(metadata) = pause_state.metadata.columns.get(target) else { + return Ok(None); + }; + if pause_state + .nbc_null_bitmap + .as_ref() + .is_some_and(|bitmap| bitmap[target / 8] & (1 << (target % 8)) != 0) + { + return Ok(Some(ColumnValues::Null)); + } + if pause_state.decryptor.is_some() { + return Ok(None); + } + + let decoder = GenericDecoder::default(); + let Some((value, used)) = + decoder.try_decode_buffered(self.tds_read_buffer.get_buffered_slice(), metadata)? + else { + return Ok(None); + }; + self.tds_read_buffer.consume_bytes(used); + Ok(Some(value)) + } + async fn receive_token( &mut self, context: &ParserContext, @@ -1628,11 +1715,15 @@ pub(crate) mod tests { use crate::connection::transport::network_transport::Stream; use crate::connection::transport::ssl_handler::SslHandler; use crate::core::EncryptionOptions; + use crate::datatypes::row_writer::DefaultRowWriter; + use crate::datatypes::sqldatatypes::{TdsDataType, TypeInfo}; use crate::message::messages::PacketType; + use crate::query::metadata::ColumnMetadata; use crate::test_packet_support::{ TestPacketBuilder, create_network_transport_with_chunked_data, create_network_transport_with_data, encode_utf16_le, }; + use crate::token::tokens::ColMetadataToken; use bytes::Bytes; use futures::SinkExt; use futures::StreamExt; @@ -1644,6 +1735,27 @@ pub(crate) mod tests { // The test would keep the payload lower than this size to make sure that the duplex stream can handle it. pub(crate) const MAX_BUFFER_SIZE: usize = 8192; + fn int4_row_context(column_count: usize) -> ParserContext { + ParserContext::ColumnMetadata( + Arc::new(ColMetadataToken { + column_count: u16::try_from(column_count).unwrap(), + columns: (0..column_count) + .map(|index| ColumnMetadata { + user_type: 0, + flags: 0, + type_info: TypeInfo::fixed_len(TdsDataType::Int4).unwrap(), + data_type: TdsDataType::Int4, + column_name: format!("value{index}"), + multi_part_name: None, + crypto_metadata: None, + }) + .collect(), + cek_table: vec![], + }), + None, + ) + } + impl Stream for DuplexStream { fn tls_handshake_starting(&mut self) { // No-op for duplex streams @@ -2762,6 +2874,158 @@ pub(crate) mod tests { assert_eq!(reader.read_uint32().await.unwrap(), 0x4433_2211); } + #[tokio::test] + async fn buffered_cursor_reads_complete_row_header_and_column() { + let expected = 0x1234_5678_i32; + let mut packet = TestPacketBuilder::new(PacketType::TabularResult); + let mut payload = vec![TokenType::Row as u8]; + payload.extend_from_slice(&expected.to_le_bytes()); + let mut reader = create_network_transport_with_data(&packet.append_bytes(&payload).build()); + reader.read_tds_packet().await.unwrap(); + + let pause_state = reader + .try_receive_row_header(&int4_row_context(1)) + .unwrap() + .expect("complete buffered row header"); + assert_eq!( + reader.try_read_buffered_column(&pause_state, 0).unwrap(), + Some(ColumnValues::Int(expected)) + ); + assert_eq!(reader.tds_read_buffer.get_remaining_byte_count(), 0); + } + + #[tokio::test] + async fn buffered_cursor_miss_preserves_bytes_for_async_continuation() { + let expected = 0x1234_5678_i32; + let value = expected.to_le_bytes(); + let mut first = TestPacketBuilder::new(PacketType::TabularResult); + let mut second = TestPacketBuilder::new(PacketType::TabularResult); + let mut first_payload = vec![TokenType::Row as u8]; + first_payload.extend_from_slice(&value[..2]); + let mut stream = first.append_bytes(&first_payload).build(); + stream.extend_from_slice(&second.append_bytes(&value[2..]).build()); + let mut reader = create_network_transport_with_data(&stream); + reader.read_tds_packet().await.unwrap(); + + let pause_state = reader + .try_receive_row_header(&int4_row_context(1)) + .unwrap() + .expect("row header is wholly buffered"); + assert_eq!(reader.tds_read_buffer.get_remaining_byte_count(), 2); + assert_eq!( + reader.try_read_buffered_column(&pause_state, 0).unwrap(), + None + ); + assert_eq!( + reader.tds_read_buffer.get_remaining_byte_count(), + 2, + "a miss must not consume the partial scalar" + ); + + let mut writer = DefaultRowWriter::new(1); + let result = reader + .resume_row_into( + pause_state, + None, + None, + ColumnPolicy::DecodeOne(0), + &mut writer, + ) + .await + .unwrap(); + assert!(matches!(result, RowReadResult::RowWritten)); + assert_eq!(writer.take_row(), vec![ColumnValues::Int(expected)]); + } + + #[tokio::test] + async fn buffered_nbcrow_null_column_needs_no_payload_bytes() { + let mut packet = TestPacketBuilder::new(PacketType::TabularResult); + let payload = [TokenType::NbcRow as u8, 0b0000_0001]; + let mut reader = create_network_transport_with_data(&packet.append_bytes(&payload).build()); + reader.read_tds_packet().await.unwrap(); + + let pause_state = reader + .try_receive_row_header(&int4_row_context(1)) + .unwrap() + .expect("complete NBCROW header"); + assert_eq!( + reader.try_read_buffered_column(&pause_state, 0).unwrap(), + Some(ColumnValues::Null) + ); + assert_eq!(reader.tds_read_buffer.get_remaining_byte_count(), 0); + } + + #[tokio::test] + async fn buffered_nbcrow_reuses_unaliased_bitmap_allocation() { + let mut packet = TestPacketBuilder::new(PacketType::TabularResult); + let payload = [ + TokenType::NbcRow as u8, + 0b0000_0001, + 0b0000_0010, + TokenType::NbcRow as u8, + 0b0000_0100, + 0b0000_1000, + ]; + let mut reader = create_network_transport_with_data(&packet.append_bytes(&payload).build()); + reader.read_tds_packet().await.unwrap(); + let context = int4_row_context(9); + + let first = reader + .try_receive_row_header(&context) + .unwrap() + .expect("first NBCROW header"); + let first_bitmap = first.nbc_null_bitmap.as_ref().expect("first bitmap"); + assert_eq!(first_bitmap.as_ref(), &[0b0000_0001, 0b0000_0010]); + let first_allocation = first_bitmap.as_ptr(); + drop(first); + + let second = reader + .try_receive_row_header(&context) + .unwrap() + .expect("second NBCROW header"); + let second_bitmap = second.nbc_null_bitmap.as_ref().expect("second bitmap"); + assert_eq!(second_bitmap.as_ref(), &[0b0000_0100, 0b0000_1000]); + assert_eq!( + second_bitmap.as_ptr(), + first_allocation, + "the uniquely owned scratch bitmap should be refilled in place" + ); + } + + #[tokio::test] + async fn buffered_nbcrow_bitmap_miss_preserves_header_for_async_continuation() { + let mut first = TestPacketBuilder::new(PacketType::TabularResult); + let mut second = TestPacketBuilder::new(PacketType::TabularResult); + let mut stream = first.append_bytes(&[TokenType::NbcRow as u8, 0]).build(); + stream.extend_from_slice(&second.append_bytes(&[0]).build()); + let mut reader = create_network_transport_with_data(&stream); + reader.read_tds_packet().await.unwrap(); + let context = int4_row_context(9); + + assert!(reader.try_receive_row_header(&context).unwrap().is_none()); + assert_eq!( + reader.tds_read_buffer.get_remaining_byte_count(), + 2, + "the token and partial bitmap must remain buffered" + ); + + let header = reader + .receive_row_header(&context, None, None) + .await + .unwrap(); + let RowHeader::Positioned(pause_state) = header else { + panic!("expected an NBCROW position"); + }; + assert_eq!( + pause_state + .nbc_null_bitmap + .as_ref() + .expect("NBCROW bitmap") + .as_ref(), + &[0, 0] + ); + } + #[tokio::test] async fn test_sync_scalar_probe_fallback_across_packet_boundaries() { let expected_uint16 = 0x1234u16; diff --git a/mssql-tds/src/datatypes/decoder.rs b/mssql-tds/src/datatypes/decoder.rs index 3a37bc41..d86e5e00 100644 --- a/mssql-tds/src/datatypes/decoder.rs +++ b/mssql-tds/src/datatypes/decoder.rs @@ -181,6 +181,73 @@ pub(crate) struct GenericDecoder { string_decoder: StringDecoder, } +struct BufferedSlice<'a> { + bytes: &'a [u8], + position: usize, +} + +impl<'a> BufferedSlice<'a> { + fn new(bytes: &'a [u8]) -> Self { + Self { bytes, position: 0 } + } + + fn take(&mut self) -> Option<[u8; N]> { + let end = self.position.checked_add(N)?; + let value = self.bytes.get(self.position..end)?.try_into().ok()?; + self.position = end; + Some(value) + } + + fn take_bytes(&mut self, len: usize) -> Option> { + let end = self.position.checked_add(len)?; + let value = self.bytes.get(self.position..end)?.to_vec(); + self.position = end; + Some(value) + } + + fn byte(&mut self) -> Option { + self.take().map(|[value]| value) + } + + fn i16(&mut self) -> Option { + self.take().map(i16::from_le_bytes) + } + + fn u16(&mut self) -> Option { + self.take().map(u16::from_le_bytes) + } + + fn u24(&mut self) -> Option { + let [b0, b1, b2] = self.take()?; + Some(u32::from_le_bytes([b0, b1, b2, 0])) + } + + fn i32(&mut self) -> Option { + self.take().map(i32::from_le_bytes) + } + + fn u32(&mut self) -> Option { + self.take().map(u32::from_le_bytes) + } + + fn u40(&mut self) -> Option { + let [b0, b1, b2, b3, b4] = self.take()?; + Some(u64::from_le_bytes([b0, b1, b2, b3, b4, 0, 0, 0])) + } + + fn i64(&mut self) -> Option { + self.take().map(i64::from_le_bytes) + } + + fn f32(&mut self) -> Option { + self.take().map(f32::from_le_bytes) + } + + fn f64(&mut self) -> Option { + self.take().map(f64::from_le_bytes) + } +} + #[cfg_attr(not(test), allow(dead_code))] #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(crate) enum PlpChunkReadLength { @@ -508,6 +575,334 @@ impl PlpColumnStream { } impl GenericDecoder { + /// Decodes one complete non-PLP column from `bytes` without awaiting. + /// + /// `None` means the buffered bytes are incomplete or this type still + /// requires the async decoder. The caller must consume `used` bytes only + /// after receiving `Some`, so a miss is non-destructive. + /// + /// Each supported arm must remain wire-compatible with [`Self::decode_into`]. + pub(crate) fn try_decode_buffered( + &self, + bytes: &[u8], + metadata: &ColumnMetadata, + ) -> TdsResult> { + if metadata.is_plp() || metadata.crypto_metadata.is_some() { + return Ok(None); + } + + let mut reader = BufferedSlice::new(bytes); + let value = match metadata.data_type { + TdsDataType::Int1 => ColumnValues::TinyInt(match reader.byte() { + Some(value) => value, + None => return Ok(None), + }), + TdsDataType::Int2 => ColumnValues::SmallInt(match reader.i16() { + Some(value) => value, + None => return Ok(None), + }), + TdsDataType::Int4 => ColumnValues::Int(match reader.i32() { + Some(value) => value, + None => return Ok(None), + }), + TdsDataType::Int8 => ColumnValues::BigInt(match reader.i64() { + Some(value) => value, + None => return Ok(None), + }), + TdsDataType::Flt4 => ColumnValues::Real(match reader.f32() { + Some(value) => value, + None => return Ok(None), + }), + TdsDataType::Flt8 => ColumnValues::Float(match reader.f64() { + Some(value) => value, + None => return Ok(None), + }), + TdsDataType::Bit => ColumnValues::Bit(match reader.byte() { + Some(value) => value == 1, + None => return Ok(None), + }), + TdsDataType::IntN => { + let Some(length) = reader.byte() else { + return Ok(None); + }; + match length { + 0 => ColumnValues::Null, + 1 => ColumnValues::TinyInt(match reader.byte() { + Some(value) => value, + None => return Ok(None), + }), + 2 => ColumnValues::SmallInt(match reader.i16() { + Some(value) => value, + None => return Ok(None), + }), + 4 => ColumnValues::Int(match reader.i32() { + Some(value) => value, + None => return Ok(None), + }), + 8 => ColumnValues::BigInt(match reader.i64() { + Some(value) => value, + None => return Ok(None), + }), + _ => { + return Err(crate::error::Error::ProtocolError(format!( + "Invalid IntN length: {length}" + ))); + } + } + } + TdsDataType::FltN => { + let Some(length) = reader.byte() else { + return Ok(None); + }; + match length { + 0 => ColumnValues::Null, + 4 => ColumnValues::Real(match reader.f32() { + Some(value) => value, + None => return Ok(None), + }), + _ => ColumnValues::Float(match reader.f64() { + Some(value) => value, + None => return Ok(None), + }), + } + } + TdsDataType::BitN => { + let Some(length) = reader.byte() else { + return Ok(None); + }; + match length { + 0 => ColumnValues::Null, + _ => ColumnValues::Bit(match reader.byte() { + Some(value) => value == 1, + None => return Ok(None), + }), + } + } + TdsDataType::Money4 => ColumnValues::SmallMoney(SqlSmallMoney { + int_val: match reader.i32() { + Some(value) => value, + None => return Ok(None), + }, + }), + TdsDataType::Money => { + let Some(msb_part) = reader.i32() else { + return Ok(None); + }; + let Some(lsb_part) = reader.i32() else { + return Ok(None); + }; + ColumnValues::Money(SqlMoney { lsb_part, msb_part }) + } + TdsDataType::MoneyN => { + let Some(length) = reader.byte() else { + return Ok(None); + }; + match length { + 0 => ColumnValues::Null, + 4 => ColumnValues::SmallMoney(SqlSmallMoney { + int_val: match reader.i32() { + Some(value) => value, + None => return Ok(None), + }, + }), + 8 => { + let Some(msb_part) = reader.i32() else { + return Ok(None); + }; + let Some(lsb_part) = reader.i32() else { + return Ok(None); + }; + ColumnValues::Money(SqlMoney { lsb_part, msb_part }) + } + _ => { + return Err(crate::error::Error::ProtocolError(format!( + "Invalid MoneyN length: {length}" + ))); + } + } + } + TdsDataType::DateTim4 => { + let Some(days) = reader.u16() else { + return Ok(None); + }; + let Some(time) = reader.u16() else { + return Ok(None); + }; + ColumnValues::SmallDateTime(SqlSmallDateTime { days, time }) + } + TdsDataType::DateTime => { + let Some(days) = reader.i32() else { + return Ok(None); + }; + let Some(time) = reader.u32() else { + return Ok(None); + }; + ColumnValues::DateTime(SqlDateTime { days, time }) + } + TdsDataType::DateTimeN => { + let Some(length) = reader.byte() else { + return Ok(None); + }; + match length { + 0 => ColumnValues::Null, + 4 => { + let Some(days) = reader.u16() else { + return Ok(None); + }; + let Some(time) = reader.u16() else { + return Ok(None); + }; + ColumnValues::SmallDateTime(SqlSmallDateTime { days, time }) + } + _ => { + let Some(days) = reader.i32() else { + return Ok(None); + }; + let Some(time) = reader.u32() else { + return Ok(None); + }; + ColumnValues::DateTime(SqlDateTime { days, time }) + } + } + } + TdsDataType::DateN => { + let Some(length) = reader.byte() else { + return Ok(None); + }; + match length { + 0 => ColumnValues::Null, + _ => ColumnValues::Date(SqlDate::unchecked_create(match reader.u24() { + Some(value) => value, + None => return Ok(None), + })), + } + } + TdsDataType::TimeN | TdsDataType::DateTime2N | TdsDataType::DateTimeOffsetN => { + let Some(length) = reader.byte() else { + return Ok(None); + }; + if length == 0 { + ColumnValues::Null + } else { + let scale = metadata.get_scale().ok_or_else(|| { + crate::error::Error::ImplementationError(format!( + "{:?} type should have scale", + metadata.data_type + )) + })?; + let trailing = match metadata.data_type { + TdsDataType::TimeN => 0, + TdsDataType::DateTime2N => 3, + TdsDataType::DateTimeOffsetN => 5, + _ => { + return Err(crate::error::Error::ImplementationError( + "unexpected buffered time type".to_string(), + )); + } + }; + let time_length = length.checked_sub(trailing).ok_or_else(|| { + crate::error::Error::ProtocolError(format!( + "Invalid {:?} length: {length}", + metadata.data_type + )) + })?; + let scaled = match time_length { + 3 => reader.u24().map(u64::from), + 4 => reader.u32().map(u64::from), + _ => reader.u40(), + }; + let Some(scaled) = scaled else { + return Ok(None); + }; + let time = SqlTime { + time_nanoseconds: scaled * 10_u64.pow(u32::from(7 - scale.min(7))), + scale, + }; + match metadata.data_type { + TdsDataType::TimeN => ColumnValues::Time(time), + TdsDataType::DateTime2N => { + let Some(days) = reader.u24() else { + return Ok(None); + }; + ColumnValues::DateTime2(SqlDateTime2 { days, time }) + } + TdsDataType::DateTimeOffsetN => { + let Some(days) = reader.u24() else { + return Ok(None); + }; + let Some(offset) = reader.i16() else { + return Ok(None); + }; + ColumnValues::DateTimeOffset(SqlDateTimeOffset { + datetime2: SqlDateTime2 { days, time }, + offset, + }) + } + _ => { + return Err(crate::error::Error::ImplementationError( + "unexpected buffered time type".to_string(), + )); + } + } + } + } + TdsDataType::Guid => { + let Some(length) = reader.byte() else { + return Ok(None); + }; + if length == 0 { + ColumnValues::Null + } else { + if length != 16 { + return Err(crate::error::Error::ProtocolError(format!( + "Invalid GUID length: expected 16 bytes, got {length}" + ))); + } + let Some(bytes) = reader.take_bytes(16) else { + return Ok(None); + }; + ColumnValues::Uuid(uuid::Uuid::from_slice_le(&bytes).map_err(|error| { + crate::error::Error::ProtocolError(format!("Failed to parse UUID: {error}")) + })?) + } + } + TdsDataType::NChar + | TdsDataType::NVarChar + | TdsDataType::BigChar + | TdsDataType::BigVarChar + | TdsDataType::Char + | TdsDataType::VarChar => { + let Some(length) = reader.u16() else { + return Ok(None); + }; + if length == u16::MAX { + ColumnValues::Null + } else { + let Some(bytes) = reader.take_bytes(usize::from(length)) else { + return Ok(None); + }; + ColumnValues::String(SqlString::new(bytes, get_encoding_type(metadata))) + } + } + TdsDataType::BigBinary | TdsDataType::BigVarBinary => { + let Some(length) = reader.u16() else { + return Ok(None); + }; + if length == u16::MAX { + ColumnValues::Null + } else { + let Some(bytes) = reader.take_bytes(usize::from(length)) else { + return Ok(None); + }; + ColumnValues::Bytes(bytes) + } + } + _ => return Ok(None), + }; + + Ok(Some((value, reader.position))) + } + #[cfg(test)] const SHORTLEN_MAXVALUE: usize = 65535; const SQL_PLP_NULL: usize = 0xffffffffffffffff; @@ -3282,7 +3677,9 @@ mod test { use byteorder::{ByteOrder, LittleEndian}; use crate::core::TdsResult; - use crate::datatypes::column_values::{ColumnValues, SqlDateTime, SqlSmallDateTime}; + use crate::datatypes::column_values::{ + ColumnValues, SqlDateTime, SqlMoney, SqlSmallDateTime, SqlTime, + }; use crate::datatypes::decoder::{ GenericDecoder, MAX_PLP_SIZE, PlpChunkReadLength, PlpChunkStreamReader, PlpColumnStream, SqlTypeDecode, @@ -3439,6 +3836,114 @@ mod test { } } + fn buffered_value(bytes: &[u8], metadata: &ColumnMetadata) -> (ColumnValues, usize) { + GenericDecoder::default() + .try_decode_buffered(bytes, metadata) + .unwrap() + .expect("complete buffered value") + } + + #[test] + fn buffered_decode_intn_handles_null_value_and_partial_payload() { + let metadata = varlen_metadata(TdsDataType::IntN, 4); + assert_eq!(buffered_value(&[0], &metadata), (ColumnValues::Null, 1)); + assert_eq!( + GenericDecoder::default() + .try_decode_buffered(&[4, 0x78, 0x56], &metadata) + .unwrap(), + None + ); + assert_eq!( + buffered_value(&[4, 0x78, 0x56, 0x34, 0x12], &metadata), + (ColumnValues::Int(0x1234_5678), 5) + ); + } + + #[test] + fn buffered_decode_money_preserves_wire_word_order() { + let metadata = fixed_metadata(TdsDataType::Money, 8); + let mut bytes = 7_i32.to_le_bytes().to_vec(); + bytes.extend_from_slice(&11_i32.to_le_bytes()); + assert_eq!( + buffered_value(&bytes, &metadata), + ( + ColumnValues::Money(SqlMoney { + lsb_part: 11, + msb_part: 7, + }), + 8, + ) + ); + } + + #[test] + fn buffered_decode_time_applies_fractional_scale() { + let metadata = ColumnMetadata { + user_type: 0, + flags: 0, + data_type: TdsDataType::TimeN, + type_info: TypeInfo::var_len_scale(TdsDataType::TimeN, 4, 3).unwrap(), + column_name: String::new(), + multi_part_name: None, + crypto_metadata: None, + }; + let mut bytes = vec![4]; + bytes.extend_from_slice(&1234_u32.to_le_bytes()); + assert_eq!( + buffered_value(&bytes, &metadata), + ( + ColumnValues::Time(SqlTime { + time_nanoseconds: 12_340_000, + scale: 3, + }), + 5, + ) + ); + } + + #[test] + fn buffered_decode_guid_uses_tds_little_endian_layout() { + let metadata = varlen_metadata(TdsDataType::Guid, 16); + let expected = uuid::Uuid::from_u128(0x0011_2233_4455_6677_8899_aabb_ccdd_eeff); + let mut bytes = vec![16]; + bytes.extend_from_slice(&expected.to_bytes_le()); + assert_eq!( + buffered_value(&bytes, &metadata), + (ColumnValues::Uuid(expected), 17) + ); + } + + #[test] + fn buffered_decode_nvarchar_waits_for_the_complete_payload() { + let metadata = ColumnMetadata { + user_type: 0, + flags: 0, + data_type: TdsDataType::NVarChar, + type_info: TypeInfo::var_len_string(TdsDataType::NVarChar, 100, None).unwrap(), + column_name: String::new(), + multi_part_name: None, + crypto_metadata: None, + }; + let payload = "row_1" + .encode_utf16() + .flat_map(u16::to_le_bytes) + .collect::>(); + let mut bytes = u16::try_from(payload.len()).unwrap().to_le_bytes().to_vec(); + bytes.extend_from_slice(&payload); + assert_eq!( + GenericDecoder::default() + .try_decode_buffered(&bytes[..bytes.len() - 1], &metadata) + .unwrap(), + None + ); + let (value, used) = buffered_value(&bytes, &metadata); + let ColumnValues::String(value) = value else { + panic!("expected string"); + }; + assert_eq!(value.to_utf8_string(), "row_1"); + assert_eq!(used, bytes.len()); + } + /// Runs both decode() and decode_into() on the same bytes and asserts /// that decode_into via DefaultRowWriter produces the same ColumnValues /// as decode(). diff --git a/mssql-tds/src/io/token_stream.rs b/mssql-tds/src/io/token_stream.rs index cad0d504..2daeafda 100644 --- a/mssql-tds/src/io/token_stream.rs +++ b/mssql-tds/src/io/token_stream.rs @@ -2,6 +2,7 @@ // Licensed under the MIT License. use crate::core::{CancelHandle, TdsResult}; +use crate::datatypes::column_values::ColumnValues; use crate::datatypes::decoder::{GenericDecoder, PlpColumnStream, decrypt_encrypted_column}; use crate::datatypes::row_writer::{DiscardRowWriter, RowWriter, write_column_value}; use crate::io::packet_reader::TdsPacketReader; @@ -192,6 +193,27 @@ impl PlpPauseState { #[async_trait] #[cfg(not(fuzzing))] pub(crate) trait TdsTokenStreamReader { + /// Attempts to read a complete row header from bytes already buffered by the + /// transport. Returns `None` without consuming bytes when async I/O or an + /// unsupported synchronous parser is required. + fn try_receive_row_header( + &mut self, + _context: &ParserContext, + ) -> TdsResult> { + Ok(None) + } + + /// Attempts to decode `target` from bytes already buffered by the transport. + /// Returns `None` without consuming bytes when async I/O or an unsupported + /// synchronous decoder is required. + fn try_read_buffered_column( + &mut self, + _pause_state: &RowPauseState, + _target: usize, + ) -> TdsResult> { + Ok(None) + } + async fn receive_token( &mut self, context: &ParserContext, @@ -247,6 +269,21 @@ pub(crate) trait TdsTokenStreamReader { #[async_trait] #[cfg(fuzzing)] pub trait TdsTokenStreamReader { + fn try_receive_row_header( + &mut self, + _context: &ParserContext, + ) -> TdsResult> { + Ok(None) + } + + fn try_read_buffered_column( + &mut self, + _pause_state: &RowPauseState, + _target: usize, + ) -> TdsResult> { + Ok(None) + } + async fn receive_token( &mut self, context: &ParserContext,