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,