Record FedAuth token deterministically in mock server

[Saurabh Singh (saurabh500)]

Summary

Fixes two design flaws in mssql-mock-tds that made downstream FedAuth regression tests (microsoft/mssql-python#652: test_unique_access_token_transmitted_exactly, test_distinct_tokens_on_sequential_connects) flaky/failing only on Linux.

Targets main directly and contains only the mock-server token-capture changes (this supersedes #89, which was stacked on #80).

The two flaws

1. Token recorded only at connection teardown, asynchronously. Each connection is handled in a spawned task; the token was captured into the per-connection ConnectionProcessor during login but only copied into the shared, queryable ConnectionStore after the read loop broke on client EOF. That insert ran on the server task after the client's close() had already returned, so tests had no barrier to wait on and papered over it with time.sleep.
2. Store keyed by client socket address. ConnectionStore was HashMap<SocketAddr, ConnectionInfo>. Two sequential connects from the same client process frequently reuse the same ephemeral local port on Linux, so the second connection's insert overwrote the first under the same key — has_received_token(first) then returned False.

The fix

• Record eagerly, before LoginAck. ConnectionProcessor now holds an Option<Arc<Mutex<ConnectionStore>>> and upserts its state into the shared store in both FedAuth paths (inline-token Login7 and challenged FedAuthToken) before the response bytes are returned. Because the client blocks on the LoginAck read, the token is guaranteed visible the moment connect() returns — no teardown dependency, no sleeps. The existing teardown store.store() calls remain as a harmless final upsert.
• Re-key by a unique per-connection id. An AtomicU64 counter on the server assigns a conn_id at accept() time, threaded into the ConnectionProcessor. ConnectionStore now uses BTreeMap<u64, ConnectionInfo> so iteration and .values().last() are ordered and get_last_access_token() deterministically returns the most recent connection. store() upserts by conn_id. ConnectionInfo.addr is retained for info; ConnectionStore::get now takes a conn_id.
• Removed the time.sleep workarounds in mssql-py-core/tests/rs-only-tests/test_mock_server_fedauth.py; conn.close() calls kept.
• Added a regression test test_token_recorded_before_connection_close that asserts the token is present in the store while the client connection is still open (guards against a regression to teardown-only recording; mirrors the ODBC pooled-connection scenario where close() returns the socket to a pool without EOF).

Validation

• cargo nextest run -p mssql-tds --test test_mock_server_fedauth — 6/6 pass (incl. the new guard).
• cargo clippy -p mssql-tds --tests --all-features -- -D warnings — clean.
• rustfmt applied.

Note: The in-repo Python fedauth tests could not be run to green on Windows due to a pre-existing, Windows-only native-tls/schannel handshake limitation with the self-signed test cert (connect times out during the TLS handshake, before any login/token logic). This is orthogonal to this change; the downstream tests run on Linux, and the Rust integration tests validate the same code paths.

[github-actions]

📊 Code Coverage Report

🔥 Diff Coverage 66%

🎯 Overall Coverage 91.5%

📦 Project: mssql-tds ℹ️ Note: diff coverage is reported, not enforced.

---

Diff Coverage

Diff: main...HEAD, staged and unstaged changes

• mssql-mock-tds/src/server.rs (66.1%): Missing lines 81-86,88,98,178,516-517,850,1052,1060-1065

Summary

• Total: 56 lines
• Missing: 19 lines
• Coverage: 66%

mssql-mock-tds/src/server.rs

  77 }
  78 
  79 impl ConnectionProcessor {
  80     /// Create a new connection processor
! 81     pub fn new(
! 82         conn_id: u64,
! 83         addr: SocketAddr,
! 84         query_registry: Arc<Mutex<QueryRegistry>>,
! 85         connection_store: Option<Arc<Mutex<ConnectionStore>>>,
! 86     ) -> Self {
  87         Self {
! 88             conn_id,
  89             addr,
  90             is_authenticated: false,
  91             received_token: None,
  92             user_agent: None,

---

   94             awaiting_fedauth_token: false,
   95             query_registry,
   96             buffer: BytesMut::with_capacity(4096),
   97             redirection: None,
!  98             connection_store,
   99         }
  100     }
  101 
  102     /// Create a new connection processor with redirection configuration

---

  174     /// moment the client's blocking LoginAck read returns.
  175     async fn record_to_store(&self) {
  176         if let Some(store) = &self.connection_store {
  177             store.lock().await.store(self);
! 178         }
  179     }
  180 
  181     /// Process a single packet from the buffer and return the response
  182     pub async fn process_packet(&mut self) -> Result<Option<BytesMut>, ProtocolError> {

---

  512         self.connections.insert(processor.conn_id(), info);
  513     }
  514 
  515     /// Get connection info by connection id
! 516     pub fn get(&self, conn_id: u64) -> Option<&ConnectionInfo> {
! 517         self.connections.get(&conn_id)
  518     }
  519 
  520     /// Get all connection infos
  521     pub fn all(&self) -> &BTreeMap<u64, ConnectionInfo> {

---

  846             info!("TLS handshake successful for {}", addr);
  847             handle_encrypted_tds_wrapped_connection(
  848                 tls_stream,
  849                 addr,
! 850                 conn_id,
  851                 query_registry,
  852                 connection_store,
  853                 redirection,
  854             )

---

  1048 /// Always supports FedAuth and username/password authentication.
  1049 async fn handle_encrypted_tds_wrapped_connection(
  1050     mut socket: TlsStream<crate::tds_tls_wrapper::TdsTlsWrapper>,
  1051     addr: SocketAddr,
! 1052     conn_id: u64,
  1053     query_registry: Arc<Mutex<QueryRegistry>>,
  1054     connection_store: Arc<Mutex<ConnectionStore>>,
  1055     redirection: Option<Arc<RedirectionConfig>>,
  1056 ) -> Result<(), ProtocolError> {

---

  1056 ) -> Result<(), ProtocolError> {
  1057     let redir_config = redirection
  1058         .as_ref()
  1059         .map(|r| RedirectionConfig::new(r.redirect_host.clone(), r.redirect_port));
! 1060     let mut processor = ConnectionProcessor::new_with_redirection(
! 1061         conn_id,
! 1062         addr,
! 1063         query_registry,
! 1064         Some(Arc::clone(&connection_store)),
! 1065         redir_config,
  1066     );
  1067 
  1068     loop {
  1069         // Read data from TLS socket (which wraps TdsTlsWrapper)

---

---

🔗 Quick Links

View Azure DevOps Build · Coverage Report

[Copilot]

Pull request overview

This PR makes FedAuth access-token capture in mssql-mock-tds deterministic and queryable immediately after login, addressing Linux-only flakiness caused by teardown-only recording and socket-address key collisions. It also updates upstream tests to stop relying on timing sleeps and adds a Rust regression guard ensuring tokens are visible while the client connection is still open.

Changes:

• Record connection state (incl. FedAuth token) eagerly into a shared store during login, and re-key the store by a monotonic per-connection conn_id.
• Remove time.sleep(...) workarounds from the Python FedAuth mock-server tests.
• Add a Rust regression test asserting the token is recorded before the client connection is closed.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.

File	Description
mssql-mock-tds/src/server.rs	Adds conn_id-keyed ConnectionStore and eager store upserts during FedAuth login paths.
mssql-tds/tests/test_mock_server_fedauth.rs	Adds a regression test to ensure token visibility before connection teardown.
mssql-py-core/tests/rs-only-tests/test_mock_server_fedauth.py	Removes timing sleeps now that token recording is eager/deterministic.