Fix mssql TDS packet splitting over TLS connections

Author: lovasoa

CHANGELOG.md

@@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Add sqlx version information to pre-login message in mssql
 - Add support for encrypted Microsoft SQL server connections (using TLS)
+- Add support for the `SSLKEYLOGFILE` environment variable for TLS decryption in Wireshark
 
 ## 0.6.27
 

sqlx-core/src/mssql/connection/establish.rs

@@ -26,18 +26,18 @@ impl MssqlConnection {
 
         log::debug!("Sending T-SQL PRELOGIN with encryption: {:?}", encryption);
 
-        stream.write_packet(
-            PacketType::PreLogin,
-            PreLogin {
-                version: Version::default(),
-                encryption,
-                instance: options.instance.clone(),
-
-                ..Default::default()
-            },
-        );
-
-        stream.flush().await?;
+        stream
+            .write_packet_and_flush(
+                PacketType::PreLogin,
+                PreLogin {
+                    version: Version::default(),
+                    encryption,
+                    instance: options.instance.clone(),
+
+                    ..Default::default()
+                },
+            )
+            .await?;
 
         let (_, packet) = stream.recv_packet().await?;
         let prelogin_response = PreLogin::decode(packet)?;
@@ -47,31 +47,36 @@ impl MssqlConnection {
             Encrypt::Required | Encrypt::On
         ) {
             stream.setup_encryption().await?;
+        } else if encryption == Encrypt::Required {
+            return Err(Error::Tls(Box::new(std::io::Error::new(
+                std::io::ErrorKind::Other,
+                "TLS encryption required but not supported by server",
+            ))));
         }
 
         // LOGIN7 defines the authentication rules for use between client and server
 
-        stream.write_packet(
-            PacketType::Tds7Login,
-            Login7 {
-                // FIXME: use a version constant
-                version: 0x74000004, // SQL Server 2012 - SQL Server 2019
-                client_program_version: options.client_program_version,
-                client_pid: options.client_pid,
-                packet_size: options.requested_packet_size, // max allowed size of TDS packet
-                hostname: &options.hostname,
-                username: &options.username,
-                password: options.password.as_deref().unwrap_or_default(),
-                app_name: &options.app_name,
-                server_name: &options.server_name,
-                client_interface_name: &options.client_interface_name,
-                language: &options.language,
-                database: &*options.database,
-                client_id: [0; 6],
-            },
-        );
-
-        stream.flush().await?;
+        stream
+            .write_packet_and_flush(
+                PacketType::Tds7Login,
+                Login7 {
+                    // FIXME: use a version constant
+                    version: 0x74000004, // SQL Server 2012 - SQL Server 2019
+                    client_program_version: options.client_program_version,
+                    client_pid: options.client_pid,
+                    packet_size: options.requested_packet_size, // max allowed size of TDS packet
+                    hostname: &options.hostname,
+                    username: &options.username,
+                    password: options.password.as_deref().unwrap_or_default(),
+                    app_name: &options.app_name,
+                    server_name: &options.server_name,
+                    client_interface_name: &options.client_interface_name,
+                    language: &options.language,
+                    database: &*options.database,
+                    client_id: [0; 6],
+                },
+            )
+            .await?;
 
         loop {
             // NOTE: we should receive an [Error] message if something goes wrong, otherwise,

sqlx-core/src/mssql/connection/executor.rs

@@ -41,27 +41,29 @@ impl MssqlConnection {
                 proc_args.append(&mut arguments);
             }
 
-            self.stream.write_packet(
-                PacketType::Rpc,
-                RpcRequest {
-                    transaction_descriptor: self.stream.transaction_descriptor,
-                    arguments: &proc_args,
-                    procedure: proc,
-                    options: OptionFlags::empty(),
-                },
-            );
+            self.stream
+                .write_packet_and_flush(
+                    PacketType::Rpc,
+                    RpcRequest {
+                        transaction_descriptor: self.stream.transaction_descriptor,
+                        arguments: &proc_args,
+                        procedure: proc,
+                        options: OptionFlags::empty(),
+                    },
+                )
+                .await?;
         } else {
-            self.stream.write_packet(
-                PacketType::SqlBatch,
-                SqlBatch {
-                    transaction_descriptor: self.stream.transaction_descriptor,
-                    sql: query,
-                },
-            );
+            self.stream
+                .write_packet_and_flush(
+                    PacketType::SqlBatch,
+                    SqlBatch {
+                        transaction_descriptor: self.stream.transaction_descriptor,
+                        sql: query,
+                    },
+                )
+                .await?;
         }
 
-        self.stream.flush().await?;
-
         Ok(())
     }
 }

