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 {