From 4c8968c4c30de69fcd0f75ee9bd3062b2f668170 Mon Sep 17 00:00:00 2001 From: Copilot <223556219+Copilot@users.noreply.github.com> Date: Wed, 1 Jul 2026 20:42:46 -0700 Subject: [PATCH 1/3] Record FedAuth token deterministically in mock server The mock server previously recorded connection tokens into the shared ConnectionStore only at connection teardown, and keyed the store by the client socket address. Both flaws made downstream fedauth regression tests flaky on Linux: tokens were not visible when connect() returned (no barrier, papered over with sleeps), and sequential connects reusing the same ephemeral port overwrote each other under the same key. Record the token eagerly during login, before the LoginAck is sent, by giving ConnectionProcessor a handle to the shared store. Since the client blocks on LoginAck, the token is guaranteed visible once connect() returns. Re-key ConnectionStore by a unique per-connection id from an AtomicU64 counter using a BTreeMap so iteration is ordered and the most recent connection is deterministic. Drop the time.sleep workarounds in the Python fedauth tests. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- mssql-mock-tds/src/server.rs | 110 +++++++++++++++--- .../rs-only-tests/test_mock_server_fedauth.py | 31 +---- 2 files changed, 96 insertions(+), 45 deletions(-) diff --git a/mssql-mock-tds/src/server.rs b/mssql-mock-tds/src/server.rs index 6ba695c7..bfe15db5 100644 --- a/mssql-mock-tds/src/server.rs +++ b/mssql-mock-tds/src/server.rs @@ -13,9 +13,10 @@ use crate::protocol::{ use crate::query_response::QueryRegistry; use bytes::BytesMut; use native_tls::Identity; -use std::collections::HashMap; +use std::collections::BTreeMap; use std::net::SocketAddr; use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::{TcpListener, TcpStream}; use tokio::sync::Mutex; @@ -51,6 +52,8 @@ impl RedirectionConfig { /// Each connection gets its own processor instance. /// FedAuth and username/password authentication are always supported. pub struct ConnectionProcessor { + /// Unique per-connection id assigned at accept time + conn_id: u64, /// Client socket address addr: SocketAddr, /// Whether the client has authenticated @@ -69,12 +72,20 @@ pub struct ConnectionProcessor { buffer: BytesMut, /// Optional redirection configuration redirection: Option, + /// Shared store used to record connection state as soon as it is known + connection_store: Option>>, } impl ConnectionProcessor { /// Create a new connection processor - pub fn new(addr: SocketAddr, query_registry: Arc>) -> Self { + pub fn new( + conn_id: u64, + addr: SocketAddr, + query_registry: Arc>, + connection_store: Option>>, + ) -> Self { Self { + conn_id, addr, is_authenticated: false, received_token: None, @@ -84,16 +95,20 @@ impl ConnectionProcessor { query_registry, buffer: BytesMut::with_capacity(4096), redirection: None, + connection_store, } } /// Create a new connection processor with redirection configuration pub fn new_with_redirection( + conn_id: u64, addr: SocketAddr, query_registry: Arc>, + connection_store: Option>>, redirection: Option, ) -> Self { Self { + conn_id, addr, is_authenticated: false, received_token: None, @@ -103,9 +118,15 @@ impl ConnectionProcessor { query_registry, buffer: BytesMut::with_capacity(4096), redirection, + connection_store, } } + /// Get the unique connection id + pub fn conn_id(&self) -> u64 { + self.conn_id + } + /// Get the client address pub fn addr(&self) -> SocketAddr { self.addr @@ -148,6 +169,15 @@ impl ConnectionProcessor { &mut self.buffer } + /// Upsert this connection's current state into the shared store. + /// Called eagerly during login so tokens are visible to callers the + /// moment the client's blocking LoginAck read returns. + async fn record_to_store(&self) { + if let Some(store) = &self.connection_store { + store.lock().await.store(self); + } + } + /// Process a single packet from the buffer and return the response pub async fn process_packet(&mut self) -> Result, ProtocolError> { if self.buffer.len() < PACKET_HEADER_SIZE { @@ -269,6 +299,7 @@ impl ConnectionProcessor { resp_header.write(&mut packet); packet.extend_from_slice(&response); + self.record_to_store().await; Some(packet) } } @@ -305,6 +336,7 @@ impl ConnectionProcessor { resp_header.write(&mut packet); packet.extend_from_slice(&response); + self.record_to_store().await; Some(packet) } Err(e) => { @@ -418,8 +450,9 @@ impl ConnectionProcessor { /// This allows tests to access per-connection state after connections complete. #[derive(Debug, Default)] pub struct ConnectionStore { - /// Completed connection processors keyed by client socket address - connections: HashMap, + /// Connection info keyed by unique connection id, ordered so that + /// iteration and `.values().last()` yield the most recent connection. + connections: BTreeMap, } /// Captured information from a completed connection @@ -461,11 +494,13 @@ impl ConnectionInfo { impl ConnectionStore { pub fn new() -> Self { Self { - connections: HashMap::new(), + connections: BTreeMap::new(), } } - /// Store connection info when a connection completes + /// Upsert connection info keyed by the connection's unique id. + /// Called both eagerly during login and again at connection teardown; + /// both updates target the same entry. pub fn store(&mut self, processor: &ConnectionProcessor) { let info = ConnectionInfo { addr: processor.addr(), @@ -474,16 +509,16 @@ impl ConnectionStore { user_agent: processor.user_agent.clone(), received_server_name: processor.received_server_name().map(|s| s.to_string()), }; - self.connections.insert(processor.addr(), info); + self.connections.insert(processor.conn_id(), info); } - /// Get connection info by address - pub fn get(&self, addr: &SocketAddr) -> Option<&ConnectionInfo> { - self.connections.get(addr) + /// Get connection info by connection id + pub fn get(&self, conn_id: u64) -> Option<&ConnectionInfo> { + self.connections.get(&conn_id) } /// Get all connection infos - pub fn all(&self) -> &HashMap { + pub fn all(&self) -> &BTreeMap { &self.connections } @@ -513,6 +548,8 @@ pub struct MockTdsServer { connection_store: Arc>, /// Optional redirection configuration for testing client redirection behavior redirection: Option, + /// Monotonic counter assigning a unique id to each accepted connection + connection_counter: Arc, } impl MockTdsServer { @@ -626,6 +663,7 @@ impl MockTdsServer { strict_mode, connection_store: Arc::new(Mutex::new(ConnectionStore::new())), redirection, + connection_counter: Arc::new(AtomicU64::new(0)), }) } @@ -653,11 +691,13 @@ impl MockTdsServer { let strict_mode = self.strict_mode; let connection_store = self.connection_store; let redirection = self.redirection.map(Arc::new); + let connection_counter = self.connection_counter; loop { let (socket, addr) = listener.accept().await?; info!("New connection from {}", addr); + let conn_id = connection_counter.fetch_add(1, Ordering::SeqCst); let registry_clone = Arc::clone(®istry); let tls_acceptor_clone = tls_acceptor.clone(); let store_clone = Arc::clone(&connection_store); @@ -668,6 +708,7 @@ impl MockTdsServer { if let Err(e) = handle_connection_with_tls( socket, addr, + conn_id, registry_clone, tls_acceptor_clone, strict_mode, @@ -693,6 +734,7 @@ impl MockTdsServer { let strict_mode = self.strict_mode; let connection_store = self.connection_store; let redirection = self.redirection.map(Arc::new); + let connection_counter = self.connection_counter; tokio::select! { result = async { @@ -703,13 +745,14 @@ impl MockTdsServer { info!("New connection from {}", addr); drop(listener); // Release lock before spawning + let conn_id = connection_counter.fetch_add(1, Ordering::SeqCst); let registry_clone = Arc::clone(®istry); let tls_acceptor_clone = tls_acceptor.clone(); let store_clone = Arc::clone(&connection_store); let redirection_clone = redirection.clone(); tokio::spawn(async move { - if let Err(e) = handle_connection_with_tls(socket, addr, registry_clone, tls_acceptor_clone, strict_mode, store_clone, redirection_clone).await { + if let Err(e) = handle_connection_with_tls(socket, addr, conn_id, registry_clone, tls_acceptor_clone, strict_mode, store_clone, redirection_clone).await { error!("Error handling connection from {}: {}", addr, e); } }); @@ -731,9 +774,11 @@ impl MockTdsServer { /// Handle a connection with optional TLS support. /// FedAuth and username/password authentication are always supported. +#[allow(clippy::too_many_arguments)] async fn handle_connection_with_tls( socket: TcpStream, addr: SocketAddr, + conn_id: u64, query_registry: Arc>, tls_acceptor: Option>, strict_mode: bool, @@ -764,6 +809,7 @@ async fn handle_connection_with_tls( handle_strict_encrypted_connection( tls_stream, addr, + conn_id, query_registry, connection_store, redirection, @@ -801,6 +847,7 @@ async fn handle_connection_with_tls( handle_encrypted_tds_wrapped_connection( tls_stream, addr, + conn_id, query_registry, connection_store, redirection, @@ -811,6 +858,7 @@ async fn handle_connection_with_tls( handle_unencrypted_connection( prelogin_socket, addr, + conn_id, query_registry, connection_store, redirection, @@ -876,6 +924,7 @@ async fn handle_prelogin_negotiation( async fn handle_strict_encrypted_connection( mut socket: TlsStream, addr: SocketAddr, + conn_id: u64, query_registry: Arc>, connection_store: Arc>, redirection: Option>, @@ -883,8 +932,13 @@ async fn handle_strict_encrypted_connection( let redir_config = redirection .as_ref() .map(|r| RedirectionConfig::new(r.redirect_host.clone(), r.redirect_port)); - let mut processor = - ConnectionProcessor::new_with_redirection(addr, query_registry, redir_config); + let mut processor = ConnectionProcessor::new_with_redirection( + conn_id, + addr, + query_registry, + Some(Arc::clone(&connection_store)), + redir_config, + ); let mut prelogin_handled = false; loop { @@ -952,10 +1006,16 @@ async fn handle_strict_encrypted_connection( async fn handle_encrypted_connection( mut socket: TlsStream, addr: SocketAddr, + conn_id: u64, query_registry: Arc>, connection_store: Arc>, ) -> Result<(), ProtocolError> { - let mut processor = ConnectionProcessor::new(addr, query_registry); + let mut processor = ConnectionProcessor::new( + conn_id, + addr, + query_registry, + Some(Arc::clone(&connection_store)), + ); loop { // Read data from TLS socket @@ -989,6 +1049,7 @@ async fn handle_encrypted_connection( async fn handle_encrypted_tds_wrapped_connection( mut socket: TlsStream, addr: SocketAddr, + conn_id: u64, query_registry: Arc>, connection_store: Arc>, redirection: Option>, @@ -996,8 +1057,13 @@ async fn handle_encrypted_tds_wrapped_connection( let redir_config = redirection .as_ref() .map(|r| RedirectionConfig::new(r.redirect_host.clone(), r.redirect_port)); - let mut processor = - ConnectionProcessor::new_with_redirection(addr, query_registry, redir_config); + let mut processor = ConnectionProcessor::new_with_redirection( + conn_id, + addr, + query_registry, + Some(Arc::clone(&connection_store)), + redir_config, + ); loop { // Read data from TLS socket (which wraps TdsTlsWrapper) @@ -1032,6 +1098,7 @@ async fn handle_encrypted_tds_wrapped_connection( async fn handle_unencrypted_connection( mut socket: TcpStream, addr: SocketAddr, + conn_id: u64, query_registry: Arc>, connection_store: Arc>, redirection: Option>, @@ -1039,8 +1106,13 @@ async fn handle_unencrypted_connection( let redir_config = redirection .as_ref() .map(|r| RedirectionConfig::new(r.redirect_host.clone(), r.redirect_port)); - let mut processor = - ConnectionProcessor::new_with_redirection(addr, query_registry, redir_config); + let mut processor = ConnectionProcessor::new_with_redirection( + conn_id, + addr, + query_registry, + Some(Arc::clone(&connection_store)), + redir_config, + ); loop { // Read data from plain socket diff --git a/mssql-py-core/tests/rs-only-tests/test_mock_server_fedauth.py b/mssql-py-core/tests/rs-only-tests/test_mock_server_fedauth.py index 21743b0b..3e904c20 100644 --- a/mssql-py-core/tests/rs-only-tests/test_mock_server_fedauth.py +++ b/mssql-py-core/tests/rs-only-tests/test_mock_server_fedauth.py @@ -98,11 +98,7 @@ def test_connect_with_access_token(self, mock_server_port): # Clean up conn.close() assert not conn.is_connected() - - # Give the server a moment to process the connection info - import time - time.sleep(0.1) - + # Verify the server received the correct token assert server.connection_count() >= 1, "Server should have recorded at least one connection" assert server.has_received_token(mock_token), \ @@ -121,7 +117,6 @@ def test_connect_with_unique_access_token(self, mock_server_port): properly sent through the TDS protocol and received by the server. """ import mssql_py_core - import time # Generate a unique token for this test unique_token = f"unique_token_{secrets.token_hex(16)}" @@ -141,10 +136,7 @@ def test_connect_with_unique_access_token(self, mock_server_port): assert conn is not None assert conn.is_connected() conn.close() - - # Wait for connection info to be stored - time.sleep(0.1) - + # Verify the unique token was received received_token = server.get_last_access_token() assert received_token == unique_token, \ @@ -174,7 +166,6 @@ def test_execute_query_with_access_token(self, mock_server_port): Verifies both the query result and that the token was received. """ import mssql_py_core - import time mock_token = "mock_token_for_query_execution" @@ -209,10 +200,7 @@ def test_execute_query_with_access_token(self, mock_server_port): # Ensure references are dropped so TdsClient is fully released del cursor del conn - - # Wait for connection info to be stored - time.sleep(0.3) - + # Verify token was received assert server.has_received_token(mock_token), \ "Server should have received the access token used for query execution" @@ -220,7 +208,6 @@ def test_execute_query_with_access_token(self, mock_server_port): def test_get_all_connections(self, mock_server_port): """Test retrieving all connection info from the server.""" import mssql_py_core - import time token = "test_token_for_connection_list" @@ -237,9 +224,7 @@ def test_get_all_connections(self, mock_server_port): conn = mssql_py_core.PyCoreConnection(client_context) conn.close() - - time.sleep(0.1) - + # Get all connections connections = server.get_connections() assert len(connections) >= 1, "Should have at least one connection" @@ -252,7 +237,6 @@ def test_get_all_connections(self, mock_server_port): def test_clear_connections(self, mock_server_port): """Test clearing stored connection info.""" import mssql_py_core - import time server = mssql_mock_tds_py.PyMockTdsServer(port=mock_server_port, tls=True) @@ -267,9 +251,7 @@ def test_clear_connections(self, mock_server_port): conn = mssql_py_core.PyCoreConnection(client_context) conn.close() - - time.sleep(0.1) - + # Verify we have a connection assert server.connection_count() >= 1 @@ -282,7 +264,6 @@ def test_clear_connections(self, mock_server_port): def test_user_agent_format(self, mock_server_port): """Test that MS-PYTHON is correctly sent as the driver name in the user agent.""" import mssql_py_core - import time server = mssql_mock_tds_py.PyMockTdsServer(port=mock_server_port, tls=True) @@ -297,8 +278,6 @@ def test_user_agent_format(self, mock_server_port): conn = mssql_py_core.PyCoreConnection(client_context) conn.close() - time.sleep(0.1) - connections = server.get_connections() assert len(connections) >= 1, "Should have at least one connection" From 35dfca6d6ce1bfc9140ded8e9219ab72c1eb329b Mon Sep 17 00:00:00 2001 From: Saurabh Singh <1623701+saurabh500@users.noreply.github.com> Date: Thu, 2 Jul 2026 10:45:10 -0700 Subject: [PATCH 2/3] Add regression test for eager FedAuth token recording Assert the FedAuth token is visible in the shared ConnectionStore while the client connection is still open, guarding against a regression to teardown-only recording. Mirrors the downstream ODBC pooled-connection scenario where close() does not send EOF. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- mssql-tds/tests/test_mock_server_fedauth.rs | 59 +++++++++++++++++++++ 1 file changed, 59 insertions(+) diff --git a/mssql-tds/tests/test_mock_server_fedauth.rs b/mssql-tds/tests/test_mock_server_fedauth.rs index 108e05fa..7d439723 100644 --- a/mssql-tds/tests/test_mock_server_fedauth.rs +++ b/mssql-tds/tests/test_mock_server_fedauth.rs @@ -469,4 +469,63 @@ mod mock_server_fedauth_tests { Ok(()) } + + /// Regression guard for eager token recording: the token must be visible in the + /// shared store WHILE the client connection is still open, i.e. without relying on + /// connection teardown. This mirrors the downstream ODBC pooled-connection scenario + /// where close() returns the socket to a pool without sending EOF, so a teardown-only + /// record would never fire. + #[tokio::test] + async fn test_token_recorded_before_connection_close() -> Result<(), Box> + { + init_tracing(); + + let server = MockTdsServer::new("127.0.0.1:0").await?; + let server_addr = server.local_addr(); + let connection_store = server.connection_store(); + + let (shutdown_tx, shutdown_rx) = oneshot::channel(); + let server_handle = + tokio::spawn(async move { server.run_with_shutdown(shutdown_rx).await }); + + tokio::time::sleep(tokio::time::Duration::from_millis(100)).await; + + let access_token = "eager_record_guard_token_67890".to_string(); + + let datasource = format!("tcp:{},{}", server_addr.ip(), server_addr.port()); + let mut context = ClientContext::default(); + context.access_token = Some(access_token.clone()); + context.tds_authentication_method = TdsAuthenticationMethod::AccessToken; + context.database = "master".to_string(); + context.encryption_options = EncryptionOptions { + mode: EncryptionSetting::PreferOff, + trust_server_certificate: true, + host_name_in_cert: None, + server_certificate: None, + }; + + let provider = TdsConnectionProvider {}; + let client = provider.create_client(context, &datasource, None).await?; + + // Assert the token is recorded WITHOUT closing/dropping the connection first. + let recorded = { + let store = connection_store.lock().await; + store + .all() + .values() + .any(|c| c.received_token_as_string().as_deref() == Some(&access_token)) + }; + assert!( + recorded, + "FedAuth token must be recorded before the client connection is closed (eager-record guard)" + ); + + // Keep the connection alive until after the assertion. + drop(client); + + let _ = shutdown_tx.send(()); + let _ = tokio::time::timeout(tokio::time::Duration::from_secs(2), server_handle).await; + + Ok(()) + } } From 1b794057506ae2db00fd1e132204553983088728 Mon Sep 17 00:00:00 2001 From: Saurabh Singh <1623701+saurabh500@users.noreply.github.com> Date: Mon, 6 Jul 2026 13:16:30 -0700 Subject: [PATCH 3/3] Narrow ConnectionProcessor::conn_id visibility to pub(crate) Only used internally by ConnectionStore::store; keep the crate's public surface deliberate per repo conventions. Addresses PR review feedback. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- mssql-mock-tds/src/server.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mssql-mock-tds/src/server.rs b/mssql-mock-tds/src/server.rs index bfe15db5..4f3f4339 100644 --- a/mssql-mock-tds/src/server.rs +++ b/mssql-mock-tds/src/server.rs @@ -123,7 +123,7 @@ impl ConnectionProcessor { } /// Get the unique connection id - pub fn conn_id(&self) -> u64 { + pub(crate) fn conn_id(&self) -> u64 { self.conn_id }