sqlx-core/src/mssql/connection/prepare.rs

@@ -52,21 +52,22 @@ pub(crate) async fn prepare(
     args.add_unnamed(sql);
     args.add_unnamed(0x0001_i32); // 1 = SEND_METADATA
 
-    conn.stream.write_packet(
-        PacketType::Rpc,
-        RpcRequest {
-            transaction_descriptor: conn.stream.transaction_descriptor,
-            arguments: &args,
-            // [sp_prepare] will emit the column meta data
-            // small issue is that we need to declare all the used placeholders with a "fallback" type
-            // we currently use regex to collect them; false positives are *okay* but false
-            // negatives would break the query
-            procedure: Either::Right(Procedure::Prepare),
-            options: OptionFlags::empty(),
-        },
-    );
-
-    conn.stream.flush().await?;
+    conn.stream
+        .write_packet_and_flush(
+            PacketType::Rpc,
+            RpcRequest {
+                transaction_descriptor: conn.stream.transaction_descriptor,
+                arguments: &args,
+                // [sp_prepare] will emit the column meta data
+                // small issue is that we need to declare all the used placeholders with a "fallback" type
+                // we currently use regex to collect them; false positives are *okay* but false
+                // negatives would break the query
+                procedure: Either::Right(Procedure::Prepare),
+                options: OptionFlags::empty(),
+            },
+        )
+        .await?;
+
     conn.stream.wait_until_ready().await?;
     conn.stream.pending_done_count += 1;
 
@@ -100,17 +101,18 @@ pub(crate) async fn prepare(
         let mut args = MssqlArguments::default();
         args.add_unnamed(id);
 
-        conn.stream.write_packet(
-            PacketType::Rpc,
-            RpcRequest {
-                transaction_descriptor: conn.stream.transaction_descriptor,
-                arguments: &args,
-                procedure: Either::Right(Procedure::Unprepare),
-                options: OptionFlags::empty(),
-            },
-        );
+        conn.stream
+            .write_packet_and_flush(
+                PacketType::Rpc,
+                RpcRequest {
+                    transaction_descriptor: conn.stream.transaction_descriptor,
+                    arguments: &args,
+                    procedure: Either::Right(Procedure::Unprepare),
+                    options: OptionFlags::empty(),
+                },
+            )
+            .await?;
 
-        conn.stream.flush().await?;
         conn.stream.wait_until_ready().await?;
         conn.stream.pending_done_count += 1;
 

sqlx-core/src/mssql/connection/stream.rs

@@ -71,11 +71,39 @@ impl MssqlStream {
         })
     }
 
-    // writes the packet out to the write buffer
+    // writes the packet out to the write buffer, but does not flush
+    // WARNING: if the packet is large, and we are over an encrypted connection, this will fail, since we would
+    // need to flush each packet individually to the encryption layer
     pub(crate) fn write_packet<'en, T: Encode<'en>>(&mut self, ty: PacketType, payload: T) {
         write_packets(&mut self.inner.wbuf, self.max_packet_size, ty, payload)
     }
 
+    // writes the packet out to the write buffer, splitting it if neccessary, and flushing TDS packets one at a time
+    pub(crate) async fn write_packet_and_flush<'en, T: Encode<'en>>(
+        &mut self,
+        ty: PacketType,
+        payload: T,
+    ) -> Result<(), Error> {
+        if !self.inner.wbuf.is_empty() {
+            self.flush().await?;
+        }
+        self.write_packet(ty, payload);
+        self.flush().await?;
+        Ok(())
+    }
+
+    // writes the packet out to the write buffer, splitting it if neccessary, and flushing TDS packets one at a time
+    pub(crate) async fn flush(&mut self) -> Result<(), Error> {
+        // flush self.max_packet_size bytes at a time
+        while self.inner.wbuf.len() > self.max_packet_size {
+            let rest = self.inner.wbuf.split_off(self.max_packet_size);
+            self.inner.flush().await?;
+            self.inner.wbuf = rest;
+        }
+        self.inner.flush().await?;
+        Ok(())
+    }
+
     // receive the next packet from the database
     // blocks until a packet is available
     pub(super) async fn recv_packet(&mut self) -> Result<(PacketHeader, Bytes), Error> {

sqlx-core/src/mssql/connection/tls_prelogin_stream_wrapper.rs

@@ -62,10 +62,12 @@ impl<S> TlsPreloginWrapper<S> {
     }
 
     pub fn start_handshake(&mut self) {
+        log::trace!("Handshake starting");
         self.pending_handshake = true;
     }
 
     pub fn handshake_complete(&mut self) {
+        log::trace!("Handshake complete");
         self.pending_handshake = false;
     }
 }
