From 9df51f01fc85f89d56f03fd3fd64d4a15c01ce60 Mon Sep 17 00:00:00 2001 From: Srinivasan S A Date: Wed, 12 Aug 2026 11:07:33 +0000 Subject: [PATCH] POC: reduce per-row cost in the TDS row decode path Proof of concept, not for merge. Shows five changes and what each is worth. Measured on 1,499,000 rows of 48 columns. Wall clock 7.45s -> 3.92s. CPU on the reading process 7.35s -> 2.58s. CPU inside mssql-tds 4.69s -> 0.97s. The server stopped blocking on ASYNC_NETWORK_IO, from 4.54s to none. The five changes: 1. Resolve each column's type, precision, scale and collation once per result set into a DecodeOp plan, instead of rederiving it from ColumnMetadata on every row. 2. Add a synchronous row path that decodes a whole ROW or NBCROW token from one borrow of the reader's buffer. It runs two passes, the first to measure and bounds check and the second to decode without checks. It declines when the row crosses a packet boundary, and the async path handles that case. 3. Remove the two tracing callsites that fired once per row. This was done to measure their cost, not as a proposal. 4. Offer decoded strings and binaries to the RowWriter as borrows, so a consumer that does not need to own them avoids a copy. 5. Let the RowWriter supply the destination for a PLP value whose length is declared up front, so large values are read into consumer storage directly. Known gaps are listed in POC_README.md. The largest is that NOT NULL columns arrive as fixed length TDS types, which are not in the decode plan, so a table of NOT NULL columns gets none of this. --- POC_README.md | 168 +++++ .../src/connection/metadata_retriever.rs | 1 + mssql-tds/src/connection/tds_client.rs | 25 +- .../connection/transport/network_transport.rs | 66 ++ mssql-tds/src/datatypes/decoder.rs | 595 ++++++++++++++---- mssql-tds/src/datatypes/row_writer.rs | 32 +- mssql-tds/src/datatypes/sql_string.rs | 98 +-- mssql-tds/src/io/packet_reader.rs | 147 +++++ mssql-tds/src/io/token_stream.rs | 445 ++++++++++++- .../src/token/parsers/colmetadata_parser.rs | 1 + mssql-tds/src/token/parsers/nbcrow_parser.rs | 1 + mssql-tds/src/token/parsers/row_parser.rs | 1 + mssql-tds/src/token/tokens.rs | 104 +++ 13 files changed, 1516 insertions(+), 168 deletions(-) create mode 100644 POC_README.md diff --git a/POC_README.md b/POC_README.md new file mode 100644 index 00000000..c2ef501f --- /dev/null +++ b/POC_README.md @@ -0,0 +1,168 @@ +# TDS row decoding proof of concept + +This branch shows five changes to `mssql-tds` that cut the cost of reading rows. + +It is a proof of concept. The code quality is not production ready. The intent +is to prove these changes work and to show what each one is worth. + +We found these while making a PostgreSQL foreign data wrapper faster. The +wrapper reads large result sets from SQL Server. Past a few hundred thousand +rows the client became the bottleneck, not the server. + +## What we measured + +One query. 1,499,000 rows, 48 columns each, 39 integers and 9 short strings. + +| Measure | Before | After | +| --- | ---: | ---: | +| Wall clock | 7.45 s | 3.92 s | +| CPU on the reading process | 7.35 s | 2.58 s | +| CPU inside `mssql-tds` | 4.69 s | 0.97 s | + +The server stopped waiting on us too. Before the change it spent 4.54 seconds +blocked on `ASYNC_NETWORK_IO`. After the change it spent none. + +The unit tests pass. Run them with `cargo test -p mssql-tds --lib`. + +These numbers were taken before this branch was rebased onto current `main`. +The rebase changed which code runs on the async path, so treat the split +between the groups as indicative. The buffered path itself is unchanged. + +## What changed + +### 1. Work out each column once + +The old path read the type, precision, scale and collation from the column +metadata on every row. None of that changes while a result set is open. + +Now it is worked out once and stored as a small plan. See `DecodeOp` in +`mssql-tds/src/token/tokens.rs`. + +**Why it is built on the first row.** The plan sits in a `OnceLock` on the +column metadata token and is filled the first time a row needs it. That was a +proof of concept choice. It kept the change to one field and one method, and it +needed no changes in the parser. The plan is built once per result set, so the +cost of building it disappears against the per row work it removes. + +Where the cache lives is an open question, and the lazy shape is not the +recommendation. Building it in the COLMETADATA parser would also work and would +be cleaner. That parser already walks every column and already holds the type +info, so the facts could be plain fields with no lock and no laziness. It would +also put them within reach of the async path, which is the gap described below. + +### 2. Decode a whole row from the buffer + +The old path read every column with its own `await`. Each of those reads goes +through `#[async_trait]`, which allocates a boxed future on the heap. A 48 +column row cost 96 of them, and the row was usually already in memory. + +The new path takes the whole row under one borrow. It runs two passes. The +first measures and bounds checks. The second decodes with no checks, because +the first pass already proved the offsets are good. + +If the row runs past the end of the buffer, the new path steps aside and the old +one handles it. See `try_receive_row_into_buffered` in +`mssql-tds/src/io/token_stream.rs`. + +It also steps aside in two other cases. Cursor reads pause part way through a +row, so anything other than `ColumnPolicy::DecodeAll` goes to the old path. So +does any result set with an Always Encrypted decryptor, because the new path +would hand you ciphertext. + +### 3. Take the tracing off the row path + +Two callsites fired once per row. At 1,499,000 rows that added up. + +We deleted them to measure the cost. That is not what we suggest you do. There +are better options, and they are listed in the notes we shared separately. + +### 4. Hand values over by reference + +The old path copied every string into a `Vec`, wrapped it in a `SqlString`, and +handed that to the consumer. The consumer often copied it again. + +Now the bytes are offered as a borrow. A consumer that needs to keep them can +still copy. See `write_string_ref` in `mssql-tds/src/datatypes/row_writer.rs`. + +### 5. Let the consumer own the buffer for large values + +A `varbinary(max)` or `varchar(max)` value has its length declared up front. The +consumer can allocate its own storage first and let the decoder read into it. + +That removes one full copy of the value. See `string_destination` and +`bytes_destination` in `mssql-tds/src/datatypes/row_writer.rs`. + +## What this does not do + +Read this part before you judge the numbers. + +### NOT NULL columns are not covered + +SQL Server sends a different type code for a column that cannot be null. + +A nullable `int` arrives as `IntN`. It carries a length byte before the value. A +`NOT NULL int` arrives as `Int4`. It is four bytes with no length byte at all. + +The decode plan only covers the nullable forms. `Int4`, `Flt8`, `Bit`, `Money`, +`DateTime` and the rest of the fixed length types all fall through to +`DecodeOp::Generic`. + +That matters more than it sounds. One `Generic` column sends the whole row back +to the old path. A table of `NOT NULL` columns gets none of this speedup. + +The good news is that these types are easier to handle than the ones we did. +There is no length byte to read and no length to validate. The same two pass +pattern works, with less code per type. + +### Large values are not in the fast row path + +A column holding `varchar(max)` or `varbinary(max)` is marked `Generic` as well. +Those values can span many packets, so they need the async path. + +Change 5 above still helps them. It just happens on the old path, not the new +one. + +### Some ordinary types are still missing + +`uniqueidentifier`, `datetimeoffset`, `money`, `sql_variant`, `xml`, `json` and +`vector` are all `Generic` today. Each is a small addition. + +### The async path does not use the plan yet + +Change 1 gives every column a `DecodeOp` tag. Only the buffered path reads it. + +The async path still works the type out from `ColumnMetadata` on every column, +inside `drive_row_columns`. That function also handles cursor policies, Always +Encrypted and paused rows, so wiring the plan into it is a real piece of work +rather than a small edit. + +`decode_op_into` in `mssql-tds/src/datatypes/decoder.rs` shows the shape it +would take. Nothing calls it today. + +### One workload, one platform + +We measured on Linux against SQL Server 2022, with one query shape. Your results +will differ. Treat the numbers as a direction, not a promise. + +## Where to look + +| File | What is in it | +| --- | --- | +| `mssql-tds/src/io/token_stream.rs` | The two pass row decoder | +| `mssql-tds/src/token/tokens.rs` | The per column decode plan | +| `mssql-tds/src/io/packet_reader.rs` | Non blocking reads and their helpers | +| `mssql-tds/src/datatypes/decoder.rs` | Per column decoding and the PLP paths | +| `mssql-tds/src/datatypes/row_writer.rs` | The consumer side contract | + +Start with `try_receive_row_into_buffered`. It has a long comment at the top +that explains the two passes and the rules they have to follow. + +## Why it is not production ready + +- The two passes have to agree on every field width. Nothing enforces that today + except care and tests. +- The fallback path is exercised by our workload, but not by a test that forces a + row to straddle a packet boundary. +- The guards that send cursor reads and encrypted rows to the old path have no + test of their own. +- The tracing change deletes diagnostics that other callers may want. diff --git a/mssql-tds/src/connection/metadata_retriever.rs b/mssql-tds/src/connection/metadata_retriever.rs index f9dac927..7a9c01cd 100644 --- a/mssql-tds/src/connection/metadata_retriever.rs +++ b/mssql-tds/src/connection/metadata_retriever.rs @@ -506,6 +506,7 @@ mod tests { column_count: columns.len() as u16, columns, cek_table: Vec::new(), + decode_plan: std::sync::OnceLock::new(), }, collation_names, } diff --git a/mssql-tds/src/connection/tds_client.rs b/mssql-tds/src/connection/tds_client.rs index 21bc3251..4d03a025 100644 --- a/mssql-tds/src/connection/tds_client.rs +++ b/mssql-tds/src/connection/tds_client.rs @@ -3073,7 +3073,6 @@ 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), @@ -3130,7 +3129,6 @@ impl TdsClient { match result { RowReadResult::RowWritten => { writer.end_row(); - info!("Row Received"); return Ok(true); } RowReadResult::RowPaused(_) | RowReadResult::PlpPaused(_) => { @@ -3503,6 +3501,29 @@ impl TdsClient { } } + /// Decodes the next row into `writer` without going through the + /// [`ResultSet`] trait. + /// + /// `ResultSet::next_row_into` is an `#[async_trait]` method, so every call + /// to it allocates a boxed future. This inherent form reaches the same + /// `get_next_row_into`. Both are kept, because a caller holding a + /// `dyn ResultSet` still needs the trait method. + /// + /// # Returns + /// + /// `true` when a row was written, and `false` once the result set is + /// drained. + pub async fn next_row_into_unboxed( + &mut self, + writer: &mut (dyn RowWriter + Send), + ) -> TdsResult { + if self.maybe_has_unread_rows() { + self.get_next_row_into(writer).await + } else { + Ok(false) + } + } + /// Returns a clone of all [`ReturnValue`]s collected during the current /// batch — output parameters and UDF return values. /// diff --git a/mssql-tds/src/connection/transport/network_transport.rs b/mssql-tds/src/connection/transport/network_transport.rs index e0b14746..45c08dd3 100644 --- a/mssql-tds/src/connection/transport/network_transport.rs +++ b/mssql-tds/src/connection/transport/network_transport.rs @@ -1109,6 +1109,72 @@ impl TdsPacketReader for NetworkTransport { self.tds_read_buffer.reset_to_length(0); } + fn try_read_byte(&mut self) -> Option { + if !self.tds_read_buffer.do_we_have_enough_data(1) { + return None; + } + let value = self.tds_read_buffer.get_slice()[0]; + self.tds_read_buffer.consume_bytes(1); + Some(value) + } + + fn try_read_uint16(&mut self) -> Option { + if !self.tds_read_buffer.do_we_have_enough_data(2) { + return None; + } + let value = LittleEndian::read_u16(self.tds_read_buffer.get_slice()); + self.tds_read_buffer.consume_bytes(2); + Some(value) + } + + fn try_read_int32(&mut self) -> Option { + if !self.tds_read_buffer.do_we_have_enough_data(4) { + return None; + } + let value = LittleEndian::read_i32(self.tds_read_buffer.get_slice()); + self.tds_read_buffer.consume_bytes(4); + Some(value) + } + + fn try_read_bytes(&mut self, buffer: &mut [u8]) -> Option { + if !self.tds_read_buffer.do_we_have_enough_data(buffer.len()) { + return None; + } + buffer.copy_from_slice(&self.tds_read_buffer.get_slice()[..buffer.len()]); + self.tds_read_buffer.consume_bytes(buffer.len()); + Some(buffer.len()) + } + + fn try_read_slice(&mut self, length: usize) -> Option<&[u8]> { + if !self.tds_read_buffer.do_we_have_enough_data(length) { + return None; + } + let start = self.tds_read_buffer.buffer_position; + let end = start + length; + self.tds_read_buffer.consume_bytes(length); + // consume_bytes only moves buffer_position, so these bytes stay readable + // until the next packet read overwrites working_buffer. That read needs + // &mut self, which is what bounds the returned borrow. + Some(&self.tds_read_buffer.working_buffer[start..end]) + } + + fn buffered_slice(&self) -> Option<&[u8]> { + // Always Some, and empty when the buffer is drained. Callers treat an + // empty slice as not enough data and fall back to the async path. + Some( + &self.tds_read_buffer.working_buffer + [self.tds_read_buffer.buffer_position..self.tds_read_buffer.buffer_length], + ) + } + + fn consume_buffered(&mut self, length: usize) -> bool { + if !self.tds_read_buffer.do_we_have_enough_data(length) { + return false; + } + self.tds_read_buffer.consume_bytes(length); + true + } + async fn read_byte(&mut self) -> TdsResult { while !self.tds_read_buffer.do_we_have_enough_data(1) { self.read_tds_packet().await?; diff --git a/mssql-tds/src/datatypes/decoder.rs b/mssql-tds/src/datatypes/decoder.rs index 91100aae..1f5716f8 100644 --- a/mssql-tds/src/datatypes/decoder.rs +++ b/mssql-tds/src/datatypes/decoder.rs @@ -12,6 +12,7 @@ use super::{ }; use crate::datatypes::sqldatatypes::TypeInfo; use crate::security::cell_decryptor::CellDecryptor; +use crate::token::tokens::DecodeOp; use crate::{ core::TdsResult, datatypes::{sql_json::SqlJson, sql_string::EncodingType, sqldatatypes::FixedLengthTypes}, @@ -21,7 +22,10 @@ use crate::{ ColumnValues, SqlDate, SqlDateTime, SqlDateTime2, SqlDateTimeOffset, SqlMoney, SqlSmallDateTime, SqlSmallMoney, SqlTime, SqlXml, }, - io::packet_reader::TdsPacketReader, + io::packet_reader::{ + TdsPacketReader, read_byte_buffered, read_bytes_buffered, read_int32_buffered, + read_uint16_buffered, + }, }; use crate::{query::metadata::ColumnMetadata, token::tokens::SqlCollation}; @@ -997,92 +1001,232 @@ impl GenericDecoder { let long_len_i64 = reader.read_int64().await?; let long_len = long_len_i64 as u64; - // If the length is SQL_PLP_NULL, it means the value is NULL. if long_len as usize == Self::SQL_PLP_NULL { - Ok(None) + return Ok(None); + } + if long_len as usize == Self::SQL_PLP_UNKNOWNLEN { + return Ok(Some(Self::read_unknown_plp_bytes(reader).await?)); + } + + let length = Self::validate_known_plp_length(long_len_i64)?; + let mut bytes = vec![0u8; length]; + Self::read_known_plp_bytes(reader, &mut bytes).await?; + Ok(Some(bytes)) + } + + /// Reads a PLP character value and hands it to the writer. + /// + /// A PLP value arrives as a declared total length followed by chunks. When + /// that length is known up front, the writer is offered the chance to + /// supply the destination through [`RowWriter::string_destination`], and + /// the chunks are read straight into the consumer's own storage. Nothing + /// is allocated here in that case. + /// + /// A writer that cannot hold the bytes in their wire encoding returns + /// `None`, and this falls back to an owned `Vec` wrapped in a [`SqlString`]. + /// + /// [`RowWriter::finish_string_destination`] is called whether the read + /// succeeded or not, so the writer can commit or discard a destination it + /// only partly filled. + /// + /// The unknown-length form gets no destination, because the consumer + /// cannot size its storage until the final chunk arrives. + async fn read_plp_string_into( + reader: &mut T, + col: usize, + writer: &mut W, + encoding_type: &EncodingType, + ) -> TdsResult<()> + where + T: TdsPacketReader + Send + Sync, + W: RowWriter + ?Sized, + { + let long_len_i64 = reader.read_int64().await?; + let long_len = long_len_i64 as u64; + + if long_len as usize == Self::SQL_PLP_NULL { + writer.write_null(col); + } else if long_len as usize == Self::SQL_PLP_UNKNOWNLEN { + let bytes = Self::read_unknown_plp_bytes(reader).await?; + writer.write_string(col, SqlString::new(bytes, encoding_type.clone())); } else { - // If the length is SQL_PLP_UNKNOWNLEN, it means the length is unknown and we have to - // gather all the chunks until we reach the end of the PLP data which is a zero length - // chunk. - let mut vector_capacity = if long_len as usize != Self::SQL_PLP_UNKNOWNLEN { - let capacity = long_len as usize; - // Check for overflow or excessively large values - // If long_len_i64 was negative, casting to u64 then usize can produce huge values - if long_len_i64 < 0 || capacity > MAX_PLP_SIZE { - return Err(crate::error::Error::ProtocolError(format!( - "PLP length {capacity} (raw i64: {long_len_i64}) exceeds maximum allowed size of {MAX_PLP_SIZE} bytes" - ))); - } - capacity + let length = Self::validate_known_plp_length(long_len_i64)?; + if let Some(destination) = writer.string_destination(col, length, encoding_type) { + let result = Self::read_known_plp_bytes(reader, destination).await; + writer.finish_string_destination(col, result.is_ok()); + result?; } else { - 0 - }; - let mut plp_buffer = vec![0u8; vector_capacity]; - let mut chunk_len = reader.read_uint32().await? as usize; - let mut offset: usize = 0; - let mut chunk_count = 0u32; - - while chunk_len > 0 { - chunk_count += 1; - - #[cfg(fuzzing)] - { - eprintln!( - "[ALLOC] read_plp_bytes: chunk #{chunk_count}, chunk_len={chunk_len}, total_capacity={vector_capacity}" - ); - } + let mut bytes = vec![0u8; length]; + Self::read_known_plp_bytes(reader, &mut bytes).await?; + writer.write_string(col, SqlString::new(bytes, encoding_type.clone())); + } + } + Ok(()) + } - if chunk_count > Self::MAX_PLP_CHUNKS { - return Err(crate::error::Error::ProtocolError(format!( - "Too many PLP chunks: {chunk_count} (max {})", - Self::MAX_PLP_CHUNKS - ))); - } + /// Reads a PLP binary value and hands it to the writer. + /// + /// Same shape as [`Self::read_plp_string_into`], using + /// [`RowWriter::bytes_destination`]. Binary needs no encoding agreement, so + /// a writer can accept the destination for any known length. + async fn read_plp_binary_into(reader: &mut T, col: usize, writer: &mut W) -> TdsResult<()> + where + T: TdsPacketReader + Send + Sync, + W: RowWriter + ?Sized, + { + let long_len_i64 = reader.read_int64().await?; + let long_len = long_len_i64 as u64; - // Limit individual chunk size - if chunk_len > Self::MAX_PLP_CHUNK_SIZE { - return Err(crate::error::Error::ProtocolError(format!( - "PLP chunk size {chunk_len} exceeds maximum allowed chunk size of {} bytes", - Self::MAX_PLP_CHUNK_SIZE - ))); - } + if long_len as usize == Self::SQL_PLP_NULL { + writer.write_null(col); + } else if long_len as usize == Self::SQL_PLP_UNKNOWNLEN { + writer.write_bytes(col, Self::read_unknown_plp_bytes(reader).await?); + } else { + let length = Self::validate_known_plp_length(long_len_i64)?; + if let Some(destination) = writer.bytes_destination(col, length) { + let result = Self::read_known_plp_bytes(reader, destination).await; + writer.finish_bytes_destination(col, result.is_ok()); + result?; + } else { + let mut bytes = vec![0u8; length]; + Self::read_known_plp_bytes(reader, &mut bytes).await?; + writer.write_bytes(col, bytes); + } + } + Ok(()) + } - if long_len as usize == Self::SQL_PLP_UNKNOWNLEN { - // Use checked_add to prevent capacity overflow - vector_capacity = vector_capacity.checked_add(chunk_len).ok_or_else(|| { - crate::error::Error::ProtocolError(format!( - "PLP chunk accumulation would overflow capacity: {vector_capacity} + {chunk_len}" - )) - })?; - // Validate against MAX_PLP_SIZE after accumulation - if vector_capacity > MAX_PLP_SIZE { - return Err(crate::error::Error::ProtocolError(format!( - "PLP accumulated size {vector_capacity} exceeds maximum allowed size of {MAX_PLP_SIZE} bytes (SQL Server limit: 2GB)" - ))); - } - plp_buffer.resize(vector_capacity, 0); - } else { - // For known length, validate that chunk fits within the allocated buffer - let end_offset = offset.checked_add(chunk_len).ok_or_else(|| { - crate::error::Error::ProtocolError(format!( - "PLP chunk offset would overflow: {offset} + {chunk_len}" - )) - })?; - if end_offset > plp_buffer.len() { - return Err(crate::error::Error::ProtocolError(format!( - "PLP chunk exceeds declared length: offset={offset}, chunk_len={chunk_len}, buffer_len={}, declared_len={long_len}", - plp_buffer.len() - ))); - } - } - let chunk_size_read = reader - .read_bytes(&mut plp_buffer[offset..offset + chunk_len]) - .await?; - offset += chunk_size_read; - chunk_len = reader.read_uint32().await? as usize; + /// Checks a declared PLP length before it is used to size an allocation. + /// + /// The wire carries the length as a signed 64-bit value. A negative one + /// would become a huge `usize` on cast, so it is rejected on the signed + /// value rather than after conversion. + fn validate_known_plp_length(long_len_i64: i64) -> TdsResult { + let length = long_len_i64 as usize; + if long_len_i64 < 0 || length > MAX_PLP_SIZE { + return Err(crate::error::Error::ProtocolError(format!( + "PLP length {length} (raw i64: {long_len_i64}) exceeds maximum allowed size of {MAX_PLP_SIZE} bytes" + ))); + } + Ok(length) + } + + /// Reads a PLP body of declared length into `destination`. + /// + /// The chunks have to add up to exactly `destination.len()`. A body that + /// ends early and one that overruns are both protocol errors, because the + /// destination may belong to the consumer and must not be left part filled + /// without the writer being told. + async fn read_known_plp_bytes(reader: &mut T, destination: &mut [u8]) -> TdsResult<()> + where + T: TdsPacketReader + Send + Sync, + { + let mut chunk_len = reader.read_uint32().await? as usize; + let mut offset = 0usize; + let mut chunk_count = 0u32; + + while chunk_len > 0 { + chunk_count += 1; + Self::validate_plp_chunk(chunk_count, chunk_len)?; + + let end_offset = offset.checked_add(chunk_len).ok_or_else(|| { + crate::error::Error::ProtocolError(format!( + "PLP chunk offset would overflow: {offset} + {chunk_len}" + )) + })?; + if end_offset > destination.len() { + return Err(crate::error::Error::ProtocolError(format!( + "PLP chunk exceeds declared length: offset={offset}, chunk_len={chunk_len}, declared_len={}", + destination.len() + ))); + } + let chunk_size_read = reader + .read_bytes(&mut destination[offset..end_offset]) + .await?; + if chunk_size_read != chunk_len { + return Err(crate::error::Error::ProtocolError(format!( + "Short PLP chunk read: expected={chunk_len}, actual={chunk_size_read}" + ))); } - Ok(Some(plp_buffer)) + offset = end_offset; + chunk_len = reader.read_uint32().await? as usize; } + + if offset != destination.len() { + return Err(crate::error::Error::ProtocolError(format!( + "PLP data shorter than declared length: actual={offset}, declared={}", + destination.len() + ))); + } + Ok(()) + } + + /// Reads a PLP body whose total length the sender did not declare. + /// + /// Chunks are appended until a zero-length chunk ends the value, so the + /// size is only known once the read finishes. That is why this form cannot + /// use a writer destination. + async fn read_unknown_plp_bytes(reader: &mut T) -> TdsResult> + where + T: TdsPacketReader + Send + Sync, + { + let mut bytes = Vec::new(); + let mut chunk_len = reader.read_uint32().await? as usize; + let mut chunk_count = 0u32; + + while chunk_len > 0 { + chunk_count += 1; + Self::validate_plp_chunk(chunk_count, chunk_len)?; + + let offset = bytes.len(); + let new_len = offset.checked_add(chunk_len).ok_or_else(|| { + crate::error::Error::ProtocolError(format!( + "PLP chunk accumulation would overflow capacity: {offset} + {chunk_len}" + )) + })?; + if new_len > MAX_PLP_SIZE { + return Err(crate::error::Error::ProtocolError(format!( + "PLP accumulated size {new_len} exceeds maximum allowed size of {MAX_PLP_SIZE} bytes (SQL Server limit: 2GB)" + ))); + } + bytes.resize(new_len, 0); + let chunk_size_read = reader.read_bytes(&mut bytes[offset..new_len]).await?; + if chunk_size_read != chunk_len { + return Err(crate::error::Error::ProtocolError(format!( + "Short PLP chunk read: expected={chunk_len}, actual={chunk_size_read}" + ))); + } + chunk_len = reader.read_uint32().await? as usize; + } + Ok(bytes) + } + + /// Caps how long an undeclared PLP stream is allowed to run. + /// + /// Without a declared total length the only protection is a bound on the + /// number of chunks and on each chunk's size. Both limits are tightened + /// under `fuzzing` so a generated input fails fast. + fn validate_plp_chunk(chunk_count: u32, chunk_len: usize) -> TdsResult<()> { + #[cfg(fuzzing)] + const MAX_PLP_CHUNKS: u32 = 1000; + #[cfg(not(fuzzing))] + const MAX_PLP_CHUNKS: u32 = 100000; + #[cfg(fuzzing)] + const MAX_CHUNK_SIZE: usize = 8 * 1024; + #[cfg(not(fuzzing))] + const MAX_CHUNK_SIZE: usize = 16 * 1024 * 1024; + + if chunk_count > MAX_PLP_CHUNKS { + return Err(crate::error::Error::ProtocolError(format!( + "Too many PLP chunks: {chunk_count} (max {MAX_PLP_CHUNKS})" + ))); + } + if chunk_len > MAX_CHUNK_SIZE { + return Err(crate::error::Error::ProtocolError(format!( + "PLP chunk size {chunk_len} exceeds maximum allowed chunk size of {MAX_CHUNK_SIZE} bytes" + ))); + } + Ok(()) } /// Decodes a column value from the wire and writes it directly into a @@ -1103,23 +1247,23 @@ 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_byte_buffered(reader).await?); } TdsDataType::Int2 => { writer.write_i16(col, reader.read_int16().await?); } TdsDataType::Int4 => { - writer.write_i32(col, reader.read_int32().await?); + writer.write_i32(col, read_int32_buffered(reader).await?); } TdsDataType::Int8 => { writer.write_i64(col, reader.read_int64().await?); } TdsDataType::IntN => { - let byte_len = reader.read_byte().await?; + let byte_len = read_byte_buffered(reader).await?; match byte_len { - 1 => writer.write_u8(col, reader.read_byte().await?), + 1 => writer.write_u8(col, read_byte_buffered(reader).await?), 2 => writer.write_i16(col, reader.read_int16().await?), - 4 => writer.write_i32(col, reader.read_int32().await?), + 4 => writer.write_i32(col, read_int32_buffered(reader).await?), 8 => writer.write_i64(col, reader.read_int64().await?), 0 => writer.write_null(col), _ => { @@ -1209,14 +1353,18 @@ impl GenericDecoder { TdsDataType::BigBinary => { let length = reader.read_uint16().await?; // 0xFFFF is the USHORTLEN NULL marker (CHARBIN_NULL). - if length == 0xFFFF { + if length == u16::MAX { writer.write_null(col); + return Ok(()); + } + if length as usize > MAX_ALLOC_SIZE { + return Err(crate::error::Error::ProtocolError(format!( + "BigBinary length {length} exceeds maximum allowed size of {MAX_ALLOC_SIZE} bytes" + ))); + } + if let Some(bytes) = reader.try_read_slice(length as usize) { + writer.write_bytes_ref(col, bytes); } else { - if length as usize > MAX_ALLOC_SIZE { - return Err(crate::error::Error::ProtocolError(format!( - "BigBinary length {length} exceeds maximum allowed size of {MAX_ALLOC_SIZE} bytes" - ))); - } let mut bytes = vec![0u8; length as usize]; reader.read_bytes(&mut bytes).await?; writer.write_bytes(col, bytes); @@ -1224,21 +1372,22 @@ impl GenericDecoder { } TdsDataType::BigVarBinary => { if metadata.is_plp() { - match GenericDecoder::read_plp_bytes(reader).await? { - Some(bytes) => writer.write_bytes(col, bytes), - None => writer.write_null(col), - } + GenericDecoder::read_plp_binary_into(reader, col, writer).await?; } else { let length = reader.read_uint16().await?; // 0xFFFF is the USHORTLEN NULL marker (CHARBIN_NULL). - if length == 0xFFFF { + if length == u16::MAX { writer.write_null(col); + return Ok(()); + } + if length as usize > MAX_ALLOC_SIZE { + return Err(crate::error::Error::ProtocolError(format!( + "BigVarBinary length {length} exceeds maximum allowed size of {MAX_ALLOC_SIZE} bytes" + ))); + } + if let Some(bytes) = reader.try_read_slice(length as usize) { + writer.write_bytes_ref(col, bytes); } else { - if length as usize > MAX_ALLOC_SIZE { - return Err(crate::error::Error::ProtocolError(format!( - "BigVarBinary length {length} exceeds maximum allowed size of {MAX_ALLOC_SIZE} bytes" - ))); - } let mut bytes = vec![0u8; length as usize]; reader.read_bytes(&mut bytes).await?; writer.write_bytes(col, bytes); @@ -1368,6 +1517,62 @@ impl GenericDecoder { } Ok(()) } + + /// Decodes one column according to a resolved [`DecodeOp`]. + /// + /// Nothing calls this yet. The async path still decodes straight from + /// [`ColumnMetadata`] in `drive_row_columns`, which also has to handle + /// cursor policies, Always Encrypted and paused rows. Applying the plan + /// there is the remaining half of this change. + /// + /// Reads go through the `*_buffered` helpers, which stay off the async + /// machinery while the bytes are already in the reader's buffer. + #[allow(dead_code)] + pub(crate) async fn decode_op_into( + &self, + reader: &mut T, + operation: &DecodeOp, + metadata: &ColumnMetadata, + col: usize, + writer: &mut W, + ) -> TdsResult<()> + where + T: TdsPacketReader + Send + Sync, + W: RowWriter + ?Sized, + { + match operation { + DecodeOp::IntN => { + let byte_len = read_byte_buffered(reader).await?; + match byte_len { + 1 => writer.write_u8(col, read_byte_buffered(reader).await?), + 2 => writer.write_i16(col, reader.read_int16().await?), + 4 => writer.write_i32(col, read_int32_buffered(reader).await?), + 8 => writer.write_i64(col, reader.read_int64().await?), + 0 => writer.write_null(col), + _ => { + return Err(crate::error::Error::from(Error::new( + std::io::ErrorKind::InvalidData, + "Invalid IntN length", + ))); + } + } + Ok(()) + } + DecodeOp::ShortString(encoding_type) => { + self.string_decoder + .decode_short_string_into(reader, encoding_type, col, writer) + .await + } + DecodeOp::BitN + | DecodeOp::FltN + | DecodeOp::Decimal { .. } + | DecodeOp::DateN + | DecodeOp::TimeN { .. } + | DecodeOp::DateTime2N { .. } + | DecodeOp::ShortBinary + | DecodeOp::Generic => self.decode_into(reader, metadata, col, writer).await, + } + } } #[async_trait] @@ -1716,10 +1921,7 @@ impl StringDecoder { let encoding_type = get_encoding_type(metadata); if metadata.is_plp() { - match GenericDecoder::read_plp_bytes(reader).await? { - Some(bytes) => writer.write_string(col, SqlString::new(bytes, encoding_type)), - None => writer.write_null(col), - } + GenericDecoder::read_plp_string_into(reader, col, writer, &encoding_type).await?; } else if Self::is_long_len_type(metadata.data_type) { let text_ptr_len = reader.read_byte().await? as usize; @@ -1748,14 +1950,39 @@ impl StringDecoder { }; writer.write_string(col, sql_string); } else { - let length = reader.read_uint16().await? as usize; - if length == 0xFFFF { - writer.write_null(col); - } else { - let mut buffer = vec![0u8; length]; - reader.read_bytes(&mut buffer).await?; - writer.write_string(col, SqlString::new(buffer, encoding_type)); - } + self.decode_short_string_into(reader, &encoding_type, col, writer) + .await?; + } + Ok(()) + } + + /// Decodes a non-PLP character value. + /// + /// When the bytes are already buffered they reach the writer as a borrow, + /// and only a writer that needs to keep them pays for a copy. Otherwise the + /// value is read into an owned buffer. + /// + /// A length of `0xFFFF` is the wire encoding for NULL, not a real length. + async fn decode_short_string_into( + &self, + reader: &mut T, + encoding_type: &EncodingType, + col: usize, + writer: &mut W, + ) -> TdsResult<()> + where + T: TdsPacketReader + Send + Sync, + W: RowWriter + ?Sized, + { + let length = read_uint16_buffered(reader).await? as usize; + if length == 0xFFFF { + writer.write_null(col); + } else if let Some(bytes) = reader.try_read_slice(length) { + writer.write_string_ref(col, bytes, encoding_type); + } else { + let mut buffer = vec![0u8; length]; + read_bytes_buffered(reader, &mut buffer).await?; + writer.write_string(col, SqlString::new(buffer, encoding_type.clone())); } Ok(()) } @@ -3142,17 +3369,24 @@ mod test { use byteorder::{ByteOrder, LittleEndian}; use crate::core::TdsResult; - use crate::datatypes::column_values::{ColumnValues, SqlDateTime, SqlSmallDateTime}; + use crate::datatypes::column_values::{ + ColumnValues, SqlDate, SqlDateTime, SqlDateTime2, SqlDateTimeOffset, SqlMoney, + SqlSmallDateTime, SqlSmallMoney, SqlTime, SqlXml, + }; use crate::datatypes::decoder::{ - GenericDecoder, MAX_PLP_SIZE, PlpChunkReadLength, PlpChunkStreamReader, + DecimalParts, GenericDecoder, MAX_PLP_SIZE, PlpChunkReadLength, PlpChunkStreamReader, PlpColumnStream, SqlTypeDecode, }; - use crate::datatypes::row_writer::DefaultRowWriter; + use crate::datatypes::row_writer::{DefaultRowWriter, RowWriter}; + use crate::datatypes::sql_json::SqlJson; + use crate::datatypes::sql_string::SqlString; + use crate::datatypes::sql_vector::SqlVector; use crate::datatypes::sqldatatypes::{ PartialLengthType, TdsDataType, TypeInfo, TypeInfoVariant, VariableLengthTypes, }; use crate::io::packet_reader::TdsPacketReader; use crate::query::metadata::ColumnMetadata; + use uuid::Uuid; /// Byte-buffer backed mock implementing every `TdsPacketReader` method /// used by the decoder. @@ -3300,6 +3534,88 @@ mod test { } } + fn plp_binary_metadata() -> ColumnMetadata { + ColumnMetadata { + user_type: 0, + flags: 0, + data_type: TdsDataType::BigVarBinary, + type_info: TypeInfo { + tds_type: TdsDataType::BigVarBinary, + length: usize::MAX, + type_info_variant: TypeInfoVariant::PartialLen( + PartialLengthType::BigVarBinary, + None, + None, + None, + None, + ), + }, + column_name: String::new(), + multi_part_name: None, + crypto_metadata: None, + } + } + + fn plp_wire_chunked(declared_length: i64, chunks: &[&[u8]]) -> Vec { + let mut bytes = declared_length.to_le_bytes().to_vec(); + for chunk in chunks { + bytes.extend_from_slice(&(chunk.len() as u32).to_le_bytes()); + bytes.extend_from_slice(chunk); + } + bytes.extend_from_slice(&0u32.to_le_bytes()); + bytes + } + + #[derive(Default)] + struct BinaryDestinationWriter { + bytes: Vec, + is_null: bool, + used_destination: bool, + used_owned_fallback: bool, + destination_success: Option, + } + + impl RowWriter for BinaryDestinationWriter { + fn write_null(&mut self, _col: usize) { + self.is_null = true; + } + fn write_bool(&mut self, _col: usize, _val: bool) {} + fn write_u8(&mut self, _col: usize, _val: u8) {} + fn write_i16(&mut self, _col: usize, _val: i16) {} + fn write_i32(&mut self, _col: usize, _val: i32) {} + fn write_i64(&mut self, _col: usize, _val: i64) {} + fn write_f32(&mut self, _col: usize, _val: f32) {} + fn write_f64(&mut self, _col: usize, _val: f64) {} + fn write_string(&mut self, _col: usize, _val: SqlString) {} + fn write_bytes(&mut self, _col: usize, val: Vec) { + self.bytes = val; + self.used_owned_fallback = true; + } + fn bytes_destination(&mut self, _col: usize, length: usize) -> Option<&mut [u8]> { + self.bytes.resize(length, 0); + self.used_destination = true; + Some(&mut self.bytes) + } + fn finish_bytes_destination(&mut self, _col: usize, success: bool) { + self.destination_success = Some(success); + } + fn write_decimal(&mut self, _col: usize, _val: DecimalParts) {} + fn write_numeric(&mut self, _col: usize, _val: DecimalParts) {} + fn write_date(&mut self, _col: usize, _val: SqlDate) {} + fn write_time(&mut self, _col: usize, _val: SqlTime) {} + fn write_datetime(&mut self, _col: usize, _val: SqlDateTime) {} + fn write_smalldatetime(&mut self, _col: usize, _val: SqlSmallDateTime) {} + fn write_datetime2(&mut self, _col: usize, _val: SqlDateTime2) {} + fn write_datetimeoffset(&mut self, _col: usize, _val: SqlDateTimeOffset) {} + fn write_money(&mut self, _col: usize, _val: SqlMoney) {} + fn write_smallmoney(&mut self, _col: usize, _val: SqlSmallMoney) {} + fn write_uuid(&mut self, _col: usize, _val: Uuid) {} + fn write_xml(&mut self, _col: usize, _val: SqlXml) {} + fn write_json(&mut self, _col: usize, _val: SqlJson) {} + fn write_vector(&mut self, _col: usize, _val: SqlVector) {} + fn end_row(&mut self) {} + } + /// Runs both decode() and decode_into() on the same bytes and asserts /// that decode_into via DefaultRowWriter produces the same ColumnValues /// as decode(). @@ -4044,6 +4360,55 @@ mod test { assert_eq!(val, ColumnValues::Bytes(vec![0xDE, 0xAD, 0xBE, 0xEF])); } + #[tokio::test] + async fn decode_into_known_plp_binary_uses_destination() { + let metadata = plp_binary_metadata(); + let mut reader = ByteReader::new(plp_wire_chunked(6, &[b"ab", b"cdef"])); + let mut writer = BinaryDestinationWriter::default(); + + GenericDecoder::default() + .decode_into(&mut reader, &metadata, 0, &mut writer) + .await + .unwrap(); + + assert_eq!(writer.bytes, b"abcdef"); + assert!(writer.used_destination); + assert!(!writer.used_owned_fallback); + assert_eq!(writer.destination_success, Some(true)); + } + + #[tokio::test] + async fn decode_into_unknown_plp_binary_uses_owned_fallback() { + let metadata = plp_binary_metadata(); + let mut reader = ByteReader::new(plp_wire_chunked(-2, &[b"ab", b"cdef"])); + let mut writer = BinaryDestinationWriter::default(); + + GenericDecoder::default() + .decode_into(&mut reader, &metadata, 0, &mut writer) + .await + .unwrap(); + + assert_eq!(writer.bytes, b"abcdef"); + assert!(!writer.used_destination); + assert!(writer.used_owned_fallback); + assert_eq!(writer.destination_success, None); + } + + #[tokio::test] + async fn decode_into_known_plp_binary_rejects_length_mismatch() { + let metadata = plp_binary_metadata(); + let mut reader = ByteReader::new(plp_wire_chunked(3, &[b"abcd"])); + let mut writer = BinaryDestinationWriter::default(); + + let error = GenericDecoder::default() + .decode_into(&mut reader, &metadata, 0, &mut writer) + .await + .unwrap_err(); + + assert!(error.to_string().contains("exceeds declared length")); + assert_eq!(writer.destination_success, Some(false)); + } + #[tokio::test] async fn decode_into_nvarchar() { // non-PLP, non-LOB: u16 length + bytes diff --git a/mssql-tds/src/datatypes/row_writer.rs b/mssql-tds/src/datatypes/row_writer.rs index 53c146fd..ffcc940c 100644 --- a/mssql-tds/src/datatypes/row_writer.rs +++ b/mssql-tds/src/datatypes/row_writer.rs @@ -7,7 +7,7 @@ use crate::datatypes::column_values::{ }; use crate::datatypes::decoder::DecimalParts; use crate::datatypes::sql_json::SqlJson; -use crate::datatypes::sql_string::SqlString; +use crate::datatypes::sql_string::{EncodingType, SqlString}; use crate::datatypes::sql_vector::SqlVector; use uuid::Uuid; @@ -35,8 +35,38 @@ pub trait RowWriter { fn write_f64(&mut self, col: usize, val: f64); /// Writes a character string value. fn write_string(&mut self, col: usize, val: SqlString); + /// Writes packet-backed encoded string bytes before the packet buffer is reused. + fn write_string_ref(&mut self, col: usize, bytes: &[u8], encoding_type: &EncodingType) { + self.write_string(col, SqlString::new(bytes.to_vec(), encoding_type.clone())); + } /// Writes a binary value. fn write_bytes(&mut self, col: usize, val: Vec); + /// Writes packet-backed binary bytes before the packet buffer is reused. + fn write_bytes_ref(&mut self, col: usize, bytes: &[u8]) { + self.write_bytes(col, bytes.to_vec()); + } + /// Returns a destination for a known-length binary value, when the consumer + /// can allocate its final representation before the decoder reads the bytes. + fn bytes_destination(&mut self, _col: usize, _length: usize) -> Option<&mut [u8]> { + None + } + /// Completes or abandons a destination returned by [`Self::bytes_destination`]. + fn finish_bytes_destination(&mut self, _col: usize, _success: bool) {} + /// Returns a destination for a known-length character value, when the consumer + /// can allocate its final representation before the decoder reads the bytes. + /// + /// `length` counts encoded bytes on the wire, not characters, so a consumer + /// whose storage differs in length under `encoding_type` returns `None`. + fn string_destination( + &mut self, + _col: usize, + _length: usize, + _encoding_type: &EncodingType, + ) -> Option<&mut [u8]> { + None + } + /// Completes or abandons a destination returned by [`Self::string_destination`]. + fn finish_string_destination(&mut self, _col: usize, _success: bool) {} /// Writes a `decimal` value. fn write_decimal(&mut self, col: usize, val: DecimalParts); /// Writes a `numeric` value. diff --git a/mssql-tds/src/datatypes/sql_string.rs b/mssql-tds/src/datatypes/sql_string.rs index 0c00ad07..dd1fb66a 100644 --- a/mssql-tds/src/datatypes/sql_string.rs +++ b/mssql-tds/src/datatypes/sql_string.rs @@ -3,7 +3,7 @@ use crate::{query::metadata::ColumnMetadata, token::tokens::SqlCollation}; use core::fmt; -use std::{fmt::Debug, fmt::Display}; +use std::{borrow::Cow, fmt::Debug, fmt::Display}; use tracing::warn; use super::{ @@ -12,7 +12,7 @@ use super::{ }; /// Character encoding used by a [`SqlString`]. -#[derive(PartialEq, Clone)] +#[derive(Debug, PartialEq, Clone)] pub enum EncodingType { /// UTF-8 encoding. Utf8, @@ -54,50 +54,7 @@ impl SqlString { /// Decodes the stored bytes into a Rust `String` according to the encoding type. pub fn to_utf8_string(&self) -> String { - match self.encoding_type { - // TODO: Investigation needed. When creating a Utf8 strings from the vector, the string is weirdly encoded. - // UTF16 decode works better. - EncodingType::Utf8 => String::from_utf8(self.bytes.clone()).unwrap(), - EncodingType::Utf16 => { - // Use encoding_rs for efficient UTF-16LE decoding without intermediate Vec allocation - let (decoded, _, _) = encoding_rs::UTF_16LE.decode(&self.bytes); - decoded.into_owned() - } - EncodingType::LcidBased(collation) => { - // Extract LCID from the lower 20 bits of collation.info - let lcid = collation.info & 0x000F_FFFF; - - // Map LCID to encoding - let encoding = match lcid_to_encoding(lcid) { - Ok(enc) => enc, - Err(e) => { - warn!( - "Unsupported LCID 0x{:04X} ({}), falling back to Windows-1252. Error: {}", - lcid, lcid, e - ); - // Fall back to Windows-1252 for unsupported LCIDs - encoding_rs::WINDOWS_1252 - } - }; - - // Decode bytes using the determined encoding - let (decoded, _used_encoding, had_errors) = encoding.decode(&self.bytes); - - if had_errors { - warn!( - "Encountered decoding errors while converting LCID 0x{:04X} ({}) encoded data. \ - Some characters may have been replaced with U+FFFD.", - lcid, lcid - ); - } - - decoded.into_owned() - } - EncodingType::DelayedSet => { - // DelayedSet encoding is not defined, so we return the bytes as a UTF-8 string. - unimplemented!("DelayedSet encoding conversion to UTF8 not implemented"); - } - } + decode_bytes_to_utf8(&self.bytes, &self.encoding_type).into_owned() } /// Returns true if this SqlString is already encoded as UTF-16 @@ -135,6 +92,55 @@ impl SqlString { } } +/// Decodes encoded SQL string bytes into UTF-8. +/// +/// The result borrows the input whenever no transcoding is needed. That covers +/// the UTF-8 case and the single-byte collation case when the content is ASCII. +/// UTF-16 always allocates. Returning a `Cow` lets a caller that only reads the +/// text avoid a copy, while [`SqlString::to_utf8_string`] keeps its owned +/// return by calling `into_owned`. +/// +/// # Panics +/// +/// Panics when `EncodingType::Utf8` bytes are not valid UTF-8, and on +/// `EncodingType::DelayedSet`, which has no conversion defined. +pub fn decode_bytes_to_utf8<'a>(bytes: &'a [u8], encoding_type: &EncodingType) -> Cow<'a, str> { + match encoding_type { + // TODO: Investigation needed. When creating a Utf8 string from the bytes, + // the string is weirdly encoded. UTF16 decode works better. + EncodingType::Utf8 => Cow::Borrowed(std::str::from_utf8(bytes).unwrap()), + EncodingType::Utf16 => { + let (decoded, _, _) = encoding_rs::UTF_16LE.decode(bytes); + decoded + } + EncodingType::LcidBased(collation) => { + let lcid = collation.info & 0x000F_FFFF; + let encoding = match lcid_to_encoding(lcid) { + Ok(enc) => enc, + Err(e) => { + warn!( + "Unsupported LCID 0x{:04X} ({}), falling back to Windows-1252. Error: {}", + lcid, lcid, e + ); + encoding_rs::WINDOWS_1252 + } + }; + let (decoded, _used_encoding, had_errors) = encoding.decode(bytes); + if had_errors { + warn!( + "Encountered decoding errors while converting LCID 0x{:04X} ({}) encoded data. \ + Some characters may have been replaced with U+FFFD.", + lcid, lcid + ); + } + decoded + } + EncodingType::DelayedSet => { + unimplemented!("DelayedSet encoding conversion to UTF8 not implemented"); + } + } +} + impl Debug for SqlString { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self.encoding_type { diff --git a/mssql-tds/src/io/packet_reader.rs b/mssql-tds/src/io/packet_reader.rs index da238125..b1c8a06c 100644 --- a/mssql-tds/src/io/packet_reader.rs +++ b/mssql-tds/src/io/packet_reader.rs @@ -38,6 +38,68 @@ pub(crate) trait TdsPacketReader { async fn skip_bytes(&mut self, skip_count: usize) -> TdsResult<()>; async fn cancel_read_stream(&mut self) -> TdsResult<()>; fn reset_reader(&mut self); + + /* + * Non-blocking reads. + * + * Each of these returns a value when the reader's buffer already holds + * enough bytes, and None otherwise. None is not an error. It means the + * caller has to fall back to the matching async method, which is able to + * wait for the next packet. + * + * They exist because one TDS packet normally carries many rows. Once the + * packet has landed, every read inside it is a memory move, but the async + * form still allocates a boxed future on each call to find that out. + * + * The defaults all decline, so a reader that keeps no buffer of its own + * compiles unchanged and never takes the fast path. + */ + + /// Reads one byte if the buffer already holds it. + fn try_read_byte(&mut self) -> Option { + None + } + + /// Reads a little-endian `u16` if the buffer already holds it. + fn try_read_uint16(&mut self) -> Option { + None + } + + /// Reads a little-endian `i32` if the buffer already holds it. + fn try_read_int32(&mut self) -> Option { + None + } + + /// Fills `buffer` if that many bytes are already held. + fn try_read_bytes(&mut self, _buffer: &mut [u8]) -> Option { + None + } + + /// Borrows `length` bytes and consumes them. + /// + /// The borrow ends the next time the reader is used, so the caller has to + /// copy or hand off the bytes before then. + fn try_read_slice(&mut self, _length: usize) -> Option<&[u8]> { + None + } + + /// Borrows everything the buffer currently holds, consuming nothing. + /// + /// The caller decides afterwards how much it used and reports that through + /// [`Self::consume_buffered`]. This is what lets a whole row be measured + /// and then decoded under one borrow. + fn buffered_slice(&self) -> Option<&[u8]> { + None + } + + /// Advances past `length` bytes that were seen through + /// [`Self::buffered_slice`]. + /// + /// Returns `false` when the buffer no longer holds that many bytes, which + /// means the caller's view of it went stale. + fn consume_buffered(&mut self, _length: usize) -> bool { + false + } } /// Low-level TDS packet reading operations (public under `fuzzing` cfg). @@ -69,6 +131,91 @@ pub trait TdsPacketReader { async fn skip_bytes(&mut self, skip_count: usize) -> TdsResult<()>; async fn cancel_read_stream(&mut self) -> TdsResult<()>; fn reset_reader(&mut self); + + fn try_read_byte(&mut self) -> Option { + None + } + + fn try_read_uint16(&mut self) -> Option { + None + } + + fn try_read_int32(&mut self) -> Option { + None + } + + fn try_read_bytes(&mut self, _buffer: &mut [u8]) -> Option { + None + } + + fn try_read_slice(&mut self, _length: usize) -> Option<&[u8]> { + None + } + + fn buffered_slice(&self) -> Option<&[u8]> { + None + } + + fn consume_buffered(&mut self, _length: usize) -> bool { + false + } +} + +/* + * Buffered read helpers. + * + * Each one tries the non-blocking form and falls back to the async form. They + * are free functions rather than trait methods so that the buffered case stays + * a direct call. A trait method under #[async_trait] would allocate a boxed + * future first and only then discover that the bytes were already in hand. + */ + +/// Reads one byte, from the buffer when it is there and from the network otherwise. +#[inline(always)] +pub(crate) async fn read_byte_buffered(reader: &mut T) -> TdsResult +where + T: TdsPacketReader + Send + Sync + ?Sized, +{ + match reader.try_read_byte() { + Some(value) => Ok(value), + None => reader.read_byte().await, + } +} + +/// Reads a little-endian `u16`, from the buffer when it is there. +#[inline(always)] +pub(crate) async fn read_uint16_buffered(reader: &mut T) -> TdsResult +where + T: TdsPacketReader + Send + Sync + ?Sized, +{ + match reader.try_read_uint16() { + Some(value) => Ok(value), + None => reader.read_uint16().await, + } +} + +/// Reads a little-endian `i32`, from the buffer when it is there. +#[inline(always)] +pub(crate) async fn read_int32_buffered(reader: &mut T) -> TdsResult +where + T: TdsPacketReader + Send + Sync + ?Sized, +{ + match reader.try_read_int32() { + Some(value) => Ok(value), + None => reader.read_int32().await, + } +} + +/// Fills `buffer`, from the reader's buffer when it holds that many bytes. +#[inline(always)] +pub(crate) async fn read_bytes_buffered(reader: &mut T, buffer: &mut [u8]) -> TdsResult +where + T: TdsPacketReader + Send + Sync + ?Sized, +{ + match reader.try_read_bytes(buffer) { + Some(bytes_read) => Ok(bytes_read), + None => reader.read_bytes(buffer).await, + } } // Blanket implementation for Box to enable dynamic dispatch diff --git a/mssql-tds/src/io/token_stream.rs b/mssql-tds/src/io/token_stream.rs index 87a84e48..d364b8b3 100644 --- a/mssql-tds/src/io/token_stream.rs +++ b/mssql-tds/src/io/token_stream.rs @@ -2,9 +2,12 @@ // Licensed under the MIT License. use crate::core::{CancelHandle, TdsResult}; -use crate::datatypes::decoder::{GenericDecoder, PlpColumnStream, decrypt_encrypted_column}; +use crate::datatypes::column_values::{SqlDate, SqlDateTime2, SqlTime}; +use crate::datatypes::decoder::{ + DecimalParts, GenericDecoder, PlpColumnStream, decrypt_encrypted_column, +}; use crate::datatypes::row_writer::{DiscardRowWriter, RowWriter, write_column_value}; -use crate::io::packet_reader::TdsPacketReader; +use crate::io::packet_reader::{TdsPacketReader, read_byte_buffered}; use crate::query::metadata::ColumnMetadata; use crate::security::cell_decryptor::CellDecryptor; use crate::token::parsers::TokenParser; @@ -15,8 +18,10 @@ use crate::token::parsers::{ OrderTokenParser, ReturnStatusTokenParser, ReturnValueTokenParser, RowTokenParser, SessionStateTokenParser, SspiTokenParser, TabNameTokenParser, }; +use crate::token::tokens::DecodeOp; use crate::token::tokens::{ColMetadataToken, TokenType, Tokens}; use async_trait::async_trait; +use byteorder::{ByteOrder, LittleEndian}; use core::convert::From; use std::collections::HashMap; use std::sync::Arc; @@ -279,6 +284,414 @@ pub trait TdsTokenStreamReader { ) -> TdsResult; } +/// Decodes one ROW or NBCROW token straight out of the reader's current buffer, +/// without entering the async packet path. +/// +/// # Why this path exists +/// +/// The async path reads every column with its own `await`, and each of those +/// reads crosses `#[async_trait]`, which allocates a boxed future. A column +/// costs two such reads, one for the length prefix and one for the data. The +/// row is normally sitting whole in the read buffer the entire time, so all of +/// that work happens to move bytes that have already arrived. +/// +/// This function takes the whole row from that buffer under a single borrow. +/// +/// # The two passes +/// +/// Pass one measures the row. Every field is bounds checked. If the row runs +/// past the end of the buffer the pass returns `Ok(None)`, and the caller falls +/// back to the async path, which is able to wait for the next packet. At that +/// point nothing has been written to the writer and nothing has been consumed +/// from the reader, so the fallback starts from a clean state. +/// +/// Pass two walks the same bytes and decodes them. It carries no bounds checks +/// and no early exits, because pass one already proved every offset is in +/// range. That is the whole reason for the split. A single pass would have to +/// check bounds at every column, and would still have to undo its writes when +/// it discovered the row was short. +/// +/// The two passes must stay in step. Any field whose width pass one computes +/// one way and pass two computes another will silently misread the rest of the +/// row. Change them together. +/// +/// `DecodeOp::Generic` marks a column this path does not implement. Pass one +/// declines on it, which is why pass two can treat it as `unreachable!()`. +/// +/// # Reader contract +/// +/// The reader is touched exactly twice. Once to borrow the buffer, and once to +/// consume the byte count pass two settled on. The borrow is held across the +/// whole body, so nothing in between can move the buffer underneath it. +/// +/// # Returns +/// +/// * `Ok(Some(RowReadResult::RowWritten))` - a row was decoded and consumed. +/// * `Ok(None)` - this path declined. Not an error. The caller must run the +/// async path instead. +fn try_receive_row_into_buffered( + reader: &mut R, + context: &ParserContext, + writer: &mut (dyn RowWriter + Send), +) -> TdsResult> { + let (metadata, decryptor) = extract_column_metadata(context)?; + // Encrypted cells arrive as ciphertext and only the async path knows how to + // decrypt them, so decline the whole row when a decryptor is present. + if decryptor.is_some() { + return Ok(None); + } + let columns = &metadata.columns; + let plan = metadata.decode_plan(); + + let consumed = { + let Some(buffer) = reader.buffered_slice() else { + return Ok(None); + }; + if buffer.is_empty() { + return Ok(None); + } + + let token = TokenType::try_from(buffer[0])?; + let is_nbc = match token { + TokenType::Row => false, + TokenType::NbcRow => true, + _ => return Ok(None), + }; + let bitmap_len = if is_nbc { columns.len().div_ceil(8) } else { 0 }; + let data_start = 1 + bitmap_len; + if buffer.len() < data_start { + return Ok(None); + } + + let is_null = |column: usize| is_nbc && buffer[1 + column / 8] & (1 << (column % 8)) != 0; + + /* + * pass 1: measure the row + * + * Bounds check every field. Return Ok(None) the moment the row runs + * past the buffer, before anything has been written or consumed. + */ + let mut position = data_start; + for (column, operation) in plan.iter().enumerate() { + if is_null(column) { + continue; + } + match operation { + DecodeOp::IntN => { + if position >= buffer.len() { + return Ok(None); + } + let length = buffer[position] as usize; + if !matches!(length, 0 | 1 | 2 | 4 | 8) { + return Err(crate::error::Error::ProtocolError(format!( + "Invalid IntN length: {length}" + ))); + } + position += 1; + if position + length > buffer.len() { + return Ok(None); + } + position += length; + } + DecodeOp::BitN => { + if position >= buffer.len() { + return Ok(None); + } + let length = buffer[position] as usize; + if !matches!(length, 0 | 1) { + return Err(crate::error::Error::ProtocolError(format!( + "Invalid BitN length: {length}" + ))); + } + position += 1; + if position + length > buffer.len() { + return Ok(None); + } + position += length; + } + DecodeOp::FltN => { + if position >= buffer.len() { + return Ok(None); + } + let length = buffer[position] as usize; + if !matches!(length, 0 | 4 | 8) { + return Err(crate::error::Error::ProtocolError(format!( + "Invalid FltN length: {length}" + ))); + } + position += 1; + if position + length > buffer.len() { + return Ok(None); + } + position += length; + } + DecodeOp::Decimal { .. } => { + if position >= buffer.len() { + return Ok(None); + } + let length = buffer[position] as usize; + if !matches!(length, 0 | 5 | 9 | 13 | 17) { + return Err(crate::error::Error::ProtocolError(format!( + "Invalid decimal length: {length}" + ))); + } + position += 1; + if position + length > buffer.len() { + return Ok(None); + } + position += length; + } + DecodeOp::DateN => { + if position >= buffer.len() { + return Ok(None); + } + let length = buffer[position] as usize; + if !matches!(length, 0 | 3) { + return Err(crate::error::Error::ProtocolError(format!( + "Invalid DateN length: {length}" + ))); + } + position += 1; + if position + length > buffer.len() { + return Ok(None); + } + position += length; + } + DecodeOp::TimeN { .. } => { + if position >= buffer.len() { + return Ok(None); + } + let length = buffer[position] as usize; + if !matches!(length, 0 | 3 | 4 | 5) { + return Err(crate::error::Error::ProtocolError(format!( + "Invalid TimeN length: {length}" + ))); + } + position += 1; + if position + length > buffer.len() { + return Ok(None); + } + position += length; + } + DecodeOp::DateTime2N { .. } => { + if position >= buffer.len() { + return Ok(None); + } + let length = buffer[position] as usize; + if !matches!(length, 0 | 6 | 7 | 8) { + return Err(crate::error::Error::ProtocolError(format!( + "Invalid DateTime2N length: {length}" + ))); + } + position += 1; + if position + length > buffer.len() { + return Ok(None); + } + position += length; + } + DecodeOp::ShortBinary => { + if position + 2 > buffer.len() { + return Ok(None); + } + let length = LittleEndian::read_u16(&buffer[position..]) as usize; + position += 2; + if length != 0xFFFF { + if position + length > buffer.len() { + return Ok(None); + } + position += length; + } + } + DecodeOp::ShortString(_) => { + if position + 2 > buffer.len() { + return Ok(None); + } + let length = LittleEndian::read_u16(&buffer[position..]) as usize; + position += 2; + if length != 0xFFFF { + if position + length > buffer.len() { + return Ok(None); + } + position += length; + } + } + DecodeOp::Generic => return Ok(None), + } + } + + /* + * pass 2: decode the row + * + * No bounds checks below. Pass 1 walked these exact offsets and proved + * they are in range. + */ + position = data_start; + for (column, operation) in plan.iter().enumerate() { + if is_null(column) { + writer.write_null(column); + continue; + } + match operation { + DecodeOp::IntN => { + let length = buffer[position] as usize; + position += 1; + match length { + 0 => writer.write_null(column), + 1 => writer.write_u8(column, buffer[position]), + 2 => writer.write_i16(column, LittleEndian::read_i16(&buffer[position..])), + 4 => writer.write_i32(column, LittleEndian::read_i32(&buffer[position..])), + 8 => writer.write_i64(column, LittleEndian::read_i64(&buffer[position..])), + _ => unreachable!(), + } + position += length; + } + DecodeOp::BitN => { + let length = buffer[position] as usize; + position += 1; + if length == 0 { + writer.write_null(column); + } else { + writer.write_bool(column, buffer[position] == 1); + position += 1; + } + } + DecodeOp::FltN => { + let length = buffer[position] as usize; + position += 1; + match length { + 0 => writer.write_null(column), + 4 => writer.write_f32(column, LittleEndian::read_f32(&buffer[position..])), + 8 => writer.write_f64(column, LittleEndian::read_f64(&buffer[position..])), + _ => unreachable!(), + } + position += length; + } + DecodeOp::Decimal { + precision, + scale, + numeric, + } => { + let length = buffer[position] as usize; + position += 1; + if length == 0 { + writer.write_null(column); + } else { + let is_positive = buffer[position] == 1; + position += 1; + let mut int_parts = Vec::with_capacity((length - 1) / 4); + for _ in 0..(length - 1) / 4 { + int_parts.push(LittleEndian::read_i32(&buffer[position..])); + position += 4; + } + let parts = DecimalParts { + is_positive, + scale: *scale, + precision: *precision, + int_parts, + }; + if *numeric { + writer.write_numeric(column, parts); + } else { + writer.write_decimal(column, parts); + } + } + } + DecodeOp::DateN => { + let length = buffer[position] as usize; + position += 1; + if length == 0 { + writer.write_null(column); + } else { + let date = SqlDate::create(LittleEndian::read_u24(&buffer[position..]))?; + position += 3; + writer.write_date(column, date); + } + } + DecodeOp::TimeN { scale } => { + let length = buffer[position] as usize; + position += 1; + if length == 0 { + writer.write_null(column); + } else { + let time = decode_buffered_time(&buffer[position..], length, *scale); + position += length; + writer.write_time(column, time); + } + } + DecodeOp::DateTime2N { scale } => { + let length = buffer[position] as usize; + position += 1; + if length == 0 { + writer.write_null(column); + } else { + let time_len = length - 3; + let time = decode_buffered_time(&buffer[position..], time_len, *scale); + position += time_len; + let days = LittleEndian::read_u24(&buffer[position..]); + position += 3; + writer.write_datetime2(column, SqlDateTime2 { days, time }); + } + } + DecodeOp::ShortBinary => { + let length = LittleEndian::read_u16(&buffer[position..]) as usize; + position += 2; + if length == 0xFFFF { + writer.write_null(column); + } else { + writer.write_bytes_ref(column, &buffer[position..position + length]); + position += length; + } + } + DecodeOp::ShortString(encoding) => { + let length = LittleEndian::read_u16(&buffer[position..]) as usize; + position += 2; + if length == 0xFFFF { + writer.write_null(column); + } else { + writer.write_string_ref( + column, + &buffer[position..position + length], + encoding, + ); + position += length; + } + } + DecodeOp::Generic => unreachable!(), + } + } + position + }; + + // The buffer borrow ends here, so this is the first point at which the + // reader can be advanced. It cannot fail today, because nothing released + // bytes while the borrow was held. It is checked so that a future change + // which breaks that invariant fails loudly instead of skipping a row. + if !reader.consume_buffered(consumed) { + return Err(crate::error::Error::ImplementationError( + "Buffered row changed before commit".to_string(), + )); + } + Ok(Some(RowReadResult::RowWritten)) +} + +fn decode_buffered_time(buffer: &[u8], length: usize, scale: u8) -> SqlTime { + let scaled = LittleEndian::read_uint(buffer, length); + let multiplier = match scale { + 0 => 10_000_000, + 1 => 1_000_000, + 2 => 100_000, + 3 => 10_000, + 4 => 1_000, + 5 => 100, + 6 => 10, + _ => 1, + }; + SqlTime { + time_nanoseconds: scaled * multiplier, + scale, + } +} + #[cfg(fuzzing)] pub struct TokenStreamReader where @@ -353,6 +766,19 @@ fn extract_row_context(context: &ParserContext) -> TdsResult TdsResult<(&ColMetadataToken, Option<&Arc>)> { + match context { + ParserContext::ColumnMetadata(metadata, decryptor) => Ok((metadata, decryptor.as_ref())), + _ => Err(crate::error::Error::ProtocolError( + "Expected ColumnMetadata in context for row decoding".to_string(), + )), + } +} + pub(crate) async fn dispatch_token( reader: &mut R, registry: &impl TokenParserRegistry, @@ -584,9 +1010,17 @@ pub(crate) async fn receive_row_into_internal( plan: ColumnPolicy, writer: &mut (dyn RowWriter + Send), ) -> TdsResult { - let token_type_byte = reader.read_byte().await?; + // The buffered path decodes every column of one whole row, so it can only + // stand in for DecodeAll. Cursor policies pause mid row and are left to + // drive_row_columns. + if matches!(plan, ColumnPolicy::DecodeAll) + && let Some(result) = try_receive_row_into_buffered(reader, context, writer)? + { + return Ok(result); + } + + let token_type_byte = read_byte_buffered(reader).await?; let token_type: TokenType = token_type_byte.try_into()?; - debug!("Parsing token type: {:?}", &token_type); match token_type { TokenType::Row => { @@ -1314,6 +1748,7 @@ mod tests { column_count: 1, columns: vec![metadata], cek_table: vec![], + ..Default::default() }), None, ); @@ -1361,6 +1796,7 @@ mod tests { plp_varbinary_metadata("c2", None), ], cek_table: vec![], + ..Default::default() }), None, ); @@ -1397,6 +1833,7 @@ mod tests { column_count: 1, columns: vec![plp_varbinary_metadata("c1", Some(ae_crypto_metadata()))], cek_table: vec![], + ..Default::default() }), None, ); diff --git a/mssql-tds/src/token/parsers/colmetadata_parser.rs b/mssql-tds/src/token/parsers/colmetadata_parser.rs index 70f144d9..d365a2bf 100644 --- a/mssql-tds/src/token/parsers/colmetadata_parser.rs +++ b/mssql-tds/src/token/parsers/colmetadata_parser.rs @@ -244,6 +244,7 @@ where column_count: col_count, columns: column_metadata, cek_table, + decode_plan: std::sync::OnceLock::new(), }; Ok(Tokens::from(metadata)) } diff --git a/mssql-tds/src/token/parsers/nbcrow_parser.rs b/mssql-tds/src/token/parsers/nbcrow_parser.rs index 5df89010..7bdd1051 100644 --- a/mssql-tds/src/token/parsers/nbcrow_parser.rs +++ b/mssql-tds/src/token/parsers/nbcrow_parser.rs @@ -154,6 +154,7 @@ mod tests { column_count: columns.len() as u16, columns, cek_table: Vec::new(), + decode_plan: std::sync::OnceLock::new(), }), None, ) diff --git a/mssql-tds/src/token/parsers/row_parser.rs b/mssql-tds/src/token/parsers/row_parser.rs index c4919aa8..831c5387 100644 --- a/mssql-tds/src/token/parsers/row_parser.rs +++ b/mssql-tds/src/token/parsers/row_parser.rs @@ -232,6 +232,7 @@ mod tests { column_count: columns.len() as u16, columns, cek_table: Vec::new(), + decode_plan: std::sync::OnceLock::new(), }), None, ) diff --git a/mssql-tds/src/token/tokens.rs b/mssql-tds/src/token/tokens.rs index 97392df3..9a55c03a 100644 --- a/mssql-tds/src/token/tokens.rs +++ b/mssql-tds/src/token/tokens.rs @@ -2,6 +2,7 @@ // Licensed under the MIT License. use std::fmt::{self, Debug}; +use std::sync::OnceLock; use super::{ fed_auth_info::{FedAuthInfoToken, SspiToken}, @@ -9,6 +10,7 @@ use super::{ tokenitems::ReturnValueStatus, }; use crate::datatypes::column_values::ColumnValues; +use crate::datatypes::sql_string::{EncodingType, get_encoding_type}; use crate::{ error::Error, message::login::{FeatureExtension, RoutingInfo}, @@ -311,6 +313,108 @@ pub(crate) struct ColMetadataToken { /// references entries here by ordinal. #[allow(dead_code)] // Consumed by CEK decryption in a later phase. pub cek_table: Vec, + /// Per-column decode plan, built on first use. See [`Self::decode_plan`]. + pub(crate) decode_plan: OnceLock>, +} + +/// What the fast row path should do with one column. +/// +/// A column's type, precision, scale and collation are fixed for the whole +/// result set, but the async decoder works them out from [`ColumnMetadata`] +/// again on every row. `DecodeOp` settles that once and reduces each column to +/// a single tag the row loop can switch on. +/// +/// `Generic` is the escape hatch. It marks a column the fast path does not +/// handle, and one such column sends the whole row to the async decoder. +#[derive(Debug, Clone)] +pub(crate) enum DecodeOp { + IntN, + BitN, + FltN, + Decimal { + precision: u8, + scale: u8, + numeric: bool, + }, + DateN, + TimeN { + scale: u8, + }, + DateTime2N { + scale: u8, + }, + ShortBinary, + ShortString(EncodingType), + Generic, +} + +impl ColMetadataToken { + /// Returns the per-column decode plan, building it on first use. + /// + /// The plan is derived from metadata that cannot change while the result + /// set is open, so it is built once and reused by every row. `OnceLock` + /// keeps that lazy without making the token mutable, which matters because + /// the token is shared behind an `Arc`. + /// + /// A column whose metadata does not match a shape the fast path knows + /// becomes [`DecodeOp::Generic`] rather than an error. Declining is always + /// safe, because the async decoder handles every type. + pub(crate) fn decode_plan(&self) -> &[DecodeOp] { + self.decode_plan.get_or_init(|| { + self.columns + .iter() + .map(|metadata| match metadata.data_type { + crate::datatypes::sqldatatypes::TdsDataType::IntN => DecodeOp::IntN, + crate::datatypes::sqldatatypes::TdsDataType::BitN => DecodeOp::BitN, + crate::datatypes::sqldatatypes::TdsDataType::FltN => DecodeOp::FltN, + crate::datatypes::sqldatatypes::TdsDataType::DecimalN + | crate::datatypes::sqldatatypes::TdsDataType::NumericN => { + let crate::datatypes::sqldatatypes::TypeInfoVariant::VarLenPrecisionScale( + _, + _, + precision, + scale, + ) = metadata.type_info.type_info_variant + else { + return DecodeOp::Generic; + }; + DecodeOp::Decimal { + precision, + scale, + numeric: metadata.data_type + == crate::datatypes::sqldatatypes::TdsDataType::NumericN, + } + } + crate::datatypes::sqldatatypes::TdsDataType::DateN => DecodeOp::DateN, + crate::datatypes::sqldatatypes::TdsDataType::TimeN => DecodeOp::TimeN { + scale: metadata.get_scale().unwrap_or(7), + }, + crate::datatypes::sqldatatypes::TdsDataType::DateTime2N => { + DecodeOp::DateTime2N { + scale: metadata.get_scale().unwrap_or(7), + } + } + crate::datatypes::sqldatatypes::TdsDataType::BigBinary + | crate::datatypes::sqldatatypes::TdsDataType::BigVarBinary + if !metadata.is_plp() => + { + DecodeOp::ShortBinary + } + crate::datatypes::sqldatatypes::TdsDataType::NChar + | crate::datatypes::sqldatatypes::TdsDataType::NVarChar + | crate::datatypes::sqldatatypes::TdsDataType::BigChar + | crate::datatypes::sqldatatypes::TdsDataType::BigVarChar + | crate::datatypes::sqldatatypes::TdsDataType::Char + | crate::datatypes::sqldatatypes::TdsDataType::VarChar + if !metadata.is_plp() => + { + DecodeOp::ShortString(get_encoding_type(metadata)) + } + _ => DecodeOp::Generic, + }) + .collect() + }) + } } impl Token for ColMetadataToken {