From dd774e5b3839ed5dbf9a4b9244c30744f01d212f Mon Sep 17 00:00:00 2001 From: gargsaumya Date: Wed, 12 Aug 2026 08:27:49 +0000 Subject: [PATCH] Seed session-recovery reconnect baseline from FEATUREEXTACK The server delivers the recoverable session-state baseline in the SESSIONRECOVERY FEATUREEXTACK at login, but the driver captured that payload and dropped it instead of parsing it. The reconnect LOGIN7 therefore omitted those state entries and the server rejected it (error 17897, state 81: "session recovery feature data ... structurally or semantically invalid"), so transparent reconnect never succeeded. Parse the FEATUREEXTACK baseline (StateId, Length, Value entries) into initial_state so it is echoed on reconnect, alongside the initial database/language/collation captured at login. Also fix transparent_reconnect_after_kill: it asserted the SPID changes after recovery, but SQL Server recycles SPID numbers, so a recovered session often reuses the same one; rely on connection_recovery_count instead. Re-enable the previously skipped C++ e2e session recovery test. --- .../tests/e2e/tests/session_recovery_test.cpp | 5 - mssql-tds/src/connection/session_recovery.rs | 191 +++++++++++++++++- mssql-tds/src/connection/tds_client.rs | 20 +- .../tds_connection_provider.rs | 32 ++- mssql-tds/src/fuzz_support.rs | 1 + mssql-tds/src/handler/handler_factory.rs | 57 +++++- .../src/message/features/session_recovery.rs | 9 + mssql-tds/src/message/login.rs | 6 + mssql-tds/src/test_client_support.rs | 1 + mssql-tds/tests/connectivity.rs | 13 +- 10 files changed, 309 insertions(+), 26 deletions(-) diff --git a/mssql-odbc/tests/e2e/tests/session_recovery_test.cpp b/mssql-odbc/tests/e2e/tests/session_recovery_test.cpp index 14d31a7e..dbf0abf1 100644 --- a/mssql-odbc/tests/e2e/tests/session_recovery_test.cpp +++ b/mssql-odbc/tests/e2e/tests/session_recovery_test.cpp @@ -22,11 +22,6 @@ class SessionRecoveryLiveTest : public ODBCTest { protected: void SetUp() override { - GTEST_SKIP() << "Disabled pending mssql-tds transparent-reconnect fix: " - "the reconnection LOGIN7 is rejected by the server, so " - "SQLExecute-after-KILL cannot recover yet. Re-enable when " - "transparent reconnect succeeds end-to-end."; - ODBCTest::SetUp(); ASSERT_TRUE(ODBCTestConfig::Instance().HasConnection()) << "No connection configured – set ODBC_TEST_SERVER or " diff --git a/mssql-tds/src/connection/session_recovery.rs b/mssql-tds/src/connection/session_recovery.rs index b5315be1..b50a0ab0 100644 --- a/mssql-tds/src/connection/session_recovery.rs +++ b/mssql-tds/src/connection/session_recovery.rs @@ -67,7 +67,9 @@ impl RecoveryContext { /// Initialize recovery context with connection-time settings. /// Called after a successful login to capture the original connection parameters - /// needed for reconnection validation and orchestration. + /// needed for reconnection validation and orchestration, and to seed the + /// baseline session state a reconnect's LOGIN7 must replay. + #[allow(clippy::too_many_arguments)] pub fn initialize( &mut self, client_context: ClientContext, @@ -75,12 +77,26 @@ impl RecoveryContext { server_version: Option, encryption_level: NegotiatedEncryptionSetting, mars_enabled: bool, + session_recovery_negotiated: bool, + initial_database: String, + initial_language: String, + initial_collation: SqlCollation, + initial_state_tokens: &[SessionStateToken], + initial_state_ack: Option<&[u8]>, ) { self.client_context = Some(Box::new(client_context)); self.original_tds_version = tds_version; self.original_server_version = server_version; self.original_encryption_level = Some(encryption_level); self.original_mars_enabled = mars_enabled; + self.session_recovery_negotiated = session_recovery_negotiated; + self.session_state_table.seed_initial_state( + initial_database, + initial_language, + initial_collation, + initial_state_tokens, + initial_state_ack, + ); } /// Check whether session recovery can be attempted. @@ -285,6 +301,82 @@ impl SessionStateTable { }); } + /// Seed the baseline snapshot captured at login. + /// + /// Called once, right after a successful login, with the database/language/ + /// collation negotiated at connect time and any SESSIONSTATE tokens the + /// server sent during login. Without this, `initial_*` stays empty and a + /// reconnect's LOGIN7 session-recovery block goes out blank, which servers + /// reject. + pub fn seed_initial_state( + &mut self, + database: String, + language: String, + collation: SqlCollation, + tokens: &[SessionStateToken], + feature_ack_initial_state: Option<&[u8]>, + ) { + self.initial_database = database; + self.initial_language = language; + self.initial_collation = collation; + + for token in tokens { + if token.sequence_number == u32::MAX { + self.master_recovery_disabled = true; + continue; + } + for entry in &token.states { + if !entry.recoverable { + self.unrecoverable_state_count += 1; + } + self.initial_state[entry.state_id as usize] = Some(SessionStateRecord { + recoverable: entry.recoverable, + sequence: token.sequence_number, + data: entry.data.clone(), + }); + } + } + + // The server delivers the recoverable baseline as a packed FEATUREEXTACK + // payload, not SESSIONSTATE tokens. Each entry must be echoed in the + // reconnect LOGIN7 or the server rejects it as semantically invalid + // (error 17897, state 81). + if let Some(data) = feature_ack_initial_state { + self.seed_initial_state_from_feature_ack(data); + } + } + + /// Parse the packed FEATUREEXTACK session-recovery baseline into + /// `initial_state`. Entry format matches the reconnect wire format: + /// StateId (1 byte), Length (1 byte, or `0xFF` followed by a u32), Value. + fn seed_initial_state_from_feature_ack(&mut self, data: &[u8]) { + let mut i = 0; + while i < data.len() { + let state_id = data[i] as usize; + i += 1; + let Some(&len_byte) = data.get(i) else { break }; + i += 1; + let len = if len_byte == 0xFF { + let Some(bytes) = data.get(i..i + 4) else { + break; + }; + i += 4; + u32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]) as usize + } else { + len_byte as usize + }; + let Some(value) = data.get(i..i + len) else { + break; + }; + i += len; + self.initial_state[state_id] = Some(SessionStateRecord { + recoverable: true, + sequence: 0, + data: value.to_vec(), + }); + } + } + /// Returns `true` if the session can be recovered after a disconnect. /// /// Recovery is blocked when the server has globally disabled it or when @@ -537,6 +629,96 @@ mod tests { assert!(table.delta[5].is_none()); } + #[test] + fn seed_from_feature_ack_parses_server_baseline() { + // The 8-entry baseline the server sent in the SESSIONRECOVERY + // FEATUREEXTACK at login (captured from the wire). Dropping any of + // these makes the reconnect LOGIN7 fail with error 17897, state 81. + let ack = vec![ + 0x00, 0x09, 0x00, 0x60, 0x81, 0x14, 0xFF, 0xE7, 0xFF, 0xFF, 0x00, // id 0, len 9 + 0x02, 0x02, 0x07, 0x01, // id 2, len 2 + 0x04, 0x01, 0x00, // id 4, len 1 + 0x05, 0x04, 0xFF, 0xFF, 0xFF, 0xFF, // id 5, len 4 + 0x06, 0x01, 0x00, // id 6, len 1 + 0x07, 0x01, 0x02, // id 7, len 1 + 0x08, 0x08, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // id 8, len 8 + 0x09, 0x04, 0xFF, 0xFF, 0xFF, 0xFF, // id 9, len 4 + ]; + let mut table = SessionStateTable::new(); + table.seed_initial_state_from_feature_ack(&ack); + + for (id, expected) in [ + ( + 0usize, + vec![0x00, 0x60, 0x81, 0x14, 0xFF, 0xE7, 0xFF, 0xFF, 0x00], + ), + (2, vec![0x07, 0x01]), + (4, vec![0x00]), + (5, vec![0xFF, 0xFF, 0xFF, 0xFF]), + (6, vec![0x00]), + (7, vec![0x02]), + (8, vec![0x00; 8]), + (9, vec![0xFF, 0xFF, 0xFF, 0xFF]), + ] { + let record = table.initial_state[id] + .as_ref() + .unwrap_or_else(|| panic!("state {id} missing")); + assert_eq!(record.data, expected); + assert!(record.recoverable); + assert_eq!(record.sequence, 0); + } + // Gaps stay empty and the session remains recoverable. + assert!(table.initial_state[1].is_none()); + assert!(table.initial_state[3].is_none()); + assert!(table.is_session_recoverable()); + } + + #[test] + fn seed_from_feature_ack_extended_length() { + // Values >= 0xFF use the 0xFF marker followed by a u32 length. + let mut ack = vec![10u8, 0xFF]; + ack.extend_from_slice(&300u32.to_le_bytes()); + ack.extend_from_slice(&[0xAB; 300]); + + let mut table = SessionStateTable::new(); + table.seed_initial_state_from_feature_ack(&ack); + + let record = table.initial_state[10].as_ref().unwrap(); + assert_eq!(record.data.len(), 300); + assert!(record.data.iter().all(|&b| b == 0xAB)); + } + + #[test] + fn seed_from_feature_ack_truncated_is_safe() { + // A complete entry followed by one whose declared length runs past the + // buffer: the good entry is kept, the truncated one dropped, no panic. + let ack = vec![3, 1, 0xAA, 5, 10, 0x01]; + let mut table = SessionStateTable::new(); + table.seed_initial_state_from_feature_ack(&ack); + + assert_eq!(table.initial_state[3].as_ref().unwrap().data, vec![0xAA]); + assert!(table.initial_state[5].is_none()); + } + + #[test] + fn seed_initial_state_forwards_feature_ack() { + let ack = vec![7, 2, 0x01, 0x02]; + let mut table = SessionStateTable::new(); + table.seed_initial_state( + "master".to_string(), + "us_english".to_string(), + SqlCollation::default(), + &[], + Some(&ack), + ); + + assert_eq!(table.initial_database, "master"); + assert_eq!( + table.initial_state[7].as_ref().unwrap().data, + vec![0x01, 0x02] + ); + } + #[test] fn reset_preserves_master_recovery_disabled() { let mut table = SessionStateTable::new(); @@ -707,8 +889,13 @@ mod tests { Some(Version::new(16, 0, 1000, 0)), NegotiatedEncryptionSetting::Mandatory, false, + true, + "master".to_string(), + "us_english".to_string(), + SqlCollation::default(), + &[], + None, ); - ctx.session_recovery_negotiated = true; ctx } diff --git a/mssql-tds/src/connection/tds_client.rs b/mssql-tds/src/connection/tds_client.rs index 48dab2c6..cc201f20 100644 --- a/mssql-tds/src/connection/tds_client.rs +++ b/mssql-tds/src/connection/tds_client.rs @@ -23,7 +23,7 @@ use crate::message::transaction_management::{ TransactionManagementType, }; use crate::query::result::ReturnValue; -use crate::token::tokens::SqlCollation; +use crate::token::tokens::{SessionStateToken, SqlCollation}; use crate::{ connection::{ execution_context::{ALREADY_EXECUTING_ERROR, ExecutionContext}, @@ -183,6 +183,7 @@ impl TdsClient { negotiated_settings: NegotiatedSettings, execution_context: ExecutionContext, client_context: ClientContext, + login_session_state_tokens: Vec, ) -> Self { let mut recovery_context = RecoveryContext::new(); recovery_context.initialize( @@ -193,6 +194,12 @@ impl TdsClient { .session_settings .negotiated_encryption_settings, negotiated_settings.session_settings.mars_enabled, + negotiated_settings.is_session_recovery_acknowledged(), + negotiated_settings.database.clone(), + negotiated_settings.language.clone(), + negotiated_settings.database_collation, + &login_session_state_tokens, + negotiated_settings.session_recovery_initial_state(), ); Self { @@ -323,7 +330,13 @@ impl TdsClient { ) .await; match connect_result { - Ok((new_transport, new_settings, new_exec_ctx, info_messages)) => { + Ok(( + new_transport, + new_settings, + new_exec_ctx, + info_messages, + _session_state_tokens, + )) => { // Validate reconnection properties match original if let Err(validation_err) = self.recovery_context.validate_reconnection(&new_settings) @@ -4421,6 +4434,7 @@ mod tests { negotiated_settings, execution_context, client_context, + Vec::new(), ) } @@ -4465,6 +4479,7 @@ mod tests { negotiated_settings, execution_context, client_context, + Vec::new(), ) } @@ -4482,6 +4497,7 @@ mod tests { negotiated_settings, execution_context, client_context, + Vec::new(), ); (client, sent) } diff --git a/mssql-tds/src/connection_provider/tds_connection_provider.rs b/mssql-tds/src/connection_provider/tds_connection_provider.rs index 3025fe47..563fb034 100644 --- a/mssql-tds/src/connection_provider/tds_connection_provider.rs +++ b/mssql-tds/src/connection_provider/tds_connection_provider.rs @@ -22,6 +22,7 @@ use crate::error::{Error, SqlInfoMessage, TimeoutErrorType}; use crate::handler::handler_factory::HandlerFactory; use crate::io::token_stream::GenericTokenParserRegistry; use crate::ssrp; +use crate::token::tokens::SessionStateToken; #[cfg(fuzzing)] use crate::io::token_stream::TdsTokenStreamReader; @@ -60,9 +61,20 @@ impl TdsConnectionProvider { + crate::io::packet_reader::TdsPacketReader + 'static, { - let (transport, negotiated_settings, execution_context, info_messages) = - Self::connect_with_transport(&context, &context.transport_context, transport).await?; - let mut client = TdsClient::new(transport, negotiated_settings, execution_context, context); + let ( + transport, + negotiated_settings, + execution_context, + info_messages, + session_state_tokens, + ) = Self::connect_with_transport(&context, &context.transport_context, transport).await?; + let mut client = TdsClient::new( + transport, + negotiated_settings, + execution_context, + context, + session_state_tokens, + ); client.extend_info_messages(info_messages); Ok(client) } @@ -177,13 +189,14 @@ impl TdsConnectionProvider { None => connect_future.await, }; match sm_result { - Ok((transport, negotiated_settings, execution_context, info_messages)) => { + Ok((transport, negotiated_settings, execution_context, info_messages, session_state_tokens)) => { debug!("Shared Memory connection succeeded, skipping SSRP"); let mut client = TdsClient::new( transport, negotiated_settings, execution_context, context.clone(), + session_state_tokens, ); client.extend_info_messages(info_messages); return Ok(client); @@ -371,13 +384,14 @@ impl TdsConnectionProvider { // Handle redirections loop { match connection_result { - Ok((transport, negotiated_settings, execution_context, info_messages)) => { + Ok((transport, negotiated_settings, execution_context, info_messages, session_state_tokens)) => { debug!("Connection successful via action chain"); let mut client = TdsClient::new( transport, negotiated_settings, execution_context, context.clone(), + session_state_tokens, ); client.extend_info_messages(info_messages); return Ok(client); @@ -544,6 +558,7 @@ impl TdsConnectionProvider { crate::handler::handler_factory::NegotiatedSettings, crate::connection::execution_context::ExecutionContext, Vec, + Vec, )> { // Create network transport directly // Convert connect_timeout from seconds to milliseconds @@ -570,7 +585,7 @@ impl TdsConnectionProvider { .await; match session_result { - Ok((negotiated_settings, info_messages)) => { + Ok((negotiated_settings, info_messages, session_state_tokens)) => { // Create execution context for the new connection let execution_context = crate::connection::execution_context::ExecutionContext::new(); @@ -580,6 +595,7 @@ impl TdsConnectionProvider { negotiated_settings, execution_context, info_messages, + session_state_tokens, )) } Err(err) => { @@ -603,6 +619,7 @@ impl TdsConnectionProvider { crate::handler::handler_factory::NegotiatedSettings, crate::connection::execution_context::ExecutionContext, Vec, + Vec, )> where T: TdsTransport @@ -621,7 +638,7 @@ impl TdsConnectionProvider { .await; match session_result { - Ok((negotiated_settings, info_messages)) => { + Ok((negotiated_settings, info_messages, session_state_tokens)) => { // Create execution context for the new connection let execution_context = crate::connection::execution_context::ExecutionContext::new(); @@ -631,6 +648,7 @@ impl TdsConnectionProvider { negotiated_settings, execution_context, info_messages, + session_state_tokens, )) } Err(err) => { diff --git a/mssql-tds/src/fuzz_support.rs b/mssql-tds/src/fuzz_support.rs index 4c1d4cec..48be06b2 100644 --- a/mssql-tds/src/fuzz_support.rs +++ b/mssql-tds/src/fuzz_support.rs @@ -756,5 +756,6 @@ pub fn create_fuzz_tds_client( negotiated_settings, execution_context, client_context, + Vec::new(), ) } diff --git a/mssql-tds/src/handler/handler_factory.rs b/mssql-tds/src/handler/handler_factory.rs index 080a067a..18b30827 100644 --- a/mssql-tds/src/handler/handler_factory.rs +++ b/mssql-tds/src/handler/handler_factory.rs @@ -20,7 +20,7 @@ use crate::message::prelogin::{ EncryptionType, PreloginRequest, PreloginRequestModel, PreloginResponse, }; use crate::token::login_ack::LoginAckToken; -use crate::token::tokens::SqlCollation; +use crate::token::tokens::{SessionStateToken, SqlCollation}; use tracing::{debug, warn}; use uuid::Uuid; @@ -149,6 +149,16 @@ impl NegotiatedSettings { .any(|f| f.feature_identifier() == FeatureExtension::SRecovery && f.is_acknowledged()) } + /// Raw FEATUREEXTACK session-recovery baseline the server sent at login, + /// which must be echoed back in the reconnect LOGIN7's initial state block. + pub(crate) fn session_recovery_initial_state(&self) -> Option<&[u8]> { + self.session_settings + .supported_features + .iter() + .find(|f| f.feature_identifier() == FeatureExtension::SRecovery && f.is_acknowledged()) + .and_then(|f| f.session_recovery_initial_state()) + } + /// Check if Always Encrypted (column encryption) was acknowledged by the /// server in FEATUREEXTACK. pub(crate) fn is_column_encryption_supported(&self) -> bool { @@ -236,7 +246,11 @@ impl<'a, 'b> SessionHandler<'a, 'b> { pub(crate) async fn execute( &mut self, reader_writer: &mut T, - ) -> TdsResult<(NegotiatedSettings, Vec)> { + ) -> TdsResult<( + NegotiatedSettings, + Vec, + Vec, + )> { let pre_login_result = self.get_pre_login_result(reader_writer).await?; self.validate_prelogin_result(&pre_login_result)?; @@ -252,8 +266,9 @@ impl<'a, 'b> SessionHandler<'a, 'b> { let negotiated_settings = self.infer_negotiated_settings(&pre_login_result, &mut login_result)?; let info_messages = std::mem::take(&mut login_result.diagnostics.info_messages); + let session_state_tokens = std::mem::take(&mut login_result.session_state_tokens); reader_writer.notify_session_setting_change(&negotiated_settings.session_settings); - Ok((negotiated_settings, info_messages)) + Ok((negotiated_settings, info_messages, session_state_tokens)) } async fn get_pre_login_result( @@ -329,10 +344,15 @@ impl<'a, 'b> SessionHandler<'a, 'b> { None => (None, None), }; + let language = change_props + .language + .clone() + .unwrap_or_else(|| self.factory.context.language.clone()); + Ok(NegotiatedSettings::new( session_settings, database_collation, - "".to_string(), + language, database, change_props.char_set.clone(), login_ack_tds_version, @@ -453,6 +473,9 @@ struct LoginResult { /// ERROR tokens plus any INFO/warning messages. On success only the /// informational messages are populated; on failure the errors explain why. diagnostics: SqlServerDiagnostics, + /// SESSIONSTATE tokens received during login — the baseline snapshot a + /// reconnect's LOGIN7 session-recovery block must replay. + session_state_tokens: Vec, } pub struct LoginHandler<'a> { @@ -594,12 +617,15 @@ impl LoginHandler<'_> { .map(|f| f.clone_box()) .collect(); + let session_state_tokens = std::mem::take(&mut login_response.session_state_tokens); + Ok(LoginResult { supported_features, change_properties: login_response.change_properties, status: response_status, login_ack: login_response.success_token, diagnostics: SqlServerDiagnostics::new(errors, info_messages), + session_state_tokens, }) } @@ -760,4 +786,27 @@ mod tests { let settings = create_test_negotiated_settings_internal(); assert!(!settings.is_session_recovery_acknowledged()); } + + #[test] + fn session_recovery_initial_state_returns_ack_payload() { + let mut settings = create_test_negotiated_settings_internal(); + let mut feature = SessionRecoveryFeature::new(1); + feature.set_acknowledged(true); + feature.deserialize(&[0x07, 0x02, 0x01, 0x02]).unwrap(); + settings + .session_settings + .supported_features + .push(Box::new(feature)); + + assert_eq!( + settings.session_recovery_initial_state(), + Some([0x07, 0x02, 0x01, 0x02].as_slice()) + ); + } + + #[test] + fn session_recovery_initial_state_none_without_feature() { + let settings = create_test_negotiated_settings_internal(); + assert!(settings.session_recovery_initial_state().is_none()); + } } diff --git a/mssql-tds/src/message/features/session_recovery.rs b/mssql-tds/src/message/features/session_recovery.rs index e951892d..6b765d50 100644 --- a/mssql-tds/src/message/features/session_recovery.rs +++ b/mssql-tds/src/message/features/session_recovery.rs @@ -189,6 +189,10 @@ impl Feature for SessionRecoveryFeature { Ok(()) } + fn session_recovery_initial_state(&self) -> Option<&[u8]> { + self.initial_state_data.as_deref() + } + fn is_acknowledged(&self) -> bool { self.acknowledged } @@ -235,10 +239,15 @@ mod tests { fn deserialize_stores_initial_state_data() { let mut feature = SessionRecoveryFeature::new(1); assert!(feature.initial_state_data.is_none()); + assert!(feature.session_recovery_initial_state().is_none()); let data = vec![0x01, 0x02, 0x03]; feature.deserialize(&data).unwrap(); assert_eq!(feature.initial_state_data.as_ref().unwrap(), &data); + assert_eq!( + feature.session_recovery_initial_state(), + Some(data.as_slice()) + ); } #[test] diff --git a/mssql-tds/src/message/login.rs b/mssql-tds/src/message/login.rs index d358c1e9..6576c512 100644 --- a/mssql-tds/src/message/login.rs +++ b/mssql-tds/src/message/login.rs @@ -117,6 +117,12 @@ pub(crate) trait Feature: Send + Sync + Debug { async fn serialize(&self, packet_writer: &mut PacketWriter) -> TdsResult<()>; fn deserialize(&mut self, data: &[u8]) -> TdsResult<()>; + /// Raw FEATUREEXTACK initial session-state payload, if this feature carries + /// one. Only SessionRecovery does; it must be echoed in the reconnect LOGIN7. + fn session_recovery_initial_state(&self) -> Option<&[u8]> { + None + } + #[allow(dead_code)] // This method is not used currently, and exists for completeness. fn is_acknowledged(&self) -> bool; diff --git a/mssql-tds/src/test_client_support.rs b/mssql-tds/src/test_client_support.rs index 62072052..b3278ea0 100644 --- a/mssql-tds/src/test_client_support.rs +++ b/mssql-tds/src/test_client_support.rs @@ -187,6 +187,7 @@ pub fn tds_client_from_tokens(tokens: Vec) -> TdsClient { negotiated_settings, execution_context, client_context, + Vec::new(), ) } diff --git a/mssql-tds/tests/connectivity.rs b/mssql-tds/tests/connectivity.rs index fa0786f8..e32e3d24 100644 --- a/mssql-tds/tests/connectivity.rs +++ b/mssql-tds/tests/connectivity.rs @@ -436,13 +436,14 @@ mod connectivity { tokio::time::sleep(std::time::Duration::from_millis(500)).await; - let new_spid = get_spid(&mut client).await?; + // The query succeeding after the KILL proves the connection was + // transparently re-established (it would error otherwise). The SPID + // is deliberately NOT asserted to change: SQL Server recycles SPID + // numbers, so a recovered session frequently reuses the same one. + // connection_recovery_count is the reliable signal that a reconnect + // actually occurred. + get_spid(&mut client).await?; - assert_ne!( - original_spid, new_spid, - "SPID should change after reconnection (was {}, now {})", - original_spid, new_spid - ); assert_eq!( client.connection_recovery_count(), 1,