From fd74a7ca49e2f3a2ac4ebc1c01d5c9b710bce053 Mon Sep 17 00:00:00 2001 From: saurabh500 <1623701+saurabh500@users.noreply.github.com> Date: Sat, 15 Aug 2026 17:49:07 +0000 Subject: [PATCH] Add experimental TDS performance baseline Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 5bcfe214-06ab-401e-9b36-bbba2e053d19 --- mssql-tds/src/connection/tds_client.rs | 195 ++++++++++----- mssql-tds/src/connection/transport/buffers.rs | 134 +++++++++++ .../connection/transport/network_transport.rs | 224 +++++++++++++---- mssql-tds/src/datatypes/decoder.rs | 225 ++++++++++-------- mssql-tds/src/io/packet_reader.rs | 120 ++++++++++ 5 files changed, 694 insertions(+), 204 deletions(-) diff --git a/mssql-tds/src/connection/tds_client.rs b/mssql-tds/src/connection/tds_client.rs index 9743dc8c..2602ea1b 100644 --- a/mssql-tds/src/connection/tds_client.rs +++ b/mssql-tds/src/connection/tds_client.rs @@ -37,8 +37,8 @@ use crate::{ message::{batch::SqlBatch, messages::Request}, token::tokens::{ColMetadataToken, CurrentCommand, DoneStatus, EnvChangeTokenSubType, Tokens}, }; -use async_trait::async_trait; use std::collections::HashMap; +use std::future::Future; use std::num::NonZeroU32; use tracing::{debug, error, info, instrument, warn}; @@ -2781,7 +2781,8 @@ impl TdsClient { /// This functions returns to the next row in the result set. /// If there are no more rows, it returns None. - #[instrument(skip(self), level = "info")] + // Not instrumented: the span pushes ResultSet::next_row over the 4 KiB + // hot-path future budget. Successful rows still emit `Row Received`. pub(crate) async fn get_next_row(&mut self) -> TdsResult>> { let col_count = self .current_metadata @@ -2789,7 +2790,7 @@ impl TdsClient { .map(|m| m.columns.len()) .unwrap_or(0); let mut writer = DefaultRowWriter::new(col_count); - if self.get_next_row_into(&mut writer).await? { + if self.next_row_into(&mut writer).await? { Ok(Some(writer.take_row())) } else { Ok(None) @@ -3587,24 +3588,25 @@ impl TdsClient { /// /// Uses `receive_row_into` to decode ROW/NBCROW tokens directly through /// `decode_into`, bypassing the intermediate `RowToken { all_values }`. - #[instrument(skip(self, writer), level = "info")] - pub(crate) async fn get_next_row_into( - &mut self, - writer: &mut (dyn RowWriter + Send), - ) -> TdsResult { + /// Concrete writers stay concrete through the production transport and + /// decode chain. [`ResultSet::next_row_into`] provides the same operation + /// through statically dispatched trait calls. + // `#[instrument]` adds enough state to exceed the 4096 B budget once the + // lazy timeout future is inlined. Successful rows still emit `Row Received`. + pub async fn next_row_into(&mut self, writer: &mut (dyn RowWriter + Send)) -> TdsResult { + // Every error return below must abort the pending prepare capture. A + // wrapper that centralizes this cleanup exceeds the row-future size budget. + // End-of-set reads are idempotent even after advancing clears metadata. + if self.current_result_set_has_been_read_till_end { + return Ok(false); + } if self.current_metadata.is_none() { + self.abort_pending_prepare_capture(); 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(), )); } - // Idempotent at end-of-set: after the terminating DONE, calling again - // must report exhaustion instead of blocking on a wire read for a - // packet the server will never send (until the caller advances). - if self.current_result_set_has_been_read_till_end { - return Ok(false); - } - // The push path decodes whole rows and never pauses, so it must not run // while the pull cursor has a row *partially* read: silently draining it // would discard that row and return the *next* one, mapping the caller's @@ -3615,21 +3617,26 @@ impl TdsClient { self.active_row_read_state, ActiveRowReadState::RowPaused(_) | ActiveRowReadState::PlpPaused(_) ) { + self.abort_pending_prepare_capture(); return Err(UsageError( - "get_next_row_into called while a pull-cursor row is still active; \ + "next_row_into called while a pull-cursor row is still active; \ advance the cursor with next_row_cursor before using the push row API" .to_string(), )); } - self.drain_active_row().await?; - let metadata = Arc::clone(self.current_metadata.as_ref().unwrap()); - let decryptor = self.resolve_cell_decryptor(&metadata).await?; + let decryptor = match self.resolve_cell_decryptor(&metadata).await { + Ok(decryptor) => decryptor, + Err(error) => { + self.abort_pending_prepare_capture(); + return Err(error); + } + }; let parser_context = ParserContext::ColumnMetadata(metadata, decryptor); loop { let start = Instant::now(); - let result = self + let result = match self .transport .receive_row_into( &parser_context, @@ -3638,7 +3645,14 @@ impl TdsClient { ColumnPolicy::DecodeAll, writer, ) - .await?; + .await + { + Ok(result) => result, + Err(error) => { + self.abort_pending_prepare_capture(); + return Err(error); + } + }; self.update_remaining_timeout(start); match result { @@ -3649,13 +3663,21 @@ impl TdsClient { } RowReadResult::RowPaused(_) | RowReadResult::PlpPaused(_) => { // DecodeAll never pauses; a pause here is a protocol/logic error. + self.abort_pending_prepare_capture(); return Err(crate::error::Error::ProtocolError( "Unexpected pause while decoding a full row (ColumnPolicy::DecodeAll)" .to_string(), )); } RowReadResult::Token(token) => { - if let Some(has_row) = self.handle_row_read_token(token).await? { + let handled = match Box::pin(self.handle_row_read_token(token)).await { + Ok(handled) => handled, + Err(error) => { + self.abort_pending_prepare_capture(); + return Err(error); + } + }; + if let Some(has_row) = handled { return Ok(has_row); } } @@ -3870,7 +3892,7 @@ impl TdsClient { /// The scratch buffer is heap-allocated rather than a stack array: it is live /// across the await below, so a stack array would be stored inline in this /// future and propagate into every caller that awaits it — `read_row_column` - /// directly, plus `drain_rows`, `get_next_row_into` and `next_row_cursor` via + /// directly, plus `drain_rows`, `next_row_into` and `next_row_cursor` via /// `drain_active_row`. Abandoning a partially read PLP column is rare and /// already network-bound, so one allocation there is negligible; an 8 KiB /// per-row state machine is not. @@ -4450,7 +4472,6 @@ impl TdsClient { } } -#[async_trait] impl ResultSet for TdsClient { fn get_metadata(&self) -> &Vec { // If no metadata is available, return an empty vector @@ -4462,39 +4483,23 @@ impl ResultSet for TdsClient { .unwrap_or(&self.empty_metadata) } - #[instrument(skip(self), level = "info")] - async fn next_row(&mut self) -> TdsResult>> { - let result = if self.maybe_has_unread_rows() { - self.get_next_row().await - } else { - Ok(None) - }; - if result.is_err() { - self.abort_pending_prepare_capture(); - } - result + fn next_row(&mut self) -> impl Future>>> + Send { + self.get_next_row() } - #[instrument(skip(self, writer), level = "info")] - async fn next_row_into(&mut self, writer: &mut (dyn RowWriter + Send)) -> TdsResult { - let result = if self.maybe_has_unread_rows() { - self.get_next_row_into(writer).await - } else { - Ok(false) - }; - if result.is_err() { - self.abort_pending_prepare_capture(); - } - result + fn next_row_into( + &mut self, + writer: &mut (dyn RowWriter + Send), + ) -> impl Future> + Send { + TdsClient::next_row_into(self, writer) } fn maybe_has_unread_rows(&self) -> bool { !self.current_result_set_has_been_read_till_end } - #[instrument(skip(self), level = "info")] - async fn close(&mut self) -> TdsResult<()> { - self.close_query().await + fn close(&mut self) -> impl Future> + Send { + self.close_query() } } @@ -4686,8 +4691,18 @@ enum ResultBoundaryKind { End, } -/// Async result set iteration. -#[async_trait] +/// Async result set iteration through statically dispatched futures. +/// +/// The returned futures are native, unboxed futures with an explicit [`Send`] +/// guarantee. +/// +/// # Dyn compatibility +/// +/// This trait is intentionally not dyn-compatible and cannot be used through +/// `dyn ResultSet`. This is a breaking change for trait-object consumers and for +/// implementations written with `#[async_trait]`; concrete call sites can keep +/// awaiting the methods unchanged, while implementations must return native +/// `Send` futures. pub trait ResultSet { /// Returns the metadata of the result set. /// This metadata includes information about the columns in the result set. @@ -4695,7 +4710,7 @@ pub trait ResultSet { /// Returns the next row of data as a vector of column values. /// If there is no more data, it returns None. - async fn next_row(&mut self) -> TdsResult>>; + fn next_row(&mut self) -> impl Future>>> + Send; /// Decodes the next row directly into a [`RowWriter`], returning `true` if /// a row was written or `false` when the result set is exhausted. @@ -4710,14 +4725,17 @@ pub trait ResultSet { /// partially read. Draining that row here would silently discard it and /// return the *next* one, so callers must first finish the row with /// `next_row_cursor`. A fully-consumed or absent row is fine. - async fn next_row_into(&mut self, writer: &mut (dyn RowWriter + Send)) -> TdsResult; + fn next_row_into( + &mut self, + writer: &mut (dyn RowWriter + Send), + ) -> impl Future> + Send; /// Returns `true` if the result set may still contain unread rows. fn maybe_has_unread_rows(&self) -> bool; /// Iterates over the result set, and marks it as closed. After calling close, the next_row method, /// will always return None. - async fn close(&mut self) -> TdsResult<()>; + fn close(&mut self) -> impl Future> + Send; } #[cfg(test)] @@ -5046,7 +5064,9 @@ mod tests { let next_row_cursor = std::mem::size_of_val(&client.next_row_cursor()); let read_row_column = std::mem::size_of_val(&client.read_row_column(0)); let drain_rows = std::mem::size_of_val(&client.drain_rows()); - let get_next_row_into = std::mem::size_of_val(&client.get_next_row_into(&mut sink)); + let next_row_into = std::mem::size_of_val(&client.next_row_into(&mut sink)); + let next_row_into_dyn = + std::mem::size_of_val(&client.next_row_into(&mut sink as &mut (dyn RowWriter + Send))); let read_active_plp_chunk = std::mem::size_of_val(&client.read_active_plp_chunk(&mut plp_out)); @@ -5054,7 +5074,8 @@ mod tests { ("next_row_cursor", next_row_cursor), ("read_row_column", read_row_column), ("drain_rows", drain_rows), - ("get_next_row_into", get_next_row_into), + ("next_row_into", next_row_into), + ("next_row_into (dyn writer)", next_row_into_dyn), ("read_active_plp_chunk", read_active_plp_chunk), ] { assert!( @@ -5062,6 +5083,34 @@ mod tests { "{name} future is {size} B, expected <= {MAX} B" ); } + + let native_next_row = std::mem::size_of_val(&client.get_next_row()); + let result_set_next_row = std::mem::size_of_val(&ResultSet::next_row(&mut client)); + let native_next_row_into = + std::mem::size_of_val(&client.next_row_into(&mut sink as &mut (dyn RowWriter + Send))); + let result_set_next_row_into = + std::mem::size_of_val(&ResultSet::next_row_into(&mut client, &mut sink)); + + assert_eq!( + result_set_next_row, native_next_row, + "ResultSet::next_row must forward the native future without boxing" + ); + assert_eq!( + result_set_next_row_into, native_next_row_into, + "ResultSet::next_row_into must forward the native future without boxing" + ); + + // `close` is not checked against the per-row budget because it runs only + // once per result set; its larger future does not affect row iteration. + for (name, size) in [ + ("ResultSet::next_row", result_set_next_row), + ("ResultSet::next_row_into", result_set_next_row_into), + ] { + assert!( + size <= MAX, + "{name} future is {size} B, expected <= {MAX} B" + ); + } } #[test] @@ -5448,7 +5497,7 @@ mod tests { } #[tokio::test] - async fn get_next_row_into_rejects_active_pull_cursor_row() { + async fn next_row_into_rejects_active_pull_cursor_row() { // A row parked by the pull cursor (`next_row_cursor`) must not be // silently drained by the push path. Mixing the two would discard the // parked row and hand back the *next* one, so the earlier @@ -5465,7 +5514,7 @@ mod tests { let mut sink = DiscardRowWriter; let err = client - .get_next_row_into(&mut sink) + .next_row_into(&mut sink) .await .expect_err("push path must reject a parked pull-cursor row"); assert!( @@ -5497,6 +5546,19 @@ mod tests { assert_eq!(TdsClient::normalize_param_name(""), ""); } + #[tokio::test] + async fn next_row_into_is_idempotent_after_end_without_metadata() { + let mut client = create_test_client(); + client.current_metadata = None; + client.current_result_set_has_been_read_till_end = true; + let mut sink = DiscardRowWriter; + + assert!( + !ResultSet::next_row_into(&mut client, &mut sink) + .await + .unwrap() + ); + } #[tokio::test] async fn consume_done_token_captures_all_info_tokens() { let mut client = create_test_client_with_tokens(vec![ @@ -6686,7 +6748,22 @@ mod tests { client.current_result_set_has_been_read_till_end = false; client.pending_capture = Some(sid(1)); - assert!(client.next_row().await.is_err()); + assert!(ResultSet::next_row(&mut client).await.is_err()); + assert!(client.pending_capture.is_none()); + } + + #[tokio::test] + async fn next_row_into_error_aborts_pending_prepare_capture() { + let mut client = create_test_client(); + client.current_result_set_has_been_read_till_end = false; + client.pending_capture = Some(sid(1)); + let mut writer = DiscardRowWriter; + + assert!( + ResultSet::next_row_into(&mut client, &mut writer) + .await + .is_err() + ); assert!(client.pending_capture.is_none()); } diff --git a/mssql-tds/src/connection/transport/buffers.rs b/mssql-tds/src/connection/transport/buffers.rs index 6ed265f9..99ff5eff 100644 --- a/mssql-tds/src/connection/transport/buffers.rs +++ b/mssql-tds/src/connection/transport/buffers.rs @@ -49,6 +49,72 @@ impl TdsReadBuffer { self.buffer_length - self.buffer_position } + #[inline(always)] + fn try_read_array(&mut self) -> Option<[u8; N]> { + if !self.do_we_have_enough_data(N) { + return None; + } + + let position = self.buffer_position; + let bytes = self.working_buffer[position..position + N] + .try_into() + .expect("slice length is fixed by N"); + self.consume_bytes(N); + Some(bytes) + } + + #[inline(always)] + pub(crate) fn try_read_byte(&mut self) -> Option { + self.try_read_array().map(|[value]| value) + } + + #[inline(always)] + pub(crate) fn try_read_int16(&mut self) -> Option { + self.try_read_array().map(i16::from_le_bytes) + } + + #[inline(always)] + pub(crate) fn try_read_uint16(&mut self) -> Option { + self.try_read_array().map(u16::from_le_bytes) + } + + #[inline(always)] + pub(crate) fn try_read_uint24(&mut self) -> Option { + let [b0, b1, b2] = self.try_read_array()?; + Some(u32::from_le_bytes([b0, b1, b2, 0])) + } + + #[inline(always)] + pub(crate) fn try_read_int32(&mut self) -> Option { + self.try_read_array().map(i32::from_le_bytes) + } + + #[inline(always)] + pub(crate) fn try_read_uint32(&mut self) -> Option { + self.try_read_array().map(u32::from_le_bytes) + } + + #[inline(always)] + pub(crate) fn try_read_uint40(&mut self) -> Option { + let [b0, b1, b2, b3, b4] = self.try_read_array()?; + Some(u64::from_le_bytes([b0, b1, b2, b3, b4, 0, 0, 0])) + } + + #[inline(always)] + pub(crate) fn try_read_int64(&mut self) -> Option { + self.try_read_array().map(i64::from_le_bytes) + } + + #[inline(always)] + pub(crate) fn try_read_float32(&mut self) -> Option { + self.try_read_array().map(f32::from_le_bytes) + } + + #[inline(always)] + pub(crate) fn try_read_float64(&mut self) -> Option { + self.try_read_array().map(f64::from_le_bytes) + } + pub(crate) fn consume_bytes(&mut self, byte_count: usize) { if byte_count > (self.buffer_length - self.buffer_position) { panic!("Not enough data to consume"); @@ -426,6 +492,74 @@ mod tests { assert!(!buf.do_we_have_enough_data(401)); } + #[test] + fn test_fixed_scalar_probes_read_complete_values() { + let expected_byte = 0xAB; + let expected_int16 = -0x1234i16; + let expected_uint16 = 0x1234u16; + let expected_uint24 = 0x00A1_B2C3u32; + let expected_int32 = -0x0123_4567i32; + let expected_uint32 = 0x89AB_CDEFu32; + let expected_uint40 = 0xAB_CDEF_0123u64; + let expected_int64 = -0x0102_0304_0506_0708i64; + let expected_float32 = 1.5f32; + let expected_float64 = -2.25f64; + + let mut bytes = Vec::new(); + bytes.push(expected_byte); + bytes.extend_from_slice(&expected_int16.to_le_bytes()); + bytes.extend_from_slice(&expected_uint16.to_le_bytes()); + bytes.extend_from_slice(&expected_uint24.to_le_bytes()[..3]); + bytes.extend_from_slice(&expected_int32.to_le_bytes()); + bytes.extend_from_slice(&expected_uint32.to_le_bytes()); + bytes.extend_from_slice(&expected_uint40.to_le_bytes()[..5]); + bytes.extend_from_slice(&expected_int64.to_le_bytes()); + bytes.extend_from_slice(&expected_float32.to_le_bytes()); + bytes.extend_from_slice(&expected_float64.to_le_bytes()); + + let mut buf = TdsReadBuffer::new(4096); + buf.working_buffer[..bytes.len()].copy_from_slice(&bytes); + buf.reset_to_length(bytes.len()); + + assert_eq!(buf.try_read_byte(), Some(expected_byte)); + assert_eq!(buf.try_read_int16(), Some(expected_int16)); + assert_eq!(buf.try_read_uint16(), Some(expected_uint16)); + assert_eq!(buf.try_read_uint24(), Some(expected_uint24)); + assert_eq!(buf.try_read_int32(), Some(expected_int32)); + assert_eq!(buf.try_read_uint32(), Some(expected_uint32)); + assert_eq!(buf.try_read_uint40(), Some(expected_uint40)); + assert_eq!(buf.try_read_int64(), Some(expected_int64)); + assert_eq!(buf.try_read_float32(), Some(expected_float32)); + assert_eq!(buf.try_read_float64(), Some(expected_float64)); + assert_eq!(buf.get_remaining_byte_count(), 0); + } + + #[test] + fn test_fixed_scalar_probe_misses_do_not_consume() { + let mut buf = TdsReadBuffer::new(4096); + + macro_rules! assert_miss_does_not_consume { + ($partial_len:expr, $method:ident) => {{ + buf.working_buffer[..$partial_len].fill(0xA5); + buf.reset_to_length($partial_len); + assert_eq!(buf.$method(), None); + assert_eq!(buf.buffer_position, 0); + assert_eq!(buf.get_remaining_byte_count(), $partial_len); + }}; + } + + assert_miss_does_not_consume!(0, try_read_byte); + assert_miss_does_not_consume!(1, try_read_int16); + assert_miss_does_not_consume!(1, try_read_uint16); + assert_miss_does_not_consume!(2, try_read_uint24); + assert_miss_does_not_consume!(3, try_read_int32); + assert_miss_does_not_consume!(3, try_read_uint32); + assert_miss_does_not_consume!(4, try_read_uint40); + assert_miss_does_not_consume!(7, try_read_int64); + assert_miss_does_not_consume!(3, try_read_float32); + assert_miss_does_not_consume!(7, try_read_float64); + } + #[test] fn test_get_remaining_byte_count() { let mut buf = TdsReadBuffer::new(4096); diff --git a/mssql-tds/src/connection/transport/network_transport.rs b/mssql-tds/src/connection/transport/network_transport.rs index f8ae9770..236b48bb 100644 --- a/mssql-tds/src/connection/transport/network_transport.rs +++ b/mssql-tds/src/connection/transport/network_transport.rs @@ -1115,13 +1115,63 @@ impl TdsPacketReader for NetworkTransport { self.tds_read_buffer.reset_to_length(0); } + #[inline(always)] + fn try_read_byte(&mut self) -> Option { + self.tds_read_buffer.try_read_byte() + } + + #[inline(always)] + fn try_read_int16(&mut self) -> Option { + self.tds_read_buffer.try_read_int16() + } + + #[inline(always)] + fn try_read_uint16(&mut self) -> Option { + self.tds_read_buffer.try_read_uint16() + } + + #[inline(always)] + fn try_read_uint24(&mut self) -> Option { + self.tds_read_buffer.try_read_uint24() + } + + #[inline(always)] + fn try_read_int32(&mut self) -> Option { + self.tds_read_buffer.try_read_int32() + } + + #[inline(always)] + fn try_read_uint32(&mut self) -> Option { + self.tds_read_buffer.try_read_uint32() + } + + #[inline(always)] + fn try_read_uint40(&mut self) -> Option { + self.tds_read_buffer.try_read_uint40() + } + + #[inline(always)] + fn try_read_int64(&mut self) -> Option { + self.tds_read_buffer.try_read_int64() + } + + #[inline(always)] + fn try_read_float32(&mut self) -> Option { + self.tds_read_buffer.try_read_float32() + } + + #[inline(always)] + fn try_read_float64(&mut self) -> Option { + self.tds_read_buffer.try_read_float64() + } + async fn read_byte(&mut self) -> TdsResult { - while !self.tds_read_buffer.do_we_have_enough_data(1) { + loop { + if let Some(value) = self.try_read_byte() { + return Ok(value); + } self.read_tds_packet().await?; } - let result: u8 = self.tds_read_buffer.working_buffer[self.tds_read_buffer.buffer_position]; - self.tds_read_buffer.consume_bytes(1); - Ok(result) } async fn read_int16_big_endian(&mut self) -> TdsResult { @@ -1142,80 +1192,79 @@ impl TdsPacketReader for NetworkTransport { } async fn read_uint40(&mut self) -> TdsResult { - while !self.tds_read_buffer.do_we_have_enough_data(5) { + loop { + if let Some(value) = self.try_read_uint40() { + return Ok(value); + } self.read_tds_packet().await?; } - - let result = LittleEndian::read_uint(self.tds_read_buffer.get_slice(), 5); - self.tds_read_buffer.consume_bytes(5); - Ok(result) } async fn read_float32(&mut self) -> TdsResult { - while !self.tds_read_buffer.do_we_have_enough_data(4) { + loop { + if let Some(value) = self.try_read_float32() { + return Ok(value); + } self.read_tds_packet().await?; } - let result = LittleEndian::read_f32(self.tds_read_buffer.get_slice()); - self.tds_read_buffer.consume_bytes(4); - Ok(result) } async fn read_float64(&mut self) -> TdsResult { - while !self.tds_read_buffer.do_we_have_enough_data(8) { + loop { + if let Some(value) = self.try_read_float64() { + return Ok(value); + } self.read_tds_packet().await?; } - let result = LittleEndian::read_f64(self.tds_read_buffer.get_slice()); - self.tds_read_buffer.consume_bytes(8); - Ok(result) } async fn read_int16(&mut self) -> TdsResult { - while !self.tds_read_buffer.do_we_have_enough_data(2) { + loop { + if let Some(value) = self.try_read_int16() { + return Ok(value); + } self.read_tds_packet().await?; } - let result = LittleEndian::read_i16(self.tds_read_buffer.get_slice()); - self.tds_read_buffer.consume_bytes(2); - Ok(result) } async fn read_uint16(&mut self) -> TdsResult { - while !self.tds_read_buffer.do_we_have_enough_data(2) { + loop { + if let Some(value) = self.try_read_uint16() { + return Ok(value); + } self.read_tds_packet().await?; } - let result = LittleEndian::read_u16(self.tds_read_buffer.get_slice()); - self.tds_read_buffer.consume_bytes(2); - Ok(result) } async fn read_uint24(&mut self) -> TdsResult { - while !self.tds_read_buffer.do_we_have_enough_data(3) { + loop { + if let Some(value) = self.try_read_uint24() { + return Ok(value); + } self.read_tds_packet().await?; } - let result = LittleEndian::read_u24(self.tds_read_buffer.get_slice()); - self.tds_read_buffer.consume_bytes(3); - Ok(result) } async fn read_int32(&mut self) -> TdsResult { - while !self.tds_read_buffer.do_we_have_enough_data(4) { + loop { + if let Some(value) = self.try_read_int32() { + return Ok(value); + } self.read_tds_packet().await?; } - let result = LittleEndian::read_i32(self.tds_read_buffer.get_slice()); - self.tds_read_buffer.consume_bytes(4); - Ok(result) } async fn read_uint32(&mut self) -> TdsResult { - while !self.tds_read_buffer.do_we_have_enough_data(4) { + loop { + if let Some(value) = self.try_read_uint32() { + return Ok(value); + } self.read_tds_packet().await?; } - let result = LittleEndian::read_u32(self.tds_read_buffer.get_slice()); - self.tds_read_buffer.consume_bytes(4); - Ok(result) } async fn read_int64(&mut self) -> TdsResult { - while !self.tds_read_buffer.do_we_have_enough_data(8) { + loop { + if let Some(value) = self.try_read_int64() { + return Ok(value); + } self.read_tds_packet().await?; } - let result = LittleEndian::read_i64(self.tds_read_buffer.get_slice()); - self.tds_read_buffer.consume_bytes(8); - Ok(result) } async fn read_uint64(&mut self) -> TdsResult { while !self.tds_read_buffer.do_we_have_enough_data(8) { @@ -1392,11 +1441,6 @@ impl TdsTokenStreamReader for NetworkTransport { plan: ColumnPolicy, writer: &mut (dyn RowWriter + Send), ) -> TdsResult { - // `self` is the packet reader, so the scratch slot has to be moved out - // for the duration of the read and put back afterwards. The restore - // below must stay unconditional, and no `?` may be introduced between - // these two points: an early return would drop the cached bitmap and - // silently cost an allocation on every subsequent row. let mut nbc_bitmap_scratch = self.nbc_bitmap_scratch.take(); let result = await_within_request_timeout!( remaining_request_timeout, @@ -2718,6 +2762,94 @@ pub(crate) mod tests { assert_eq!(reader.read_uint32().await.unwrap(), 0x4433_2211); } + #[tokio::test] + async fn test_sync_scalar_probe_fallback_across_packet_boundaries() { + let expected_uint16 = 0x1234u16; + let expected_int16 = -0x1234i16; + let expected_uint24 = 0x00A1_B2C3u32; + let expected_int32 = -0x0123_4567i32; + let expected_uint32 = 0x89AB_CDEFu32; + let expected_uint40 = 0xAB_CDEF_0123u64; + let expected_int64 = -0x0102_0304_0506_0708i64; + let expected_float32 = 1.5f32; + let expected_float64 = -2.25f64; + + let uint16 = expected_uint16.to_le_bytes(); + let int16 = expected_int16.to_le_bytes(); + let uint24 = expected_uint24.to_le_bytes(); + let int32 = expected_int32.to_le_bytes(); + let uint32 = expected_uint32.to_le_bytes(); + let uint40 = expected_uint40.to_le_bytes(); + let int64 = expected_int64.to_le_bytes(); + let float32 = expected_float32.to_le_bytes(); + let float64 = expected_float64.to_le_bytes(); + + let payloads = [ + vec![0xAB, uint16[0]], + vec![uint16[1], int16[0]], + vec![int16[1], uint24[0], uint24[1]], + vec![uint24[2], int32[0], int32[1], int32[2]], + vec![int32[3], uint32[0], uint32[1], uint32[2]], + vec![uint32[3], uint40[0], uint40[1], uint40[2], uint40[3]], + vec![ + uint40[4], int64[0], int64[1], int64[2], int64[3], int64[4], int64[5], int64[6], + ], + vec![int64[7], float32[0], float32[1], float32[2]], + vec![ + float32[3], float64[0], float64[1], float64[2], float64[3], float64[4], float64[5], + float64[6], + ], + vec![float64[7]], + ]; + + let mut stream = Vec::new(); + for payload in payloads { + let mut packet = TestPacketBuilder::new(PacketType::TabularResult); + stream.extend_from_slice(&packet.append_bytes(&payload).build()); + } + + let mut reader = create_network_transport_with_data(&stream); + + assert_eq!(reader.try_read_byte(), None); + assert_eq!(reader.read_byte().await.unwrap(), 0xAB); + + assert_eq!(reader.try_read_uint16(), None); + assert_eq!(reader.tds_read_buffer.get_remaining_byte_count(), 1); + assert_eq!(reader.read_uint16().await.unwrap(), expected_uint16); + + assert_eq!(reader.try_read_int16(), None); + assert_eq!(reader.tds_read_buffer.get_remaining_byte_count(), 1); + assert_eq!(reader.read_int16().await.unwrap(), expected_int16); + + assert_eq!(reader.try_read_uint24(), None); + assert_eq!(reader.tds_read_buffer.get_remaining_byte_count(), 2); + assert_eq!(reader.read_uint24().await.unwrap(), expected_uint24); + + assert_eq!(reader.try_read_int32(), None); + assert_eq!(reader.tds_read_buffer.get_remaining_byte_count(), 3); + assert_eq!(reader.read_int32().await.unwrap(), expected_int32); + + assert_eq!(reader.try_read_uint32(), None); + assert_eq!(reader.tds_read_buffer.get_remaining_byte_count(), 3); + assert_eq!(reader.read_uint32().await.unwrap(), expected_uint32); + + assert_eq!(reader.try_read_uint40(), None); + assert_eq!(reader.tds_read_buffer.get_remaining_byte_count(), 4); + assert_eq!(reader.read_uint40().await.unwrap(), expected_uint40); + + assert_eq!(reader.try_read_int64(), None); + assert_eq!(reader.tds_read_buffer.get_remaining_byte_count(), 7); + assert_eq!(reader.read_int64().await.unwrap(), expected_int64); + + assert_eq!(reader.try_read_float32(), None); + assert_eq!(reader.tds_read_buffer.get_remaining_byte_count(), 3); + assert_eq!(reader.read_float32().await.unwrap(), expected_float32); + + assert_eq!(reader.try_read_float64(), None); + assert_eq!(reader.tds_read_buffer.get_remaining_byte_count(), 7); + assert_eq!(reader.read_float64().await.unwrap(), expected_float64); + } + /// A payload-free non-EOM packet is malformed: it neither carries payload /// nor terminates a message. #[tokio::test] diff --git a/mssql-tds/src/datatypes/decoder.rs b/mssql-tds/src/datatypes/decoder.rs index 94fd6d2e..3a37bc41 100644 --- a/mssql-tds/src/datatypes/decoder.rs +++ b/mssql-tds/src/datatypes/decoder.rs @@ -30,6 +30,15 @@ use crate::{query::metadata::ColumnMetadata, token::tokens::SqlCollation}; use super::row_writer::{RowWriter, write_column_value}; +macro_rules! read_sync_first { + ($reader:expr, $try_method:ident, $read_method:ident) => { + match ($reader).$try_method() { + Some(value) => value, + None => ($reader).$read_method().await?, + } + }; +} + /// Reads an encrypted column's cipher bytes from the wire and turns them back /// into a plaintext [`ColumnValues`]. /// @@ -209,7 +218,7 @@ impl PlpChunkStreamReader { where T: TdsPacketReader + Send + Sync, { - let raw_len_i64 = reader.read_int64().await?; + let raw_len_i64 = read_sync_first!(reader, try_read_int64, read_int64); let raw_len = raw_len_i64 as u64; let raw_len_usize = raw_len as usize; @@ -264,7 +273,7 @@ impl PlpChunkStreamReader { return Ok(true); } - let chunk_len = reader.read_uint32().await? as usize; + let chunk_len = read_sync_first!(reader, try_read_uint32, read_uint32) as usize; if chunk_len == 0 { self.reached_end = true; if let PlpChunkReadLength::Known(known_len) = self.length @@ -537,10 +546,10 @@ impl GenericDecoder { where T: TdsPacketReader + Send + Sync, { - let length = reader.read_uint32().await?; - let variant_base_type = reader.read_byte().await?; + let length = read_sync_first!(reader, try_read_uint32, read_uint32); + let variant_base_type = read_sync_first!(reader, try_read_byte, read_byte); let tds_type = TdsDataType::try_from(variant_base_type)?; - let variant_prop_bytes = reader.read_byte().await?; + let variant_prop_bytes = read_sync_first!(reader, try_read_byte, read_byte); let bytes_for_type_and_properties_byte = 2; // Use checked arithmetic to prevent integer underflow @@ -632,7 +641,7 @@ impl GenericDecoder { where T: TdsPacketReader + Send + Sync, { - let scale = reader.read_byte().await?; + let scale = read_sync_first!(reader, try_read_byte, read_byte); Ok(match tds_type { TdsDataType::TimeN => { let time_nanos = self.read_time(reader, data_length as u8, scale).await?; @@ -663,7 +672,7 @@ impl GenericDecoder { T: TdsPacketReader + Send + Sync, { // Decimal/numeric data type has 1 byte length. - let length = reader.read_byte().await?; + let length = read_sync_first!(reader, try_read_byte, read_byte); let TypeInfoVariant::VarLenPrecisionScale(_, _, precision, scale) = metadata.type_info.type_info_variant else { @@ -688,7 +697,7 @@ impl GenericDecoder { if length == 0 { return Ok(None); } - let sign = reader.read_byte().await?; + let sign = read_sync_first!(reader, try_read_byte, read_byte); let is_positive = sign == 1; // Round up: a declared length that does not cover whole 32-bit words @@ -729,8 +738,8 @@ impl GenericDecoder { where T: TdsPacketReader + Send + Sync, { - let days = reader.read_int32().await?; - let ticks = reader.read_uint32().await?; + let days = read_sync_first!(reader, try_read_int32, read_int32); + let ticks = read_sync_first!(reader, try_read_uint32, read_uint32); Ok(SqlDateTime { days, time: ticks }) } @@ -739,8 +748,8 @@ impl GenericDecoder { where T: TdsPacketReader + Send + Sync, { - let days = reader.read_uint16().await?; - let minutes = reader.read_uint16().await?; + let days = read_sync_first!(reader, try_read_uint16, read_uint16); + let minutes = read_sync_first!(reader, try_read_uint16, read_uint16); Ok(SqlSmallDateTime { days, time: minutes, @@ -751,7 +760,7 @@ impl GenericDecoder { where T: TdsPacketReader + Send + Sync, { - let days = reader.read_uint24().await?; + let days = read_sync_first!(reader, try_read_uint24, read_uint24); Ok(SqlDate::unchecked_create(days)) } @@ -760,9 +769,9 @@ impl GenericDecoder { T: TdsPacketReader + Send + Sync, { let scaled_value = match byte_len { - 3 => reader.read_uint24().await? as u64, - 4 => reader.read_uint32().await? as u64, - _ => reader.read_uint40().await?, + 3 => read_sync_first!(reader, try_read_uint24, read_uint24) as u64, + 4 => read_sync_first!(reader, try_read_uint32, read_uint32) as u64, + _ => read_sync_first!(reader, try_read_uint40, read_uint40), }; // The value from SQL Server is in scaled units based on the scale: @@ -841,7 +850,7 @@ impl GenericDecoder { ))); } }; - let offset = reader.read_int16().await?; + let offset = read_sync_first!(reader, try_read_int16, read_int16); let datetime_offset = SqlDateTimeOffset { datetime2, offset }; Ok(ColumnValues::DateTimeOffset(datetime_offset)) } @@ -851,10 +860,10 @@ impl GenericDecoder { T: TdsPacketReader + Send + Sync, { let value: ColumnValues = match byte_len { - 1 => ColumnValues::TinyInt(reader.read_byte().await?), // Some(reader.read_byte().await? as i64), - 2 => ColumnValues::SmallInt(reader.read_int16().await?), // Some(reader.read_int16().await? as i64), - 4 => ColumnValues::Int(reader.read_int32().await?), - 8 => ColumnValues::BigInt(reader.read_int64().await?), + 1 => ColumnValues::TinyInt(read_sync_first!(reader, try_read_byte, read_byte)), + 2 => ColumnValues::SmallInt(read_sync_first!(reader, try_read_int16, read_int16)), + 4 => ColumnValues::Int(read_sync_first!(reader, try_read_int32, read_int32)), + 8 => ColumnValues::BigInt(read_sync_first!(reader, try_read_int64, read_int64)), 0 => ColumnValues::Null, _ => { return Err(crate::error::Error::from(Error::new( @@ -870,7 +879,7 @@ impl GenericDecoder { where T: TdsPacketReader + Send + Sync, { - let small_money_val = reader.read_int32().await?; + let small_money_val = read_sync_first!(reader, try_read_int32, read_int32); Ok(small_money_val.into()) } @@ -880,8 +889,8 @@ impl GenericDecoder { where T: TdsPacketReader + Send + Sync, { - let msb = reader.read_int32().await?; - let lsb = reader.read_int32().await?; + let msb = read_sync_first!(reader, try_read_int32, read_int32); + let lsb = read_sync_first!(reader, try_read_int32, read_int32); Ok(SqlMoney { lsb_part: lsb, msb_part: msb, @@ -937,7 +946,7 @@ impl GenericDecoder { }; // Read length prefix (USHORTLEN format) - let length_prefix_value = reader.read_uint16().await? as usize; + let length_prefix_value = read_sync_first!(reader, try_read_uint16, read_uint16) as usize; // Handle NULL (length = 0xFFFF) if length_prefix_value == 0xFFFF { @@ -961,13 +970,13 @@ impl GenericDecoder { } // Read 8-byte header - let layout_format_byte = reader.read_byte().await?; - let layout_version_byte = reader.read_byte().await?; - let dimension_count = reader.read_uint16().await?; - let base_type_byte = reader.read_byte().await?; - let _reserved1 = reader.read_byte().await?; // Reserved - let _reserved2 = reader.read_byte().await?; // Reserved - let _reserved3 = reader.read_byte().await?; // Reserved + let layout_format_byte = read_sync_first!(reader, try_read_byte, read_byte); + let layout_version_byte = read_sync_first!(reader, try_read_byte, read_byte); + let dimension_count = read_sync_first!(reader, try_read_uint16, read_uint16); + let base_type_byte = read_sync_first!(reader, try_read_byte, read_byte); + let _reserved1 = read_sync_first!(reader, try_read_byte, read_byte); // Reserved + let _reserved2 = read_sync_first!(reader, try_read_byte, read_byte); // Reserved + let _reserved3 = read_sync_first!(reader, try_read_byte, read_byte); // Reserved // Validate header using enum conversions let _layout_format = VectorLayoutFormat::try_from(layout_format_byte)?; @@ -1031,7 +1040,7 @@ impl GenericDecoder { where T: TdsPacketReader + Send + Sync, { - let long_len_i64 = reader.read_int64().await?; + let long_len_i64 = read_sync_first!(reader, try_read_int64, read_int64); let long_len = long_len_i64 as u64; // If the length is SQL_PLP_NULL, it means the value is NULL. @@ -1055,7 +1064,7 @@ impl GenericDecoder { 0 }; let mut plp_buffer = vec![0u8; vector_capacity]; - let mut chunk_len = reader.read_uint32().await? as usize; + let mut chunk_len = read_sync_first!(reader, try_read_uint32, read_uint32) as usize; let mut offset: usize = 0; let mut chunk_count = 0u32; @@ -1116,7 +1125,7 @@ impl GenericDecoder { .read_bytes(&mut plp_buffer[offset..offset + chunk_len]) .await?; offset += chunk_size_read; - chunk_len = reader.read_uint32().await? as usize; + chunk_len = read_sync_first!(reader, try_read_uint32, read_uint32) as usize; } Ok(Some(plp_buffer)) } @@ -1140,24 +1149,30 @@ impl GenericDecoder { match metadata.data_type { // === Fixed-length integer types === TdsDataType::Int1 => { - writer.write_u8(col, reader.read_byte().await?); + writer.write_u8(col, read_sync_first!(reader, try_read_byte, read_byte)); } TdsDataType::Int2 => { - writer.write_i16(col, reader.read_int16().await?); + writer.write_i16(col, read_sync_first!(reader, try_read_int16, read_int16)); } TdsDataType::Int4 => { - writer.write_i32(col, reader.read_int32().await?); + writer.write_i32(col, read_sync_first!(reader, try_read_int32, read_int32)); } TdsDataType::Int8 => { - writer.write_i64(col, reader.read_int64().await?); + writer.write_i64(col, read_sync_first!(reader, try_read_int64, read_int64)); } TdsDataType::IntN => { - let byte_len = reader.read_byte().await?; + let byte_len = read_sync_first!(reader, try_read_byte, read_byte); match byte_len { - 1 => writer.write_u8(col, reader.read_byte().await?), - 2 => writer.write_i16(col, reader.read_int16().await?), - 4 => writer.write_i32(col, reader.read_int32().await?), - 8 => writer.write_i64(col, reader.read_int64().await?), + 1 => writer.write_u8(col, read_sync_first!(reader, try_read_byte, read_byte)), + 2 => { + writer.write_i16(col, read_sync_first!(reader, try_read_int16, read_int16)) + } + 4 => { + writer.write_i32(col, read_sync_first!(reader, try_read_int32, read_int32)) + } + 8 => { + writer.write_i64(col, read_sync_first!(reader, try_read_int64, read_int64)) + } 0 => writer.write_null(col), _ => { return Err(crate::error::Error::from(Error::new( @@ -1170,28 +1185,40 @@ impl GenericDecoder { // === Fixed-length float types === TdsDataType::Flt4 => { - writer.write_f32(col, reader.read_float32().await?); + writer.write_f32( + col, + read_sync_first!(reader, try_read_float32, read_float32), + ); } TdsDataType::Flt8 => { - writer.write_f64(col, reader.read_float64().await?); + writer.write_f64( + col, + read_sync_first!(reader, try_read_float64, read_float64), + ); } TdsDataType::FltN => { - let length = reader.read_byte().await?; + let length = read_sync_first!(reader, try_read_byte, read_byte); match length { 0 => writer.write_null(col), - 4 => writer.write_f32(col, reader.read_float32().await?), - _ => writer.write_f64(col, reader.read_float64().await?), + 4 => writer.write_f32( + col, + read_sync_first!(reader, try_read_float32, read_float32), + ), + _ => writer.write_f64( + col, + read_sync_first!(reader, try_read_float64, read_float64), + ), } } // === Bit types === TdsDataType::Bit => { - writer.write_bool(col, reader.read_byte().await? == 1); + writer.write_bool(col, read_sync_first!(reader, try_read_byte, read_byte) == 1); } TdsDataType::BitN => { - let byte_len = reader.read_byte().await?; + let byte_len = read_sync_first!(reader, try_read_byte, read_byte); if byte_len > 0 { - writer.write_bool(col, reader.read_byte().await? == 1); + writer.write_bool(col, read_sync_first!(reader, try_read_byte, read_byte) == 1); } else { writer.write_null(col); } @@ -1205,7 +1232,7 @@ impl GenericDecoder { writer.write_money(col, self.read_money8(reader).await?); } TdsDataType::MoneyN => { - let byte_len = reader.read_byte().await?; + let byte_len = read_sync_first!(reader, try_read_byte, read_byte); match byte_len { 4 => writer.write_smallmoney(col, self.read_money4(reader).await?), 8 => writer.write_money(col, self.read_money8(reader).await?), @@ -1244,7 +1271,7 @@ impl GenericDecoder { // === Binary types === TdsDataType::BigBinary => { - let length = reader.read_uint16().await?; + let length = read_sync_first!(reader, try_read_uint16, read_uint16); // 0xFFFF is the USHORTLEN NULL marker (CHARBIN_NULL). if length == 0xFFFF { writer.write_null(col); @@ -1266,7 +1293,7 @@ impl GenericDecoder { None => writer.write_null(col), } } else { - let length = reader.read_uint16().await?; + let length = read_sync_first!(reader, try_read_uint16, read_uint16); // 0xFFFF is the USHORTLEN NULL marker (CHARBIN_NULL). if length == 0xFFFF { writer.write_null(col); @@ -1288,8 +1315,8 @@ impl GenericDecoder { writer.write_datetime(col, self.read_datetime(reader).await?); } TdsDataType::DateTim4 => { - let daypart = reader.read_uint16().await?; - let timepart = reader.read_uint16().await?; + let daypart = read_sync_first!(reader, try_read_uint16, read_uint16); + let timepart = read_sync_first!(reader, try_read_uint16, read_uint16); writer.write_smalldatetime( col, SqlSmallDateTime { @@ -1299,7 +1326,7 @@ impl GenericDecoder { ); } TdsDataType::DateTimeN => { - let length = reader.read_byte().await?; + let length = read_sync_first!(reader, try_read_byte, read_byte); match length { 0 => writer.write_null(col), 4 => writer.write_smalldatetime(col, self.read_small_datetime(reader).await?), @@ -1307,7 +1334,7 @@ impl GenericDecoder { } } TdsDataType::DateN => { - let length = reader.read_byte().await?; + let length = read_sync_first!(reader, try_read_byte, read_byte); if length == 0 { writer.write_null(col); } else { @@ -1315,7 +1342,7 @@ impl GenericDecoder { } } TdsDataType::TimeN => { - let length = reader.read_byte().await?; + let length = read_sync_first!(reader, try_read_byte, read_byte); if length == 0 { writer.write_null(col); } else { @@ -1335,7 +1362,7 @@ impl GenericDecoder { } } TdsDataType::DateTime2N => { - let length = reader.read_byte().await?; + let length = read_sync_first!(reader, try_read_byte, read_byte); if length == 0 { writer.write_null(col); } else { @@ -1356,7 +1383,7 @@ impl GenericDecoder { } } TdsDataType::DateTimeOffsetN => { - let length = reader.read_byte().await?; + let length = read_sync_first!(reader, try_read_byte, read_byte); if length == 0 { writer.write_null(col); } else { @@ -1379,7 +1406,7 @@ impl GenericDecoder { // === GUID === TdsDataType::Guid => { - let length = reader.read_byte().await?; + let length = read_sync_first!(reader, try_read_byte, read_byte); if length == 0 { writer.write_null(col); } else { @@ -1414,33 +1441,33 @@ impl SqlTypeDecode for GenericDecoder { { let result = match metadata.data_type { TdsDataType::Int1 => { - let value = reader.read_byte().await?; + let value = read_sync_first!(reader, try_read_byte, read_byte); ColumnValues::from(value) } TdsDataType::Int2 => { - let value = reader.read_int16().await?; + let value = read_sync_first!(reader, try_read_int16, read_int16); ColumnValues::SmallInt(value) } TdsDataType::Int4 => { - let value = reader.read_int32().await?; + let value = read_sync_first!(reader, try_read_int32, read_int32); ColumnValues::from(value) } TdsDataType::Int8 => { - let value = reader.read_int64().await?; + let value = read_sync_first!(reader, try_read_int64, read_int64); ColumnValues::BigInt(value) } TdsDataType::Flt4 => { - let value = reader.read_float32().await?; + let value = read_sync_first!(reader, try_read_float32, read_float32); ColumnValues::Real(value) } TdsDataType::Flt8 => { - let value = reader.read_float64().await?; + let value = read_sync_first!(reader, try_read_float64, read_float64); ColumnValues::Float(value) } TdsDataType::Money4 => ColumnValues::SmallMoney(self.read_money4(reader).await?), TdsDataType::Money => ColumnValues::Money(self.read_money8(reader).await?), TdsDataType::MoneyN => { - let byte_len = reader.read_byte().await?; + let byte_len = read_sync_first!(reader, try_read_byte, read_byte); match byte_len { 4 => ColumnValues::SmallMoney(self.read_money4(reader).await?), 8 => ColumnValues::Money(self.read_money8(reader).await?), @@ -1467,7 +1494,7 @@ impl SqlTypeDecode for GenericDecoder { } } TdsDataType::Bit => { - let value = reader.read_byte().await?; + let value = read_sync_first!(reader, try_read_byte, read_byte); ColumnValues::Bit(value == 1) } TdsDataType::NChar @@ -1483,11 +1510,11 @@ impl SqlTypeDecode for GenericDecoder { ColumnValues::DateTime(value) } TdsDataType::IntN => { - let byte_len = reader.read_byte().await?; + let byte_len = read_sync_first!(reader, try_read_byte, read_byte); self.read_intn(reader, byte_len).await? } TdsDataType::BigBinary => { - let length = reader.read_uint16().await?; + let length = read_sync_first!(reader, try_read_uint16, read_uint16); // 0xFFFF is the USHORTLEN NULL marker (CHARBIN_NULL). if length == 0xFFFF { ColumnValues::Null @@ -1510,7 +1537,7 @@ impl SqlTypeDecode for GenericDecoder { None => ColumnValues::Null, } } else { - let length = reader.read_uint16().await?; + let length = read_sync_first!(reader, try_read_uint16, read_uint16); // 0xFFFF is the USHORTLEN NULL marker (CHARBIN_NULL). if length == 0xFFFF { ColumnValues::Null @@ -1552,34 +1579,34 @@ impl SqlTypeDecode for GenericDecoder { } TdsDataType::Vector => self.decode_vector(reader, metadata).await?, TdsDataType::BitN => { - let byte_len = reader.read_byte().await?; + let byte_len = read_sync_first!(reader, try_read_byte, read_byte); if byte_len > 0 { - let value = reader.read_byte().await?; + let value = read_sync_first!(reader, try_read_byte, read_byte); ColumnValues::Bit(value == 1) } else { ColumnValues::Null } } TdsDataType::Guid => { - let length = reader.read_byte().await?; + let length = read_sync_first!(reader, try_read_byte, read_byte); Self::read_guid(reader, length).await? } TdsDataType::FltN => { // This is variable length float, hence the length needs to be read first - let length = reader.read_byte().await?; + let length = read_sync_first!(reader, try_read_byte, read_byte); if length == 0 { return Ok(ColumnValues::Null); } if length == 4 { - let value = reader.read_float32().await?; + let value = read_sync_first!(reader, try_read_float32, read_float32); ColumnValues::Real(value) } else { - let value = reader.read_float64().await?; + let value = read_sync_first!(reader, try_read_float64, read_float64); ColumnValues::Float(value) } } TdsDataType::DateTimeN => { - let length = reader.read_byte().await?; + let length = read_sync_first!(reader, try_read_byte, read_byte); // If length is 0, then it is NULL if length == 0 { return Ok(ColumnValues::Null); @@ -1593,11 +1620,11 @@ impl SqlTypeDecode for GenericDecoder { } } TdsDataType::DateN => { - let length = reader.read_byte().await?; + let length = read_sync_first!(reader, try_read_byte, read_byte); return Self::read_daten(reader, length).await; } TdsDataType::TimeN => { - let length = reader.read_byte().await?; + let length = read_sync_first!(reader, try_read_byte, read_byte); match length { 0 => return Ok(ColumnValues::Null), _ => { @@ -1617,7 +1644,7 @@ impl SqlTypeDecode for GenericDecoder { } } TdsDataType::DateTime2N => { - let length = reader.read_byte().await?; + let length = read_sync_first!(reader, try_read_byte, read_byte); match length { 0 => Ok(ColumnValues::Null), _ => { @@ -1635,7 +1662,7 @@ impl SqlTypeDecode for GenericDecoder { } }?, TdsDataType::DateTimeOffsetN => { - let length = reader.read_byte().await?; + let length = read_sync_first!(reader, try_read_byte, read_byte); match length { 0 => Ok(ColumnValues::Null), _ => { @@ -1653,13 +1680,13 @@ impl SqlTypeDecode for GenericDecoder { } }?, TdsDataType::Image => { - let text_ptr_len = reader.read_byte().await? as usize; + let text_ptr_len = read_sync_first!(reader, try_read_byte, read_byte) as usize; let length = if text_ptr_len > 0 { const TIMESTAMP_BYTE_COUNT: usize = 8; reader.skip_bytes(text_ptr_len).await?; reader.skip_bytes(TIMESTAMP_BYTE_COUNT).await?; - reader.read_uint32().await? as usize + read_sync_first!(reader, try_read_uint32, read_uint32) as usize } else { 0 }; @@ -1691,8 +1718,8 @@ impl SqlTypeDecode for GenericDecoder { } TdsDataType::SsVariant => self.read_sql_variant(reader).await?, TdsDataType::DateTim4 => { - let daypart = reader.read_uint16().await?; - let timepart = reader.read_uint16().await?; + let daypart = read_sync_first!(reader, try_read_uint16, read_uint16); + let timepart = read_sync_first!(reader, try_read_uint16, read_uint16); ColumnValues::SmallDateTime(SqlSmallDateTime { days: daypart, time: timepart, @@ -1757,7 +1784,7 @@ impl StringDecoder { None => writer.write_null(col), } } else if Self::is_long_len_type(metadata.data_type) { - let text_ptr_len = reader.read_byte().await? as usize; + let text_ptr_len = read_sync_first!(reader, try_read_byte, read_byte) as usize; if text_ptr_len == 0 { writer.write_null(col); @@ -1767,7 +1794,7 @@ impl StringDecoder { const TIMESTAMP_BYTE_COUNT: usize = 8; reader.skip_bytes(text_ptr_len).await?; reader.skip_bytes(TIMESTAMP_BYTE_COUNT).await?; - let length = reader.read_uint32().await? as usize; + let length = read_sync_first!(reader, try_read_uint32, read_uint32) as usize; if length > MAX_ALLOC_SIZE { return Err(crate::error::Error::ProtocolError(format!( @@ -1784,7 +1811,7 @@ impl StringDecoder { }; writer.write_string(col, sql_string); } else { - let length = reader.read_uint16().await? as usize; + let length = read_sync_first!(reader, try_read_uint16, read_uint16) as usize; if length == 0xFFFF { writer.write_null(col); } else { @@ -1837,13 +1864,13 @@ impl SqlTypeDecode for StringDecoder { // Creates SqlString with appropriate encoding type // NULL handling works (textptr_len = 0) // LCID-based decoding implemented (see sql_string.rs) - let text_ptr_len = reader.read_byte().await? as usize; + let text_ptr_len = read_sync_first!(reader, try_read_byte, read_byte) as usize; let length = if text_ptr_len > 0 { const TIMESTAMP_BYTE_COUNT: usize = 8; reader.skip_bytes(text_ptr_len).await?; reader.skip_bytes(TIMESTAMP_BYTE_COUNT).await?; - reader.read_uint32().await? as usize + read_sync_first!(reader, try_read_uint32, read_uint32) as usize } else { // text_ptr_len == 0 means NULL value return Ok(ColumnValues::Null); @@ -1866,7 +1893,7 @@ impl SqlTypeDecode for StringDecoder { }; Ok(ColumnValues::String(sql_string)) } else { - let length = reader.read_uint16().await? as usize; + let length = read_sync_first!(reader, try_read_uint16, read_uint16) as usize; if length == 0xFFFF { Ok(ColumnValues::Null) } else { @@ -2167,7 +2194,7 @@ where Ok(match tds_type { // BIGVARBINARYTYPE, BIGBINARYTYPE TdsDataType::BigVarBinary | TdsDataType::BigBinary => { - let _max_length: u16 = reader.read_uint16().await?; + let _max_length: u16 = read_sync_first!(reader, try_read_uint16, read_uint16); if data_length as usize > MAX_ALLOC_SIZE { return Err(crate::error::Error::ProtocolError(format!( "SQL Variant binary data length {data_length} exceeds maximum allowed size of {MAX_ALLOC_SIZE} bytes" @@ -2178,8 +2205,8 @@ where ColumnValues::Bytes(buffer) } TdsDataType::NumericN | TdsDataType::DecimalN => { - let precision = reader.read_byte().await?; - let scale = reader.read_byte().await?; + let precision = read_sync_first!(reader, try_read_byte, read_byte); + let scale = read_sync_first!(reader, try_read_byte, read_byte); let decimal_parts = GenericDecoder::read_decimal_data(reader, data_length as u8, precision, scale) .await?; @@ -2222,7 +2249,7 @@ where } let mut collation_bytes = vec![0u8; 5]; reader.read_bytes(&mut collation_bytes).await?; - let _max_length = reader.read_uint16().await? as usize; + let _max_length = read_sync_first!(reader, try_read_uint16, read_uint16) as usize; let collation: SqlCollation = collation_bytes.as_slice().try_into()?; if data_length as usize > MAX_ALLOC_SIZE { return Err(crate::error::Error::ProtocolError(format!( diff --git a/mssql-tds/src/io/packet_reader.rs b/mssql-tds/src/io/packet_reader.rs index cf9c9329..617f60f3 100644 --- a/mssql-tds/src/io/packet_reader.rs +++ b/mssql-tds/src/io/packet_reader.rs @@ -10,6 +10,66 @@ pub(crate) const LENGTH_NULL: u16 = 0xffff; #[cfg(not(fuzzing))] pub(crate) trait TdsPacketReader { + /// Returns a buffered byte, or `None` without consuming data if one is unavailable. + #[inline] + fn try_read_byte(&mut self) -> Option { + None + } + + /// Returns a buffered little-endian `i16`, or `None` without consuming partial data. + #[inline] + fn try_read_int16(&mut self) -> Option { + None + } + + /// Returns a buffered little-endian `u16`, or `None` without consuming partial data. + #[inline] + fn try_read_uint16(&mut self) -> Option { + None + } + + /// Returns a buffered little-endian 24-bit integer, or `None` without consuming partial data. + #[inline] + fn try_read_uint24(&mut self) -> Option { + None + } + + /// Returns a buffered little-endian `i32`, or `None` without consuming partial data. + #[inline] + fn try_read_int32(&mut self) -> Option { + None + } + + /// Returns a buffered little-endian `u32`, or `None` without consuming partial data. + #[inline] + fn try_read_uint32(&mut self) -> Option { + None + } + + /// Returns a buffered little-endian 40-bit integer, or `None` without consuming partial data. + #[inline] + fn try_read_uint40(&mut self) -> Option { + None + } + + /// Returns a buffered little-endian `i64`, or `None` without consuming partial data. + #[inline] + fn try_read_int64(&mut self) -> Option { + None + } + + /// Returns a buffered little-endian `f32`, or `None` without consuming partial data. + #[inline] + fn try_read_float32(&mut self) -> Option { + None + } + + /// Returns a buffered little-endian `f64`, or `None` without consuming partial data. + #[inline] + fn try_read_float64(&mut self) -> Option { + None + } + fn read_byte(&mut self) -> impl Future> + Send; fn read_int16_big_endian(&mut self) -> impl Future> + Send; fn read_int32_big_endian(&mut self) -> impl Future> + Send; @@ -49,6 +109,66 @@ pub(crate) trait TdsPacketReader { /// Low-level TDS packet reading operations (public under `fuzzing` cfg). #[cfg(fuzzing)] pub trait TdsPacketReader { + /// Returns a buffered byte, or `None` without consuming data if one is unavailable. + #[inline] + fn try_read_byte(&mut self) -> Option { + None + } + + /// Returns a buffered little-endian `i16`, or `None` without consuming partial data. + #[inline] + fn try_read_int16(&mut self) -> Option { + None + } + + /// Returns a buffered little-endian `u16`, or `None` without consuming partial data. + #[inline] + fn try_read_uint16(&mut self) -> Option { + None + } + + /// Returns a buffered little-endian 24-bit integer, or `None` without consuming partial data. + #[inline] + fn try_read_uint24(&mut self) -> Option { + None + } + + /// Returns a buffered little-endian `i32`, or `None` without consuming partial data. + #[inline] + fn try_read_int32(&mut self) -> Option { + None + } + + /// Returns a buffered little-endian `u32`, or `None` without consuming partial data. + #[inline] + fn try_read_uint32(&mut self) -> Option { + None + } + + /// Returns a buffered little-endian 40-bit integer, or `None` without consuming partial data. + #[inline] + fn try_read_uint40(&mut self) -> Option { + None + } + + /// Returns a buffered little-endian `i64`, or `None` without consuming partial data. + #[inline] + fn try_read_int64(&mut self) -> Option { + None + } + + /// Returns a buffered little-endian `f32`, or `None` without consuming partial data. + #[inline] + fn try_read_float32(&mut self) -> Option { + None + } + + /// Returns a buffered little-endian `f64`, or `None` without consuming partial data. + #[inline] + fn try_read_float64(&mut self) -> Option { + None + } + fn read_byte(&mut self) -> impl Future> + Send; fn read_int16_big_endian(&mut self) -> impl Future> + Send; fn read_int32_big_endian(&mut self) -> impl Future> + Send;