nautilus_network/websocket/
auth.rs

1// -------------------------------------------------------------------------------------------------
2//  Copyright (C) 2015-2026 Nautech Systems Pty Ltd. All rights reserved.
3//  https://nautechsystems.io
4//
5//  Licensed under the GNU Lesser General Public License Version 3.0 (the "License");
6//  You may not use this file except in compliance with the License.
7//  You may obtain a copy of the License at https://www.gnu.org/licenses/lgpl-3.0.en.html
8//
9//  Unless required by applicable law or agreed to in writing, software
10//  distributed under the License is distributed on an "AS IS" BASIS,
11//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12//  See the License for the specific language governing permissions and
13//  limitations under the License.
14// -------------------------------------------------------------------------------------------------
15
16//! Adapter authentication state independent of the WebSocket transport state.
17//!
18//! [`AuthTracker`] separates a specific authentication attempt from the shared session state.
19//! [`AuthTracker::begin`] returns a oneshot receiver for the attempt and fails any earlier pending
20//! attempt as superseded. [`AuthTracker::succeed`] and [`AuthTracker::fail`] resolve the active
21//! attempt and wake state waiters. [`AuthTracker::invalidate`] returns an authenticated session to
22//! unauthenticated without resolving a pending attempt or clearing terminal failure.
23//!
24//! # Client integration
25//!
26//! Registering a tracker with the client invalidates it on reconnectable connection loss and fails
27//! it on terminal shutdown. When authentication-gated replay is enabled, ordinary buffered sends
28//! wait for `Authenticated` and are discarded on `Failed`. The adapter remains responsible for
29//! sending authentication, interpreting the response, and ordering resubscription.
30
31use std::{
32    pin::pin,
33    sync::{
34        Arc,
35        atomic::{AtomicU8, Ordering},
36    },
37    time::Duration,
38};
39
40use parking_lot::Mutex;
41
42use crate::dst;
43
44pub type AuthResultSender = tokio::sync::oneshot::Sender<Result<(), String>>;
45pub type AuthResultReceiver = tokio::sync::oneshot::Receiver<Result<(), String>>;
46
47/// Authentication state for a WebSocket session.
48#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
49#[repr(u8)]
50pub enum AuthState {
51    /// Not authenticated (initial state, after begin, or after invalidating authenticated state).
52    #[default]
53    Unauthenticated = 0,
54    /// Successfully authenticated (after succeed).
55    Authenticated = 1,
56    /// Authentication failed or became impossible (after fail).
57    Failed = 2,
58}
59
60impl AuthState {
61    #[inline]
62    #[must_use]
63    #[expect(
64        clippy::match_same_arms,
65        reason = "explicit variant listing is clearer than collapsing 0 with wildcard"
66    )]
67    fn from_u8(value: u8) -> Self {
68        match value {
69            0 => Self::Unauthenticated,
70            1 => Self::Authenticated,
71            2 => Self::Failed,
72            _ => Self::Unauthenticated,
73        }
74    }
75
76    #[inline]
77    #[must_use]
78    const fn as_u8(self) -> u8 {
79        self as u8
80    }
81}
82
83/// Tracks authentication state for WebSocket connections.
84///
85/// Each authentication attempt receives a dedicated oneshot channel that resolves when the server
86/// responds.
87///
88/// # State management
89///
90/// The tracker maintains three states:
91///
92/// - [`AuthState::Unauthenticated`]: The initial state, the state after [`Self::begin`], and the
93///   result of [`Self::invalidate`] from [`AuthState::Authenticated`].
94/// - [`AuthState::Authenticated`]: The state after [`Self::succeed`].
95/// - [`AuthState::Failed`]: The state after [`Self::fail`]. Authentication waiters return early in
96///   this state.
97///
98/// # Superseding behavior
99///
100/// If a new authentication attempt begins while another remains pending, the old attempt is
101/// cancelled with an error. This prevents responses from an earlier attempt from racing with a
102/// later attempt during rapid reconnections.
103///
104/// # Thread safety
105///
106/// Clones share the pending attempt and session state. All operations are thread-safe and can run
107/// concurrently from multiple tasks.
108#[derive(Clone, Debug)]
109pub struct AuthTracker {
110    tx: Arc<Mutex<Option<AuthResultSender>>>,
111    state: Arc<AtomicU8>,
112    state_notify: Arc<tokio::sync::Notify>,
113}
114
115impl AuthTracker {
116    /// Creates a new authentication tracker.
117    #[must_use]
118    pub fn new() -> Self {
119        Self {
120            tx: Arc::new(Mutex::new(None)),
121            state: Arc::new(AtomicU8::new(AuthState::Unauthenticated.as_u8())),
122            state_notify: Arc::new(tokio::sync::Notify::new()),
123        }
124    }
125
126    /// Returns the current authentication state.
127    #[must_use]
128    pub fn auth_state(&self) -> AuthState {
129        AuthState::from_u8(self.state.load(Ordering::Acquire))
130    }
131
132    /// Returns whether the client is currently authenticated.
133    #[must_use]
134    pub fn is_authenticated(&self) -> bool {
135        self.auth_state() == AuthState::Authenticated
136    }
137
138    /// Clears authenticated state without affecting pending auth attempts.
139    ///
140    /// Call this when a live connection drops and reconnect may authenticate
141    /// again, so operations requiring authentication are properly guarded. A
142    /// terminal [`AuthState::Failed`] state remains failed.
143    pub fn invalidate(&self) {
144        let _guard = self.tx.lock();
145
146        if self
147            .state
148            .try_update(Ordering::AcqRel, Ordering::Acquire, |state| {
149                (AuthState::from_u8(state) == AuthState::Authenticated)
150                    .then_some(AuthState::Unauthenticated.as_u8())
151            })
152            .is_ok()
153        {
154            self.state_notify.notify_waiters();
155        }
156    }
157
158    /// Begins a new authentication attempt.
159    ///
160    /// Returns a receiver that will be notified when authentication completes.
161    /// If a previous authentication attempt is still pending, it will be cancelled
162    /// with an error message indicating it was superseded.
163    ///
164    /// Transitions to `Unauthenticated` since a new attempt invalidates any
165    /// previous status.
166    #[allow(
167        clippy::must_use_candidate,
168        reason = "callers use this for side effects"
169    )]
170    pub fn begin(&self) -> AuthResultReceiver {
171        let (sender, receiver) = tokio::sync::oneshot::channel();
172        let mut guard = self.tx.lock();
173        self.state
174            .store(AuthState::Unauthenticated.as_u8(), Ordering::Release);
175
176        if let Some(old) = guard.take() {
177            log::warn!("New authentication request superseding previous pending request");
178            let _ = old.send(Err("Authentication attempt superseded".to_string()));
179        } else {
180            log::debug!("Starting new authentication request");
181        }
182
183        *guard = Some(sender);
184
185        receiver
186    }
187
188    /// Marks the current authentication attempt as successful.
189    ///
190    /// Transitions to `Authenticated` and notifies any waiting receiver
191    /// with `Ok(())`. This should be called when the server sends a successful
192    /// authentication response.
193    ///
194    /// The state is always updated even if no receiver is waiting (e.g., after
195    /// a timeout), since the server has confirmed authentication.
196    pub fn succeed(&self) {
197        let mut guard = self.tx.lock();
198        self.state
199            .store(AuthState::Authenticated.as_u8(), Ordering::Release);
200        self.state_notify.notify_waiters();
201
202        if let Some(sender) = guard.take() {
203            let _ = sender.send(Ok(()));
204        }
205    }
206
207    /// Marks the current authentication attempt as failed.
208    ///
209    /// Transitions to `Failed` and notifies any waiting receiver
210    /// with `Err(message)`. This should be called when the server sends an
211    /// authentication error response, or on terminal client shutdown.
212    ///
213    /// The state is always updated even if no receiver is waiting, since the
214    /// server has rejected authentication or future auth is impossible.
215    pub fn fail(&self, error: impl Into<String>) {
216        let message = error.into();
217        let mut guard = self.tx.lock();
218        self.state
219            .store(AuthState::Failed.as_u8(), Ordering::Release);
220        self.state_notify.notify_waiters();
221
222        if let Some(sender) = guard.take() {
223            let _ = sender.send(Err(message));
224        }
225    }
226
227    /// Waits for the authentication result with a timeout.
228    ///
229    /// Returns `Ok(())` if authentication succeeds, or an error if it fails,
230    /// times out, or the channel is closed.
231    ///
232    /// # Type Parameters
233    ///
234    /// - `E`: Error type that implements `From<String>` for error message conversion
235    ///
236    /// # Errors
237    ///
238    /// Returns an error if:
239    /// - Authentication fails (server rejects credentials).
240    /// - Authentication times out (no response within timeout duration).
241    /// - Authentication channel closes unexpectedly.
242    /// - Authentication attempt is superseded by a new attempt.
243    pub async fn wait_for_result<E>(
244        &self,
245        timeout: Duration,
246        receiver: AuthResultReceiver,
247    ) -> Result<(), E>
248    where
249        E: From<String>,
250    {
251        match dst::time::timeout(timeout, receiver).await {
252            Ok(Ok(Ok(()))) => Ok(()),
253            Ok(Ok(Err(msg))) => Err(E::from(msg)),
254            Ok(Err(_)) => Err(E::from("Authentication channel closed".to_string())),
255            Err(_) => {
256                // Don't clear the sender: a concurrent begin() may have replaced it,
257                // and guard.take() would cancel the newer sender. The next begin()
258                // call cleans up any stale sender.
259                Err(E::from("Authentication timed out".to_string()))
260            }
261        }
262    }
263
264    /// Waits for the tracker to enter the authenticated state.
265    ///
266    /// Returns `true` if authenticated within the timeout, `false` if the timeout
267    /// expires or authentication explicitly fails. Uses event-driven notification
268    /// from `succeed()` / `fail()` / `invalidate()` to avoid polling.
269    ///
270    /// Returns early with `false` when `fail()` is called (e.g., the exchange
271    /// rejects credentials), so callers are not blocked for the full timeout
272    /// on a definitive auth rejection.
273    ///
274    /// This is intended for callers on a separate task who need to gate operations
275    /// on authentication state (e.g., order sends that must wait for re-authentication
276    /// after a WebSocket reconnection).
277    pub async fn wait_for_authenticated(&self, timeout: Duration) -> bool {
278        if self.is_authenticated() {
279            return true;
280        }
281
282        dst::time::timeout(timeout, async {
283            loop {
284                // Enable before the state check: an unpolled Notified is unregistered and misses notifies
285                let mut notified = pin!(self.state_notify.notified());
286                notified.as_mut().enable();
287
288                match self.auth_state() {
289                    AuthState::Authenticated => return true,
290                    AuthState::Failed => return false,
291                    AuthState::Unauthenticated => notified.await,
292                }
293            }
294        })
295        .await
296        .unwrap_or(false)
297    }
298}
299
300impl Default for AuthTracker {
301    fn default() -> Self {
302        Self::new()
303    }
304}
305
306#[cfg(test)]
307#[cfg(not(all(feature = "simulation", madsim)))]
308mod tests {
309    use std::{
310        sync::{
311            atomic::{AtomicBool, Ordering},
312            mpsc,
313        },
314        time::Duration,
315    };
316
317    use rstest::rstest;
318
319    use super::*;
320
321    #[derive(Debug, PartialEq)]
322    struct TestError(String);
323
324    impl From<String> for TestError {
325        fn from(msg: String) -> Self {
326            Self(msg)
327        }
328    }
329
330    #[rstest]
331    fn test_begin_during_failure_keeps_receiver_and_state_consistent() {
332        struct DeferredError {
333            entered: mpsc::SyncSender<()>,
334            proceed: mpsc::Receiver<()>,
335        }
336
337        impl From<DeferredError> for String {
338            fn from(error: DeferredError) -> Self {
339                error.entered.send(()).unwrap();
340                error.proceed.recv_timeout(Duration::from_secs(2)).unwrap();
341                "credentials rejected".to_string()
342            }
343        }
344
345        let tracker = AuthTracker::new();
346        let (entered_tx, entered_rx) = mpsc::sync_channel(1);
347        let (proceed_tx, proceed_rx) = mpsc::sync_channel(1);
348
349        let mut receiver = std::thread::scope(|scope| {
350            let completion = scope.spawn(|| {
351                tracker.fail(DeferredError {
352                    entered: entered_tx,
353                    proceed: proceed_rx,
354                });
355            });
356
357            entered_rx.recv_timeout(Duration::from_secs(2)).unwrap();
358            let receiver = tracker.begin();
359            proceed_tx.send(()).unwrap();
360            completion.join().unwrap();
361            receiver
362        });
363
364        assert_eq!(
365            receiver.try_recv(),
366            Ok(Err("credentials rejected".to_string()))
367        );
368        assert_eq!(tracker.auth_state(), AuthState::Failed);
369    }
370
371    #[rstest]
372    #[case::success(AuthTracker::succeed, Ok(()), AuthState::Authenticated)]
373    #[case::failure(|tracker: &AuthTracker| tracker.fail("rejected"), Err("rejected".to_string()), AuthState::Failed)]
374    fn test_concurrent_begin_and_completion_keep_receiver_and_state_consistent(
375        #[case] complete: fn(&AuthTracker),
376        #[case] result: Result<(), String>,
377        #[case] state: AuthState,
378    ) {
379        const ATTEMPTS: usize = 2_000;
380
381        let tracker = AuthTracker::new();
382        let (begin_tx, begin_rx) = mpsc::sync_channel(1);
383        let (complete_tx, complete_rx) = mpsc::sync_channel(1);
384        let (receiver_tx, receiver_rx) = mpsc::sync_channel(1);
385        let (done_tx, done_rx) = mpsc::sync_channel(1);
386
387        let observations = std::thread::scope(|scope| {
388            let tracker = &tracker;
389
390            scope.spawn(move || {
391                for _ in 0..ATTEMPTS {
392                    begin_rx.recv_timeout(Duration::from_secs(2)).unwrap();
393                    receiver_tx.send(tracker.begin()).unwrap();
394                }
395            });
396
397            scope.spawn(move || {
398                for _ in 0..ATTEMPTS {
399                    complete_rx.recv_timeout(Duration::from_secs(2)).unwrap();
400                    complete(tracker);
401                    done_tx.send(()).unwrap();
402                }
403            });
404
405            let mut observations = Vec::with_capacity(ATTEMPTS);
406
407            for _ in 0..ATTEMPTS {
408                begin_tx.send(()).unwrap();
409                complete_tx.send(()).unwrap();
410                let mut receiver = receiver_rx.recv_timeout(Duration::from_secs(2)).unwrap();
411                done_rx.recv_timeout(Duration::from_secs(2)).unwrap();
412                observations.push((receiver.try_recv(), tracker.auth_state()));
413            }
414
415            observations
416        });
417
418        for (received, actual_state) in observations {
419            match received {
420                Ok(actual_result) => {
421                    assert_eq!(actual_result, result);
422                    assert_eq!(actual_state, state);
423                }
424                Err(tokio::sync::oneshot::error::TryRecvError::Empty) => {
425                    assert_eq!(actual_state, AuthState::Unauthenticated);
426                }
427                Err(receive_error) => panic!("authentication channel closed: {receive_error}"),
428            }
429        }
430    }
431
432    #[rstest]
433    #[tokio::test]
434    async fn test_successful_authentication() {
435        let tracker = AuthTracker::new();
436        let rx = tracker.begin();
437
438        tracker.succeed();
439
440        let result: Result<(), TestError> =
441            tracker.wait_for_result(Duration::from_secs(1), rx).await;
442
443        assert!(result.is_ok());
444    }
445
446    #[rstest]
447    #[tokio::test]
448    async fn test_failed_authentication() {
449        let tracker = AuthTracker::new();
450        let rx = tracker.begin();
451
452        tracker.fail("Invalid credentials");
453
454        let result: Result<(), TestError> =
455            tracker.wait_for_result(Duration::from_secs(1), rx).await;
456
457        assert_eq!(
458            result.unwrap_err(),
459            TestError("Invalid credentials".to_string())
460        );
461    }
462
463    #[rstest]
464    #[tokio::test]
465    async fn test_authentication_timeout() {
466        let tracker = AuthTracker::new();
467        let rx = tracker.begin();
468
469        // Don't call succeed or fail - let it timeout
470
471        let result: Result<(), TestError> =
472            tracker.wait_for_result(Duration::from_millis(50), rx).await;
473
474        assert_eq!(
475            result.unwrap_err(),
476            TestError("Authentication timed out".to_string())
477        );
478    }
479
480    #[rstest]
481    #[tokio::test]
482    async fn test_begin_supersedes_previous_sender() {
483        let tracker = AuthTracker::new();
484
485        let first = tracker.begin();
486        let second = tracker.begin();
487
488        // First receiver should get superseded error
489        let result = first.await.expect("oneshot closed unexpectedly");
490        assert_eq!(result, Err("Authentication attempt superseded".to_string()));
491
492        // Second attempt should succeed
493        tracker.succeed();
494        let result: Result<(), TestError> = tracker
495            .wait_for_result(Duration::from_secs(1), second)
496            .await;
497
498        assert!(result.is_ok());
499    }
500
501    #[rstest]
502    #[tokio::test]
503    async fn test_succeed_without_pending_auth() {
504        let tracker = AuthTracker::new();
505
506        tracker.succeed();
507
508        assert_eq!(tracker.auth_state(), AuthState::Authenticated);
509    }
510
511    #[rstest]
512    #[tokio::test]
513    async fn test_fail_without_pending_auth() {
514        let tracker = AuthTracker::new();
515
516        tracker.fail("Some error");
517
518        assert_eq!(tracker.auth_state(), AuthState::Failed);
519    }
520
521    #[rstest]
522    #[tokio::test]
523    async fn test_multiple_sequential_authentications() {
524        let tracker = AuthTracker::new();
525
526        // First auth succeeds
527        let rx1 = tracker.begin();
528        tracker.succeed();
529        let result1: Result<(), TestError> =
530            tracker.wait_for_result(Duration::from_secs(1), rx1).await;
531        assert!(result1.is_ok());
532
533        // Second auth fails
534        let rx2 = tracker.begin();
535        tracker.fail("Credentials expired");
536        let result2: Result<(), TestError> =
537            tracker.wait_for_result(Duration::from_secs(1), rx2).await;
538        assert_eq!(
539            result2.unwrap_err(),
540            TestError("Credentials expired".to_string())
541        );
542
543        // Third auth succeeds
544        let rx3 = tracker.begin();
545        tracker.succeed();
546        let result3: Result<(), TestError> =
547            tracker.wait_for_result(Duration::from_secs(1), rx3).await;
548        assert!(result3.is_ok());
549    }
550
551    #[rstest]
552    #[tokio::test]
553    async fn test_channel_closed_before_result() {
554        let tracker = AuthTracker::new();
555        let rx = tracker.begin();
556
557        drop(tracker.tx.lock().take());
558
559        let result: Result<(), TestError> =
560            tracker.wait_for_result(Duration::from_secs(1), rx).await;
561
562        assert_eq!(
563            result,
564            Err(TestError("Authentication channel closed".to_string()))
565        );
566        assert_eq!(tracker.auth_state(), AuthState::Unauthenticated);
567    }
568
569    #[rstest]
570    #[tokio::test]
571    async fn test_concurrent_auth_attempts() {
572        let tracker = Arc::new(AuthTracker::new());
573        let mut handles = vec![];
574
575        // Spawn 10 concurrent auth attempts
576        for i in 0..10 {
577            let tracker_clone = Arc::clone(&tracker);
578            let handle = tokio::spawn(async move {
579                let rx = tracker_clone.begin();
580
581                // Only the last one should succeed
582                if i == 9 {
583                    tokio::time::sleep(Duration::from_millis(10)).await;
584                    tracker_clone.succeed();
585                }
586
587                let result: Result<(), TestError> = tracker_clone
588                    .wait_for_result(Duration::from_secs(1), rx)
589                    .await;
590
591                (i, result)
592            });
593
594            handles.push(handle);
595        }
596
597        let mut successes = 0;
598        let mut superseded = 0;
599
600        for handle in handles {
601            let (i, result) = handle.await.unwrap();
602            match result {
603                Ok(()) => {
604                    // Only task 9 should succeed
605                    assert_eq!(i, 9);
606                    successes += 1;
607                }
608                Err(TestError(msg)) if msg.contains("superseded") => {
609                    superseded += 1;
610                }
611                Err(e) => panic!("Unexpected error: {e:?}"),
612            }
613        }
614
615        assert_eq!(successes, 1);
616        assert_eq!(superseded, 9);
617    }
618
619    #[rstest]
620    fn test_default_trait() {
621        let _tracker = AuthTracker::default();
622    }
623
624    #[rstest]
625    #[tokio::test]
626    async fn test_clone_trait() {
627        let tracker = AuthTracker::new();
628        let cloned = tracker.clone();
629
630        // Verify cloned instance shares state with original (Arc behavior)
631        let rx = tracker.begin();
632        cloned.succeed(); // Succeed via clone affects original
633        let result: Result<(), TestError> =
634            tracker.wait_for_result(Duration::from_secs(1), rx).await;
635        assert!(result.is_ok());
636    }
637
638    #[rstest]
639    fn test_debug_trait() {
640        let tracker = AuthTracker::new();
641        let debug_str = format!("{tracker:?}");
642        assert!(debug_str.contains("AuthTracker"));
643    }
644
645    #[rstest]
646    #[tokio::test]
647    async fn test_timeout_clears_sender() {
648        let tracker = AuthTracker::new();
649
650        // Start auth that will timeout
651        let rx1 = tracker.begin();
652        let result1: Result<(), TestError> = tracker
653            .wait_for_result(Duration::from_millis(50), rx1)
654            .await;
655        assert_eq!(
656            result1.unwrap_err(),
657            TestError("Authentication timed out".to_string())
658        );
659
660        // Verify sender was cleared - new auth should work
661        let rx2 = tracker.begin();
662        tracker.succeed();
663        let result2: Result<(), TestError> =
664            tracker.wait_for_result(Duration::from_secs(1), rx2).await;
665        assert!(result2.is_ok());
666    }
667
668    #[rstest]
669    #[tokio::test]
670    async fn test_fail_clears_sender() {
671        let tracker = AuthTracker::new();
672
673        // Auth fails
674        let rx1 = tracker.begin();
675        tracker.fail("Bad credentials");
676        let result1: Result<(), TestError> =
677            tracker.wait_for_result(Duration::from_secs(1), rx1).await;
678        assert!(result1.is_err());
679
680        // Verify sender was cleared - new auth should work
681        let rx2 = tracker.begin();
682        tracker.succeed();
683        let result2: Result<(), TestError> =
684            tracker.wait_for_result(Duration::from_secs(1), rx2).await;
685        assert!(result2.is_ok());
686    }
687
688    #[rstest]
689    #[tokio::test]
690    async fn test_succeed_clears_sender() {
691        let tracker = AuthTracker::new();
692
693        // Auth succeeds
694        let rx1 = tracker.begin();
695        tracker.succeed();
696        let result1: Result<(), TestError> =
697            tracker.wait_for_result(Duration::from_secs(1), rx1).await;
698        assert!(result1.is_ok());
699
700        // Verify sender was cleared - new auth should work
701        let rx2 = tracker.begin();
702        tracker.succeed();
703        let result2: Result<(), TestError> =
704            tracker.wait_for_result(Duration::from_secs(1), rx2).await;
705        assert!(result2.is_ok());
706    }
707
708    #[rstest]
709    #[tokio::test]
710    async fn test_rapid_begin_succeed_cycles() {
711        let tracker = AuthTracker::new();
712
713        // Rapidly cycle through auth attempts
714        for _ in 0..100 {
715            let rx = tracker.begin();
716            tracker.succeed();
717            let result: Result<(), TestError> =
718                tracker.wait_for_result(Duration::from_secs(1), rx).await;
719            assert!(result.is_ok());
720        }
721    }
722
723    #[rstest]
724    #[tokio::test]
725    async fn test_double_succeed_is_safe() {
726        let tracker = AuthTracker::new();
727        let rx = tracker.begin();
728
729        // Call succeed twice
730        tracker.succeed();
731        tracker.succeed(); // Second call should be no-op
732
733        let result: Result<(), TestError> =
734            tracker.wait_for_result(Duration::from_secs(1), rx).await;
735        assert!(result.is_ok());
736    }
737
738    #[rstest]
739    #[tokio::test]
740    async fn test_double_fail_is_safe() {
741        let tracker = AuthTracker::new();
742        let rx = tracker.begin();
743
744        // Call fail twice
745        tracker.fail("Error 1");
746        tracker.fail("Error 2"); // Second call should be no-op
747
748        let result: Result<(), TestError> =
749            tracker.wait_for_result(Duration::from_secs(1), rx).await;
750        assert_eq!(
751            result.unwrap_err(),
752            TestError("Error 1".to_string()) // Should be first error
753        );
754    }
755
756    #[rstest]
757    #[tokio::test]
758    async fn test_succeed_after_fail_is_ignored() {
759        let tracker = AuthTracker::new();
760        let rx = tracker.begin();
761
762        tracker.fail("Auth failed");
763        tracker.succeed(); // This should be no-op
764
765        let result: Result<(), TestError> =
766            tracker.wait_for_result(Duration::from_secs(1), rx).await;
767        assert!(result.is_err()); // Should still be error
768    }
769
770    #[rstest]
771    #[tokio::test]
772    async fn test_fail_after_succeed_is_ignored() {
773        let tracker = AuthTracker::new();
774        let rx = tracker.begin();
775
776        tracker.succeed();
777        tracker.fail("Auth failed"); // This should be no-op
778
779        let result: Result<(), TestError> =
780            tracker.wait_for_result(Duration::from_secs(1), rx).await;
781        assert!(result.is_ok()); // Should still be success
782    }
783
784    /// Simulates a reconnect flow where authentication must complete before resubscription.
785    ///
786    /// This is an integration-style test that verifies:
787    /// 1. On reconnect, authentication starts first
788    /// 2. Subscription logic waits for auth to complete
789    /// 3. Subscriptions only proceed after successful auth
790    #[rstest]
791    #[tokio::test]
792    async fn test_reconnect_flow_waits_for_auth() {
793        let tracker = Arc::new(AuthTracker::new());
794        let subscribed = Arc::new(tokio::sync::Notify::new());
795        let auth_completed = Arc::new(tokio::sync::Notify::new());
796
797        // Simulate reconnect handler
798        let tracker_reconnect = Arc::clone(&tracker);
799        let subscribed_reconnect = Arc::clone(&subscribed);
800        let auth_completed_reconnect = Arc::clone(&auth_completed);
801
802        let reconnect_task = tokio::spawn(async move {
803            // Step 1: Begin authentication
804            let rx = tracker_reconnect.begin();
805
806            // Step 2: Spawn resubscription task that waits for auth
807            let tracker_resub = Arc::clone(&tracker_reconnect);
808            let subscribed_resub = Arc::clone(&subscribed_reconnect);
809            let auth_completed_resub = Arc::clone(&auth_completed_reconnect);
810
811            let resub_task = tokio::spawn(async move {
812                // Wait for auth to complete
813                let result: Result<(), TestError> = tracker_resub
814                    .wait_for_result(Duration::from_secs(5), rx)
815                    .await;
816
817                if result.is_ok() {
818                    auth_completed_resub.notify_one();
819                    // Simulate resubscription
820                    tokio::time::sleep(Duration::from_millis(10)).await;
821                    subscribed_resub.notify_one();
822                }
823            });
824
825            resub_task.await.unwrap();
826        });
827
828        // Simulate server auth response after delay
829        tokio::time::sleep(Duration::from_millis(100)).await;
830        tracker.succeed();
831
832        // Wait for reconnect flow to complete
833        reconnect_task.await.unwrap();
834
835        // Verify auth completed before subscription
836        tokio::select! {
837            () = auth_completed.notified() => {
838                // Good - auth completed
839            }
840            () = tokio::time::sleep(Duration::from_secs(1)) => {
841                panic!("Auth never completed");
842            }
843        }
844
845        // Verify subscription completed
846        tokio::select! {
847            () = subscribed.notified() => {
848                // Good - subscribed
849            }
850            () = tokio::time::sleep(Duration::from_secs(1)) => {
851                panic!("Subscription never completed");
852            }
853        }
854    }
855
856    /// Verifies that failed authentication prevents resubscription in reconnect flow.
857    #[rstest]
858    #[tokio::test]
859    async fn test_reconnect_flow_blocks_on_auth_failure() {
860        let tracker = Arc::new(AuthTracker::new());
861        let subscribed = Arc::new(AtomicBool::new(false));
862
863        let tracker_reconnect = Arc::clone(&tracker);
864        let subscribed_reconnect = Arc::clone(&subscribed);
865
866        let reconnect_task = tokio::spawn(async move {
867            let rx = tracker_reconnect.begin();
868
869            // Spawn resubscription task that waits for auth
870            let tracker_resub = Arc::clone(&tracker_reconnect);
871            let subscribed_resub = Arc::clone(&subscribed_reconnect);
872
873            let resub_task = tokio::spawn(async move {
874                let result: Result<(), TestError> = tracker_resub
875                    .wait_for_result(Duration::from_secs(5), rx)
876                    .await;
877
878                // Only subscribe if auth succeeds
879                if result.is_ok() {
880                    subscribed_resub.store(true, Ordering::Relaxed);
881                }
882            });
883
884            resub_task.await.unwrap();
885        });
886
887        // Simulate server auth failure
888        tokio::time::sleep(Duration::from_millis(50)).await;
889        tracker.fail("Invalid credentials");
890
891        // Wait for reconnect flow to complete
892        reconnect_task.await.unwrap();
893
894        // Verify subscription never happened
895        tokio::time::sleep(Duration::from_millis(100)).await;
896        assert!(!subscribed.load(Ordering::Relaxed));
897    }
898
899    /// Tests state machine transitions exhaustively.
900    #[rstest]
901    #[tokio::test]
902    async fn test_state_machine_transitions() {
903        let tracker = AuthTracker::new();
904
905        // Transition 1: Initial -> Pending (begin)
906        let rx1 = tracker.begin();
907
908        // Transition 2: Pending -> Success (succeed)
909        tracker.succeed();
910        let result1: Result<(), TestError> =
911            tracker.wait_for_result(Duration::from_secs(1), rx1).await;
912        assert!(result1.is_ok());
913
914        // Transition 3: Success -> Pending (begin again)
915        let rx2 = tracker.begin();
916
917        // Transition 4: Pending -> Failure (fail)
918        tracker.fail("Error");
919        let result2: Result<(), TestError> =
920            tracker.wait_for_result(Duration::from_secs(1), rx2).await;
921        assert!(result2.is_err());
922
923        // Transition 5: Failure -> Pending (begin again)
924        let rx3 = tracker.begin();
925
926        // Transition 6: Pending -> Timeout
927        let result3: Result<(), TestError> = tracker
928            .wait_for_result(Duration::from_millis(50), rx3)
929            .await;
930        assert_eq!(
931            result3.unwrap_err(),
932            TestError("Authentication timed out".to_string())
933        );
934
935        // Transition 7: Timeout -> Pending (begin again)
936        let rx4 = tracker.begin();
937
938        // Transition 8: Pending -> Superseded (begin interrupts)
939        let rx5 = tracker.begin();
940        let result4: Result<(), TestError> =
941            tracker.wait_for_result(Duration::from_secs(1), rx4).await;
942        assert_eq!(
943            result4.unwrap_err(),
944            TestError("Authentication attempt superseded".to_string())
945        );
946
947        // Final success to clean up
948        tracker.succeed();
949        let result5: Result<(), TestError> =
950            tracker.wait_for_result(Duration::from_secs(1), rx5).await;
951        assert!(result5.is_ok());
952    }
953
954    /// Verifies no memory leaks from orphaned senders.
955    #[rstest]
956    #[tokio::test]
957    async fn test_no_sender_leaks() {
958        let tracker = AuthTracker::new();
959        for _ in 0..100 {
960            let rx = tracker.begin();
961            let _result: Result<(), TestError> =
962                tracker.wait_for_result(Duration::from_millis(1), rx).await;
963        }
964
965        let rx = tracker.begin();
966        tracker.succeed();
967        let result: Result<(), TestError> =
968            tracker.wait_for_result(Duration::from_secs(1), rx).await;
969        assert!(result.is_ok());
970    }
971
972    /// Tests concurrent success/fail calls don't cause panics.
973    #[rstest]
974    #[tokio::test]
975    async fn test_concurrent_succeed_fail_calls() {
976        let tracker = Arc::new(AuthTracker::new());
977        let rx = tracker.begin();
978
979        let mut handles = vec![];
980
981        // Spawn many tasks trying to succeed
982        for _ in 0..50 {
983            let tracker_clone = Arc::clone(&tracker);
984            handles.push(tokio::spawn(async move {
985                tracker_clone.succeed();
986            }));
987        }
988
989        // Spawn many tasks trying to fail
990        for _ in 0..50 {
991            let tracker_clone = Arc::clone(&tracker);
992            handles.push(tokio::spawn(async move {
993                tracker_clone.fail("Error");
994            }));
995        }
996
997        // Wait for all tasks
998        for handle in handles {
999            handle.await.unwrap();
1000        }
1001
1002        // Should get either success or failure, but not panic
1003        let result: Result<(), TestError> =
1004            tracker.wait_for_result(Duration::from_secs(1), rx).await;
1005        // Don't care which outcome, just that it doesn't panic
1006        let _ = result;
1007    }
1008
1009    #[rstest]
1010    fn test_is_authenticated_initial_state() {
1011        let tracker = AuthTracker::new();
1012        assert!(!tracker.is_authenticated());
1013    }
1014
1015    #[rstest]
1016    #[tokio::test]
1017    async fn test_is_authenticated_after_succeed() {
1018        let tracker = AuthTracker::new();
1019        assert!(!tracker.is_authenticated());
1020
1021        let _rx = tracker.begin();
1022        assert!(!tracker.is_authenticated());
1023
1024        tracker.succeed();
1025        assert!(tracker.is_authenticated());
1026    }
1027
1028    #[rstest]
1029    #[tokio::test]
1030    async fn test_is_authenticated_after_fail() {
1031        let tracker = AuthTracker::new();
1032        let _rx = tracker.begin();
1033        tracker.fail("error");
1034        assert!(!tracker.is_authenticated());
1035    }
1036
1037    #[rstest]
1038    #[tokio::test]
1039    async fn test_invalidate_clears_auth_state() {
1040        let tracker = AuthTracker::new();
1041        let _rx = tracker.begin();
1042        tracker.succeed();
1043        assert!(tracker.is_authenticated());
1044
1045        tracker.invalidate();
1046        assert!(!tracker.is_authenticated());
1047    }
1048
1049    #[rstest]
1050    #[case(true)]
1051    #[case(false)]
1052    #[tokio::test]
1053    async fn test_invalidate_preserves_terminal_failure(#[case] invalidate_first: bool) {
1054        let tracker = AuthTracker::new();
1055        let receiver = tracker.begin();
1056        if invalidate_first {
1057            tracker.invalidate();
1058            tracker.fail("terminal");
1059        } else {
1060            tracker.fail("terminal");
1061            tracker.invalidate();
1062        }
1063
1064        let result: Result<(), TestError> = tracker
1065            .wait_for_result(Duration::from_secs(1), receiver)
1066            .await;
1067
1068        assert_eq!(tracker.auth_state(), AuthState::Failed);
1069        assert_eq!(result.unwrap_err(), TestError("terminal".to_string()));
1070        assert!(
1071            !tracker
1072                .wait_for_authenticated(Duration::from_millis(10))
1073                .await
1074        );
1075    }
1076
1077    #[rstest]
1078    #[tokio::test]
1079    async fn test_begin_clears_auth_state() {
1080        let tracker = AuthTracker::new();
1081        let _rx1 = tracker.begin();
1082        tracker.succeed();
1083        assert!(tracker.is_authenticated());
1084
1085        let _rx2 = tracker.begin();
1086        assert!(!tracker.is_authenticated());
1087    }
1088
1089    #[rstest]
1090    fn test_is_authenticated_shared_across_clones() {
1091        let tracker = AuthTracker::new();
1092        let cloned = tracker.clone();
1093
1094        let _rx = tracker.begin();
1095        tracker.succeed();
1096
1097        assert!(cloned.is_authenticated());
1098    }
1099
1100    #[rstest]
1101    fn test_invalidate_shared_across_clones() {
1102        let tracker = AuthTracker::new();
1103        let cloned = tracker.clone();
1104
1105        let _rx = tracker.begin();
1106        tracker.succeed();
1107        assert!(tracker.is_authenticated());
1108
1109        cloned.invalidate();
1110        assert!(!tracker.is_authenticated());
1111    }
1112
1113    #[rstest]
1114    fn test_succeed_without_begin_still_updates_auth_state() {
1115        let tracker = AuthTracker::new();
1116        assert!(!tracker.is_authenticated());
1117
1118        // State updates even without begin() to handle late responses after timeout
1119        tracker.succeed();
1120        assert!(tracker.is_authenticated());
1121    }
1122
1123    #[rstest]
1124    fn test_fail_without_begin_still_updates_auth_state() {
1125        let tracker = AuthTracker::new();
1126        tracker.succeed();
1127        assert!(tracker.is_authenticated());
1128
1129        // State updates even without begin() to handle late responses
1130        tracker.fail("error");
1131        assert!(!tracker.is_authenticated());
1132    }
1133
1134    #[rstest]
1135    #[tokio::test]
1136    async fn test_auth_state_false_after_timeout_until_late_response() {
1137        let tracker = AuthTracker::new();
1138        let rx = tracker.begin();
1139        assert!(!tracker.is_authenticated());
1140
1141        let result: Result<(), TestError> =
1142            tracker.wait_for_result(Duration::from_millis(10), rx).await;
1143
1144        assert!(result.is_err());
1145        assert!(!tracker.is_authenticated());
1146
1147        // Late response after timeout still updates state
1148        tracker.succeed();
1149        assert!(tracker.is_authenticated());
1150    }
1151
1152    #[rstest]
1153    #[tokio::test]
1154    async fn test_wait_for_authenticated_already_authenticated() {
1155        let tracker = AuthTracker::new();
1156        let _rx = tracker.begin();
1157        tracker.succeed();
1158
1159        assert!(
1160            tracker
1161                .wait_for_authenticated(Duration::from_millis(50))
1162                .await
1163        );
1164    }
1165
1166    #[rstest]
1167    #[tokio::test]
1168    async fn test_wait_for_authenticated_succeeds_after_delay() {
1169        let tracker = AuthTracker::new();
1170        let _rx = tracker.begin();
1171
1172        let tracker_clone = tracker.clone();
1173
1174        tokio::spawn(async move {
1175            tokio::time::sleep(Duration::from_millis(50)).await;
1176            tracker_clone.succeed();
1177        });
1178
1179        assert!(tracker.wait_for_authenticated(Duration::from_secs(1)).await);
1180    }
1181
1182    #[rstest]
1183    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1184    async fn test_wait_for_authenticated_no_lost_wakeup_under_race() {
1185        // Regression: a succeed() landing between the waiter's state check and
1186        // its first poll of Notified must not be lost. Notified only registers
1187        // with the Notify once polled or enabled, so without enable() this
1188        // stalls for the full timeout and returns false on some iterations.
1189        for _ in 0..200 {
1190            let tracker = AuthTracker::new();
1191            let _rx = tracker.begin();
1192
1193            let succeeder = tracker.clone();
1194            let handle = std::thread::spawn(move || succeeder.succeed());
1195
1196            assert!(
1197                tracker
1198                    .wait_for_authenticated(Duration::from_millis(500))
1199                    .await,
1200                "wakeup lost despite successful authentication"
1201            );
1202            handle.join().unwrap();
1203        }
1204    }
1205
1206    #[rstest]
1207    #[tokio::test]
1208    async fn test_wait_for_authenticated_returns_false_on_failure() {
1209        let tracker = AuthTracker::new();
1210        let _rx = tracker.begin();
1211
1212        let tracker_clone = tracker.clone();
1213
1214        tokio::spawn(async move {
1215            tokio::time::sleep(Duration::from_millis(50)).await;
1216            tracker_clone.fail("rejected");
1217        });
1218
1219        let start = tokio::time::Instant::now();
1220        let result = tracker.wait_for_authenticated(Duration::from_secs(5)).await;
1221        let elapsed = start.elapsed();
1222
1223        assert!(!result);
1224        assert!(elapsed < Duration::from_secs(1));
1225    }
1226
1227    #[rstest]
1228    #[tokio::test]
1229    async fn test_wait_for_authenticated_times_out() {
1230        let tracker = AuthTracker::new();
1231        let _rx = tracker.begin();
1232
1233        assert!(
1234            !tracker
1235                .wait_for_authenticated(Duration::from_millis(50))
1236                .await
1237        );
1238    }
1239
1240    #[rstest]
1241    #[tokio::test]
1242    async fn test_wait_for_authenticated_begin_clears_failed() {
1243        let tracker = AuthTracker::new();
1244        let _rx = tracker.begin();
1245        tracker.fail("first attempt");
1246
1247        assert!(
1248            !tracker
1249                .wait_for_authenticated(Duration::from_millis(10))
1250                .await
1251        );
1252
1253        // begin() clears the failed flag, allowing a fresh wait
1254        let _rx = tracker.begin();
1255
1256        let tracker_clone = tracker.clone();
1257
1258        tokio::spawn(async move {
1259            tokio::time::sleep(Duration::from_millis(50)).await;
1260            tracker_clone.succeed();
1261        });
1262
1263        assert!(tracker.wait_for_authenticated(Duration::from_secs(1)).await);
1264    }
1265
1266    #[rstest]
1267    #[tokio::test]
1268    async fn test_wait_for_authenticated_invalidate_does_not_return_false() {
1269        let tracker = AuthTracker::new();
1270        let _rx = tracker.begin();
1271
1272        let tracker_clone = tracker.clone();
1273
1274        tokio::spawn(async move {
1275            // invalidate wakes the loop but should not cause early false return
1276            tokio::time::sleep(Duration::from_millis(20)).await;
1277            tracker_clone.invalidate();
1278            // then succeed shortly after
1279            tokio::time::sleep(Duration::from_millis(20)).await;
1280            tracker_clone.succeed();
1281        });
1282
1283        assert!(tracker.wait_for_authenticated(Duration::from_secs(1)).await);
1284    }
1285
1286    #[rstest]
1287    #[tokio::test]
1288    async fn test_wait_for_authenticated_concurrent_waiters() {
1289        let tracker = Arc::new(AuthTracker::new());
1290        let _rx = tracker.begin();
1291
1292        let mut handles = vec![];
1293
1294        for _ in 0..10 {
1295            let t = Arc::clone(&tracker);
1296            handles.push(tokio::spawn(async move {
1297                t.wait_for_authenticated(Duration::from_secs(1)).await
1298            }));
1299        }
1300
1301        tokio::time::sleep(Duration::from_millis(50)).await;
1302        tracker.succeed();
1303
1304        for handle in handles {
1305            assert!(handle.await.unwrap());
1306        }
1307    }
1308
1309    #[rstest]
1310    #[tokio::test]
1311    async fn test_wait_for_authenticated_not_authenticated_initially() {
1312        let tracker = AuthTracker::new();
1313
1314        // Not authenticated, no begin() called, no failed flag set
1315        // Should time out
1316        assert!(
1317            !tracker
1318                .wait_for_authenticated(Duration::from_millis(50))
1319                .await
1320        );
1321    }
1322}
1323
1324#[cfg(test)]
1325#[cfg(not(all(feature = "simulation", madsim)))]
1326mod proptest_tests {
1327    use std::{sync::Arc, time::Duration};
1328
1329    use proptest::prelude::*;
1330    use rstest::rstest;
1331
1332    use super::*;
1333
1334    const AUTH_FAILED: &str = "model auth failed";
1335    const AUTH_SUPERSEDED: &str = "Authentication attempt superseded";
1336
1337    #[derive(Debug, Clone)]
1338    enum AuthTraceOp {
1339        Begin,
1340        Succeed,
1341        Fail,
1342        Invalidate,
1343        WaitForAuthenticated,
1344    }
1345
1346    #[derive(Debug, Clone, PartialEq, Eq)]
1347    enum ExpectedAuthResult {
1348        Success,
1349        Failed(&'static str),
1350    }
1351
1352    #[derive(Debug)]
1353    struct AuthTraceModel {
1354        state: AuthState,
1355        pending_receiver: Option<usize>,
1356        expected_results: Vec<Option<ExpectedAuthResult>>,
1357    }
1358
1359    impl AuthTraceModel {
1360        fn new() -> Self {
1361            Self {
1362                state: AuthState::Unauthenticated,
1363                pending_receiver: None,
1364                expected_results: Vec::new(),
1365            }
1366        }
1367
1368        fn begin(&mut self) {
1369            if let Some(receiver_index) = self.pending_receiver.take() {
1370                self.expected_results[receiver_index] =
1371                    Some(ExpectedAuthResult::Failed(AUTH_SUPERSEDED));
1372            }
1373
1374            self.pending_receiver = Some(self.expected_results.len());
1375            self.expected_results.push(None);
1376            self.state = AuthState::Unauthenticated;
1377        }
1378
1379        fn succeed(&mut self) {
1380            self.state = AuthState::Authenticated;
1381
1382            if let Some(receiver_index) = self.pending_receiver.take() {
1383                self.expected_results[receiver_index] = Some(ExpectedAuthResult::Success);
1384            }
1385        }
1386
1387        fn fail(&mut self) {
1388            self.state = AuthState::Failed;
1389
1390            if let Some(receiver_index) = self.pending_receiver.take() {
1391                self.expected_results[receiver_index] =
1392                    Some(ExpectedAuthResult::Failed(AUTH_FAILED));
1393            }
1394        }
1395
1396        fn invalidate(&mut self) {
1397            if self.state == AuthState::Authenticated {
1398                self.state = AuthState::Unauthenticated;
1399            }
1400        }
1401    }
1402
1403    fn auth_trace_op_strategy() -> impl Strategy<Value = AuthTraceOp> {
1404        prop_oneof![
1405            Just(AuthTraceOp::Begin),
1406            Just(AuthTraceOp::Succeed),
1407            Just(AuthTraceOp::Fail),
1408            Just(AuthTraceOp::Invalidate),
1409            Just(AuthTraceOp::WaitForAuthenticated),
1410        ]
1411    }
1412
1413    fn assert_auth_receivers_match_model(
1414        receivers: &mut [Option<AuthResultReceiver>],
1415        model: &AuthTraceModel,
1416        step: usize,
1417    ) -> Result<(), TestCaseError> {
1418        for (receiver_index, receiver_slot) in receivers.iter_mut().enumerate() {
1419            let Some(receiver) = receiver_slot.as_mut() else {
1420                continue;
1421            };
1422
1423            let mut clear_receiver = false;
1424
1425            match &model.expected_results[receiver_index] {
1426                Some(ExpectedAuthResult::Success) => {
1427                    match receiver.try_recv() {
1428                        Ok(Ok(())) => {}
1429                        actual => prop_assert!(
1430                            false,
1431                            "receiver {} should succeed at step {}, was {:?}",
1432                            receiver_index,
1433                            step,
1434                            actual
1435                        ),
1436                    }
1437
1438                    clear_receiver = true;
1439                }
1440                Some(ExpectedAuthResult::Failed(expected)) => {
1441                    match receiver.try_recv() {
1442                        Ok(Err(actual)) => prop_assert_eq!(
1443                            actual,
1444                            *expected,
1445                            "receiver {} should fail at step {}",
1446                            receiver_index,
1447                            step
1448                        ),
1449                        actual => prop_assert!(
1450                            false,
1451                            "receiver {} should fail at step {}, was {:?}",
1452                            receiver_index,
1453                            step,
1454                            actual
1455                        ),
1456                    }
1457
1458                    clear_receiver = true;
1459                }
1460                None => {
1461                    prop_assert_eq!(
1462                        receiver.try_recv(),
1463                        Err(tokio::sync::oneshot::error::TryRecvError::Empty),
1464                        "receiver {} should stay pending at step {}",
1465                        receiver_index,
1466                        step
1467                    );
1468                }
1469            }
1470
1471            if clear_receiver {
1472                *receiver_slot = None;
1473            }
1474        }
1475
1476        Ok(())
1477    }
1478
1479    proptest! {
1480        #![proptest_config(ProptestConfig::with_cases(256))]
1481
1482        /// Property: auth attempt traces match the cycle model for success,
1483        /// failure, superseded stale receivers, invalidation, and auth waits.
1484        #[rstest]
1485        fn test_auth_tracker_trace_matches_cycle_model(
1486            ops in proptest::collection::vec(auth_trace_op_strategy(), 1..80)
1487        ) {
1488            let runtime = tokio::runtime::Builder::new_current_thread()
1489                .enable_time()
1490                .build()
1491                .unwrap();
1492            let tracker = AuthTracker::new();
1493            let mut model = AuthTraceModel::new();
1494            let mut receivers: Vec<Option<AuthResultReceiver>> = Vec::new();
1495
1496            for (step, op) in ops.iter().enumerate() {
1497                match op {
1498                    AuthTraceOp::Begin => {
1499                        receivers.push(Some(tracker.begin()));
1500                        model.begin();
1501                    }
1502                    AuthTraceOp::Succeed => {
1503                        tracker.succeed();
1504                        model.succeed();
1505                    }
1506                    AuthTraceOp::Fail => {
1507                        tracker.fail(AUTH_FAILED);
1508                        model.fail();
1509                    }
1510                    AuthTraceOp::Invalidate => {
1511                        tracker.invalidate();
1512                        model.invalidate();
1513                    }
1514                    AuthTraceOp::WaitForAuthenticated => {
1515                        let actual = runtime
1516                            .block_on(tracker.wait_for_authenticated(Duration::from_millis(0)));
1517                        let expected = model.state == AuthState::Authenticated;
1518                        prop_assert_eq!(
1519                            actual,
1520                            expected,
1521                            "wait_for_authenticated mismatch at step {}, op {:?}",
1522                            step,
1523                            op
1524                        );
1525                    }
1526                }
1527
1528                prop_assert_eq!(
1529                    tracker.auth_state(),
1530                    model.state,
1531                    "auth state mismatch at step {}, op {:?}",
1532                    step,
1533                    op
1534                );
1535                assert_auth_receivers_match_model(&mut receivers, &model, step)?;
1536            }
1537        }
1538
1539        /// Verifies that any sequence of begin/succeed/fail/invalidate calls
1540        /// leaves the tracker in a consistent state where `is_authenticated`
1541        /// agrees with the last state-setting call.
1542        #[rstest]
1543        fn test_state_consistency_after_random_operations(
1544            ops in proptest::collection::vec(0u8..4, 1..50)
1545        ) {
1546            let tracker = AuthTracker::new();
1547            let mut expected_auth = false;
1548
1549            for op in &ops {
1550                match op {
1551                    0 => {
1552                        let _rx = tracker.begin();
1553                        expected_auth = false;
1554                    }
1555                    1 => {
1556                        tracker.succeed();
1557                        expected_auth = true;
1558                    }
1559                    2 => {
1560                        tracker.fail("test");
1561                        expected_auth = false;
1562                    }
1563                    3 => {
1564                        tracker.invalidate();
1565                        expected_auth = false;
1566                    }
1567                    _ => unreachable!(),
1568                }
1569            }
1570
1571            prop_assert_eq!(tracker.is_authenticated(), expected_auth);
1572        }
1573
1574        /// Verifies that begin() always clears the failed flag regardless of
1575        /// prior state, so a new auth attempt starts clean.
1576        #[rstest]
1577        fn test_begin_always_clears_failed(
1578            prior_ops in proptest::collection::vec(0u8..4, 0..20)
1579        ) {
1580            let tracker = AuthTracker::new();
1581
1582            for op in &prior_ops {
1583                match op {
1584                    0 => { let _rx = tracker.begin(); }
1585                    1 => tracker.succeed(),
1586                    2 => tracker.fail("test"),
1587                    3 => tracker.invalidate(),
1588                    _ => unreachable!(),
1589                }
1590            }
1591
1592            let _rx = tracker.begin();
1593            // After begin(), state is Unauthenticated
1594            prop_assert_eq!(tracker.auth_state(), AuthState::Unauthenticated);
1595        }
1596
1597        /// Verifies that succeed() always transitions to Authenticated,
1598        /// regardless of prior state.
1599        #[rstest]
1600        fn test_succeed_always_sets_authenticated(
1601            prior_ops in proptest::collection::vec(0u8..4, 0..20)
1602        ) {
1603            let tracker = AuthTracker::new();
1604
1605            for op in &prior_ops {
1606                match op {
1607                    0 => { let _rx = tracker.begin(); }
1608                    1 => tracker.succeed(),
1609                    2 => tracker.fail("test"),
1610                    3 => tracker.invalidate(),
1611                    _ => unreachable!(),
1612                }
1613            }
1614
1615            tracker.succeed();
1616            prop_assert_eq!(tracker.auth_state(), AuthState::Authenticated);
1617        }
1618    }
1619
1620    /// Verifies that `wait_for_authenticated` returns within a bounded time
1621    /// when `succeed()` or `fail()` is called, regardless of the timeout value.
1622    #[rstest]
1623    #[tokio::test]
1624    async fn test_wait_responds_within_bounded_time() {
1625        for auth_result in [true, false] {
1626            let tracker = Arc::new(AuthTracker::new());
1627            let _rx = tracker.begin();
1628
1629            let tracker_clone = Arc::clone(&tracker);
1630
1631            tokio::spawn(async move {
1632                tokio::time::sleep(Duration::from_millis(30)).await;
1633
1634                if auth_result {
1635                    tracker_clone.succeed();
1636                } else {
1637                    tracker_clone.fail("rejected");
1638                }
1639            });
1640
1641            let start = tokio::time::Instant::now();
1642            let result = tracker
1643                .wait_for_authenticated(Duration::from_secs(10))
1644                .await;
1645            let elapsed = start.elapsed();
1646
1647            assert_eq!(result, auth_result);
1648            assert!(
1649                elapsed < Duration::from_millis(500),
1650                "wait_for_authenticated took {elapsed:?} for auth_result={auth_result}"
1651            );
1652        }
1653    }
1654}
1655
1656#[cfg(all(test, feature = "simulation", madsim))]
1657mod simulation_tests {
1658    use std::time::Duration;
1659
1660    use super::*;
1661
1662    #[madsim::test]
1663    async fn test_wait_for_result_succeeds_without_tokio_reactor() {
1664        let tracker = AuthTracker::new();
1665        let rx = tracker.begin();
1666        tracker.succeed();
1667        let result: Result<(), String> = tracker.wait_for_result(Duration::from_secs(1), rx).await;
1668        assert_eq!(result, Ok(()));
1669    }
1670
1671    #[madsim::test]
1672    async fn test_wait_for_result_times_out_on_virtual_clock() {
1673        let tracker = AuthTracker::new();
1674        let rx = tracker.begin();
1675        let result: Result<(), String> =
1676            tracker.wait_for_result(Duration::from_millis(10), rx).await;
1677        assert_eq!(result, Err("Authentication timed out".to_string()));
1678    }
1679
1680    #[madsim::test]
1681    async fn test_wait_for_authenticated_succeeds_without_tokio_reactor() {
1682        let tracker = AuthTracker::new();
1683        let pending = tracker.clone();
1684
1685        madsim::task::spawn(async move {
1686            madsim::time::sleep(Duration::from_millis(1)).await;
1687            pending.succeed();
1688        });
1689
1690        assert!(tracker.wait_for_authenticated(Duration::from_secs(1)).await);
1691    }
1692}
nautilus_network/websocket/auth.rs

nautilus_network/websocket/
auth.rs