@@ -159,6 +161,8 @@ impl<S: AsyncRead + AsyncWrite + Unpin + Send> AsyncWrite for TlsPreloginWrapper
             }
 
             while !inner.wr_buf.is_empty() {
+                log::trace!("Writing {} bytes of TLS handshake", inner.wr_buf.len());
+
                 let written = ready!(Pin::new(&mut inner.stream).poll_write(cx, &inner.wr_buf))?;
 
                 inner.wr_buf.drain(..written);

sqlx-core/src/mssql/protocol/packet.rs

@@ -27,18 +27,21 @@ pub(crate) struct PacketHeader {
     pub(crate) packet_id: u8,
 }
 
+impl PacketHeader {
+    fn to_array(&self) -> [u8; PACKET_HEADER_SIZE] {
+        let mut arr = [0u8; PACKET_HEADER_SIZE];
+        arr[0] = self.r#type as u8;
+        arr[1] = self.status.bits();
+        arr[2..4].copy_from_slice(&self.length.to_be_bytes());
+        arr[4..6].copy_from_slice(&self.server_process_id.to_be_bytes());
+        arr[6] = self.packet_id;
+        arr
+    }
+}
+
 impl<'s> Encode<'s, ()> for PacketHeader {
     fn encode_with(&self, buf: &mut Vec<u8>, _: ()) {
-        buf.push(self.r#type as u8);
-        buf.push(self.status.bits());
-
-        buf.extend(&self.length.to_be_bytes());
-
-        buf.extend(&self.server_process_id.to_be_bytes());
-        buf.push(self.packet_id);
-
-        // window, unused
-        buf.push(0);
+        buf.extend_from_slice(&self.to_array());
     }
 }
 

sqlx-core/src/mssql/transaction.rs

@@ -80,6 +80,7 @@ impl TransactionManager for MssqlTransactionManager {
 
             conn.stream.pending_done_count += 1;
 
+            // We cannot flush since we are in a synchronous context, but the packet will always be small here
             conn.stream.write_packet(
                 PacketType::SqlBatch,
                 SqlBatch {

sqlx-core/src/net/tls/rustls.rs

@@ -2,8 +2,8 @@ use rustls::client::danger::{HandshakeSignatureValid, ServerCertVerified, Server
 use rustls::client::WebPkiServerVerifier;
 use rustls::pki_types::{CertificateDer, PrivateKeyDer, ServerName, UnixTime};
 use rustls::{
-    CertificateError, ClientConfig, DigitallySignedStruct, Error as TlsError, RootCertStore,
-    SignatureScheme,
+    CertificateError, ClientConfig, DigitallySignedStruct, Error as TlsError, KeyLogFile,
+    RootCertStore, SignatureScheme,
 };
 use std::io::{BufReader, Cursor};
 use std::sync::Arc;
@@ -50,7 +50,7 @@ pub async fn configure_tls_connector(
     };
 
     // authentication using user's key and its associated certificate
-    let config = match (tls_config.client_cert_path, tls_config.client_key_path) {
+    let mut config = match (tls_config.client_cert_path, tls_config.client_key_path) {
         (Some(cert_path), Some(key_path)) => {
             let cert_chain = certs_from_pem(cert_path.data().await?)?;
             let key_der = private_key_from_pem(key_path.data().await?)?;
@@ -66,6 +66,9 @@ pub async fn configure_tls_connector(
         }
     };
 
+    // When SSLKEYLOGFILE is set, write the TLS keys to a file for use with Wireshark
+    config.key_log = Arc::new(KeyLogFile::new());
+
     Ok(Arc::new(config).into())
 }
 

tests/mssql/mssql.rs

@@ -205,7 +205,7 @@ async fn it_binds_string_with_special_chars() -> anyhow::Result<()> {
 async fn it_accepts_long_query_strings() -> anyhow::Result<()> {
     let mut conn = new::<Mssql>().await?;
     // try a query that does not fit in a single TDS packet
-    let (n,): (i32,) = sqlx_oldapi::query_as(&format!("SELECT {} 42", " ".repeat(0x1_00_00)))
+    let (n,): (i32,) = sqlx_oldapi::query_as(&format!("SELECT {} 42", " ".repeat(3000)))
         .fetch_one(&mut conn)
         .await?;
     assert_eq!(n, 42);