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        *guard = Some(sender);
183
184        receiver
185    }
186
187    /// Marks the current authentication attempt as successful.
188    ///
189    /// Transitions to `Authenticated` and notifies any waiting receiver
190    /// with `Ok(())`. This should be called when the server sends a successful
191    /// authentication response.
192    ///
193    /// The state is always updated even if no receiver is waiting (e.g., after
194    /// a timeout), since the server has confirmed authentication.
195    pub fn succeed(&self) {
196        let mut guard = self.tx.lock();
197        self.state
198            .store(AuthState::Authenticated.as_u8(), Ordering::Release);
199        self.state_notify.notify_waiters();
200
201        if let Some(sender) = guard.take() {
202            let _ = sender.send(Ok(()));
203        }
204    }
205
206    /// Marks the current authentication attempt as failed.
207    ///
208    /// Transitions to `Failed` and notifies any waiting receiver
209    /// with `Err(message)`. This should be called when the server sends an
210    /// authentication error response, or on terminal client shutdown.
211    ///
212    /// The state is always updated even if no receiver is waiting, since the
213    /// server has rejected authentication or future auth is impossible.
214    pub fn fail(&self, error: impl Into<String>) {
215        let message = error.into();
216        let mut guard = self.tx.lock();
217        self.state
218            .store(AuthState::Failed.as_u8(), Ordering::Release);
219        self.state_notify.notify_waiters();
220
221        if let Some(sender) = guard.take() {
222            let _ = sender.send(Err(message));
223        }
224    }
225
226    /// Waits for the authentication result with a timeout.
227    ///
228    /// Returns `Ok(())` if authentication succeeds, or an error if it fails,
229    /// times out, or the channel is closed.
230    ///
231    /// # Type Parameters
232    ///
233    /// - `E`: Error type that implements `From<String>` for error message conversion
234    ///
235    /// # Errors
236    ///
237    /// Returns an error if:
238    /// - Authentication fails (server rejects credentials).
239    /// - Authentication times out (no response within timeout duration).
240    /// - Authentication channel closes unexpectedly.
241    /// - Authentication attempt is superseded by a new attempt.
242    pub async fn wait_for_result<E>(
243        &self,
244        timeout: Duration,
245        receiver: AuthResultReceiver,
246    ) -> Result<(), E>
247    where
248        E: From<String>,
249    {
250        match dst::time::timeout(timeout, receiver).await {
251            Ok(Ok(Ok(()))) => Ok(()),
252            Ok(Ok(Err(msg))) => Err(E::from(msg)),
253            Ok(Err(_)) => Err(E::from("Authentication channel closed".to_string())),
254            Err(_) => {
255                // Don't clear the sender: a concurrent begin() may have replaced it,
256                // and guard.take() would cancel the newer sender. The next begin()
257                // call cleans up any stale sender.
258                Err(E::from("Authentication timed out".to_string()))
259            }
260        }
261    }
262
263    /// Waits for the tracker to enter the authenticated state.
264    ///
265    /// Returns `true` if authenticated within the timeout, `false` if the timeout
266    /// expires or authentication explicitly fails. Uses event-driven notification
267    /// from `succeed()` / `fail()` / `invalidate()` to avoid polling.
268    ///
269    /// Returns early with `false` when `fail()` is called (e.g., the exchange
270    /// rejects credentials), so callers are not blocked for the full timeout
271    /// on a definitive auth rejection.
272    ///
273    /// This is intended for callers on a separate task who need to gate operations
274    /// on authentication state (e.g., order sends that must wait for re-authentication
275    /// after a WebSocket reconnection).
276    pub async fn wait_for_authenticated(&self, timeout: Duration) -> bool {
277        if self.is_authenticated() {
278            return true;
279        }
280
281        dst::time::timeout(timeout, async {
282            loop {
283                // Enable before the state check: an unpolled Notified is unregistered and misses notifies
284                let mut notified = pin!(self.state_notify.notified());
285                notified.as_mut().enable();
286
287                match self.auth_state() {
288                    AuthState::Authenticated => return true,
289                    AuthState::Failed => return false,
290                    AuthState::Unauthenticated => notified.await,
291                }
292            }
293        })
294        .await
295        .unwrap_or(false)
296    }
297}
298
299impl Default for AuthTracker {
300    fn default() -> Self {
301        Self::new()
302    }
303}
304
305#[cfg(test)]
306#[cfg(not(all(feature = "simulation", madsim)))]
307mod tests {
308    use std::{
309        sync::{
310            atomic::{AtomicBool, Ordering},
311            mpsc,
312        },
313        time::Duration,
314    };
315
316    use rstest::rstest;
317
318    use super::*;
319
320    #[derive(Debug, PartialEq)]
321    struct TestError(String);
322
323    impl From<String> for TestError {
324        fn from(msg: String) -> Self {
325            Self(msg)
326        }
327    }
328
329    #[rstest]
330    fn test_begin_during_failure_keeps_receiver_and_state_consistent() {
331        struct DeferredError {
332            entered: mpsc::SyncSender<()>,
333            proceed: mpsc::Receiver<()>,
334        }
335
336        impl From<DeferredError> for String {
337            fn from(error: DeferredError) -> Self {
338                error.entered.send(()).unwrap();
339                error.proceed.recv_timeout(Duration::from_secs(2)).unwrap();
340                "credentials rejected".to_string()
341            }
342        }
343
344        let tracker = AuthTracker::new();
345        let (entered_tx, entered_rx) = mpsc::sync_channel(1);
346        let (proceed_tx, proceed_rx) = mpsc::sync_channel(1);
347
348        let mut receiver = std::thread::scope(|scope| {
349            let completion = scope.spawn(|| {
350                tracker.fail(DeferredError {
351                    entered: entered_tx,
352                    proceed: proceed_rx,
353                });
354            });
355
356            entered_rx.recv_timeout(Duration::from_secs(2)).unwrap();
357            let receiver = tracker.begin();
358            proceed_tx.send(()).unwrap();
359            completion.join().unwrap();
360            receiver
361        });
362
363        assert_eq!(
364            receiver.try_recv(),
365            Ok(Err("credentials rejected".to_string()))
366        );
367        assert_eq!(tracker.auth_state(), AuthState::Failed);
368    }
369
370    #[rstest]
371    #[case::success(AuthTracker::succeed, Ok(()), AuthState::Authenticated)]
372    #[case::failure(|tracker: &AuthTracker| tracker.fail("rejected"), Err("rejected".to_string()), AuthState::Failed)]
373    fn test_concurrent_begin_and_completion_keep_receiver_and_state_consistent(
374        #[case] complete: fn(&AuthTracker),
375        #[case] result: Result<(), String>,
376        #[case] state: AuthState,
377    ) {
378        const ATTEMPTS: usize = 2_000;
379
380        let tracker = AuthTracker::new();
381        let (begin_tx, begin_rx) = mpsc::sync_channel(1);
382        let (complete_tx, complete_rx) = mpsc::sync_channel(1);
383        let (receiver_tx, receiver_rx) = mpsc::sync_channel(1);
384        let (done_tx, done_rx) = mpsc::sync_channel(1);
385
386        let observations = std::thread::scope(|scope| {
387            let tracker = &tracker;
388
389            scope.spawn(move || {
390                for _ in 0..ATTEMPTS {
391                    begin_rx.recv_timeout(Duration::from_secs(2)).unwrap();
392                    receiver_tx.send(tracker.begin()).unwrap();
393                }
394            });
395
396            scope.spawn(move || {
397                for _ in 0..ATTEMPTS {
398                    complete_rx.recv_timeout(Duration::from_secs(2)).unwrap();
399                    complete(tracker);
400                    done_tx.send(()).unwrap();
401                }
402            });
403
404            let mut observations = Vec::with_capacity(ATTEMPTS);
405
406            for _ in 0..ATTEMPTS {
407                begin_tx.send(()).unwrap();
408                complete_tx.send(()).unwrap();
409                let mut receiver = receiver_rx.recv_timeout(Duration::from_secs(2)).unwrap();
410                done_rx.recv_timeout(Duration::from_secs(2)).unwrap();
411                observations.push((receiver.try_recv(), tracker.auth_state()));
412            }
413
414            observations
415        });
416
417        for (received, actual_state) in observations {
418            match received {
419                Ok(actual_result) => {
420                    assert_eq!(actual_result, result);
421                    assert_eq!(actual_state, state);
422                }
423                Err(tokio::sync::oneshot::error::TryRecvError::Empty) => {
424                    assert_eq!(actual_state, AuthState::Unauthenticated);
425                }
426                Err(receive_error) => panic!("authentication channel closed: {receive_error}"),
427            }
428        }
429    }
430
431    #[rstest]
432    #[tokio::test]
433    async fn test_successful_authentication() {
434        let tracker = AuthTracker::new();
435        let rx = tracker.begin();
436
437        tracker.succeed();
438
439        let result: Result<(), TestError> =
440            tracker.wait_for_result(Duration::from_secs(1), rx).await;
441
442        assert!(result.is_ok());
443    }
444
445    #[rstest]
446    #[tokio::test]
447    async fn test_failed_authentication() {
448        let tracker = AuthTracker::new();
449        let rx = tracker.begin();
450
451        tracker.fail("Invalid credentials");
452
453        let result: Result<(), TestError> =
454            tracker.wait_for_result(Duration::from_secs(1), rx).await;
455
456        assert_eq!(
457            result.unwrap_err(),
458            TestError("Invalid credentials".to_string())
459        );
460    }
461
462    #[rstest]
463    #[tokio::test]
464    async fn test_authentication_timeout() {
465        let tracker = AuthTracker::new();
466        let rx = tracker.begin();
467
468        // Don't call succeed or fail - let it timeout
469
470        let result: Result<(), TestError> =
471            tracker.wait_for_result(Duration::from_millis(50), rx).await;
472
473        assert_eq!(
474            result.unwrap_err(),
475            TestError("Authentication timed out".to_string())
476        );
477    }
478
479    #[rstest]
480    #[tokio::test]
481    async fn test_begin_supersedes_previous_sender() {
482        let tracker = AuthTracker::new();
483
484        let first = tracker.begin();
485        let second = tracker.begin();
486
487        // First receiver should get superseded error
488        let result = first.await.expect("oneshot closed unexpectedly");
489        assert_eq!(result, Err("Authentication attempt superseded".to_string()));
490
491        // Second attempt should succeed
492        tracker.succeed();
493        let result: Result<(), TestError> = tracker
494            .wait_for_result(Duration::from_secs(1), second)
495            .await;
496
497        assert!(result.is_ok());
498    }
499
500    #[rstest]
501    #[tokio::test]
502    async fn test_succeed_without_pending_auth() {
503        let tracker = AuthTracker::new();
504
505        tracker.succeed();
506
507        assert_eq!(tracker.auth_state(), AuthState::Authenticated);
508    }
509
510    #[rstest]
511    #[tokio::test]
512    async fn test_fail_without_pending_auth() {
513        let tracker = AuthTracker::new();
514
515        tracker.fail("Some error");
516
517        assert_eq!(tracker.auth_state(), AuthState::Failed);
518    }
519
520    #[rstest]
521    #[tokio::test]
522    async fn test_multiple_sequential_authentications() {
523        let tracker = AuthTracker::new();
524
525        // First auth succeeds
526        let rx1 = tracker.begin();
527        tracker.succeed();
528        let result1: Result<(), TestError> =
529            tracker.wait_for_result(Duration::from_secs(1), rx1).await;
530        assert!(result1.is_ok());
531
532        // Second auth fails
533        let rx2 = tracker.begin();
534        tracker.fail("Credentials expired");
535        let result2: Result<(), TestError> =
536            tracker.wait_for_result(Duration::from_secs(1), rx2).await;
537        assert_eq!(
538            result2.unwrap_err(),
539            TestError("Credentials expired".to_string())
540        );
541
542        // Third auth succeeds
543        let rx3 = tracker.begin();
544        tracker.succeed();
545        let result3: Result<(), TestError> =
546            tracker.wait_for_result(Duration::from_secs(1), rx3).await;
547        assert!(result3.is_ok());
548    }
549
550    #[rstest]
551    #[tokio::test]
552    async fn test_channel_closed_before_result() {
553        let tracker = AuthTracker::new();
554        let rx = tracker.begin();
555
556        drop(tracker.tx.lock().take());
557
558        let result: Result<(), TestError> =
559            tracker.wait_for_result(Duration::from_secs(1), rx).await;
560
561        assert_eq!(
562            result,
563            Err(TestError("Authentication channel closed".to_string()))
564        );
565        assert_eq!(tracker.auth_state(), AuthState::Unauthenticated);
566    }
567
568    #[rstest]
569    #[tokio::test]
570    async fn test_concurrent_auth_attempts() {
571        let tracker = Arc::new(AuthTracker::new());
572        let mut handles = vec![];
573
574        // Spawn 10 concurrent auth attempts
575        for i in 0..10 {
576            let tracker_clone = Arc::clone(&tracker);
577            let handle = tokio::spawn(async move {
578                let rx = tracker_clone.begin();
579
580                // Only the last one should succeed
581                if i == 9 {
582                    tokio::time::sleep(Duration::from_millis(10)).await;
583                    tracker_clone.succeed();
584                }
585
586                let result: Result<(), TestError> = tracker_clone
587                    .wait_for_result(Duration::from_secs(1), rx)
588                    .await;
589
590                (i, result)
591            });
592            handles.push(handle);
593        }
594
595        let mut successes = 0;
596        let mut superseded = 0;
597
598        for handle in handles {
599            let (i, result) = handle.await.unwrap();
600            match result {
601                Ok(()) => {
602                    // Only task 9 should succeed
603                    assert_eq!(i, 9);
604                    successes += 1;
605                }
606                Err(TestError(msg)) if msg.contains("superseded") => {
607                    superseded += 1;
608                }
609                Err(e) => panic!("Unexpected error: {e:?}"),
610            }
611        }
612
613        assert_eq!(successes, 1);
614        assert_eq!(superseded, 9);
615    }
616
617    #[rstest]
618    fn test_default_trait() {
619        let _tracker = AuthTracker::default();
620    }
621
622    #[rstest]
623    #[tokio::test]
624    async fn test_clone_trait() {
625        let tracker = AuthTracker::new();
626        let cloned = tracker.clone();
627
628        // Verify cloned instance shares state with original (Arc behavior)
629        let rx = tracker.begin();
630        cloned.succeed(); // Succeed via clone affects original
631        let result: Result<(), TestError> =
632            tracker.wait_for_result(Duration::from_secs(1), rx).await;
633        assert!(result.is_ok());
634    }
635
636    #[rstest]
637    fn test_debug_trait() {
638        let tracker = AuthTracker::new();
639        let debug_str = format!("{tracker:?}");
640        assert!(debug_str.contains("AuthTracker"));
641    }
642
643    #[rstest]
644    #[tokio::test]
645    async fn test_timeout_clears_sender() {
646        let tracker = AuthTracker::new();
647
648        // Start auth that will timeout
649        let rx1 = tracker.begin();
650        let result1: Result<(), TestError> = tracker
651            .wait_for_result(Duration::from_millis(50), rx1)
652            .await;
653        assert_eq!(
654            result1.unwrap_err(),
655            TestError("Authentication timed out".to_string())
656        );
657
658        // Verify sender was cleared - new auth should work
659        let rx2 = tracker.begin();
660        tracker.succeed();
661        let result2: Result<(), TestError> =
662            tracker.wait_for_result(Duration::from_secs(1), rx2).await;
663        assert!(result2.is_ok());
664    }
665
666    #[rstest]
667    #[tokio::test]
668    async fn test_fail_clears_sender() {
669        let tracker = AuthTracker::new();
670
671        // Auth fails
672        let rx1 = tracker.begin();
673        tracker.fail("Bad credentials");
674        let result1: Result<(), TestError> =
675            tracker.wait_for_result(Duration::from_secs(1), rx1).await;
676        assert!(result1.is_err());
677
678        // Verify sender was cleared - new auth should work
679        let rx2 = tracker.begin();
680        tracker.succeed();
681        let result2: Result<(), TestError> =
682            tracker.wait_for_result(Duration::from_secs(1), rx2).await;
683        assert!(result2.is_ok());
684    }
685
686    #[rstest]
687    #[tokio::test]
688    async fn test_succeed_clears_sender() {
689        let tracker = AuthTracker::new();
690
691        // Auth succeeds
692        let rx1 = tracker.begin();
693        tracker.succeed();
694        let result1: Result<(), TestError> =
695            tracker.wait_for_result(Duration::from_secs(1), rx1).await;
696        assert!(result1.is_ok());
697
698        // Verify sender was cleared - new auth should work
699        let rx2 = tracker.begin();
700        tracker.succeed();
701        let result2: Result<(), TestError> =
702            tracker.wait_for_result(Duration::from_secs(1), rx2).await;
703        assert!(result2.is_ok());
704    }
705
706    #[rstest]
707    #[tokio::test]
708    async fn test_rapid_begin_succeed_cycles() {
709        let tracker = AuthTracker::new();
710
711        // Rapidly cycle through auth attempts
712        for _ in 0..100 {
713            let rx = tracker.begin();
714            tracker.succeed();
715            let result: Result<(), TestError> =
716                tracker.wait_for_result(Duration::from_secs(1), rx).await;
717            assert!(result.is_ok());
718        }
719    }
720
721    #[rstest]
722    #[tokio::test]
723    async fn test_double_succeed_is_safe() {
724        let tracker = AuthTracker::new();
725        let rx = tracker.begin();
726
727        // Call succeed twice
728        tracker.succeed();
729        tracker.succeed(); // Second call should be no-op
730
731        let result: Result<(), TestError> =
732            tracker.wait_for_result(Duration::from_secs(1), rx).await;
733        assert!(result.is_ok());
734    }
735
736    #[rstest]
737    #[tokio::test]
738    async fn test_double_fail_is_safe() {
739        let tracker = AuthTracker::new();
740        let rx = tracker.begin();
741
742        // Call fail twice
743        tracker.fail("Error 1");
744        tracker.fail("Error 2"); // Second call should be no-op
745
746        let result: Result<(), TestError> =
747            tracker.wait_for_result(Duration::from_secs(1), rx).await;
748        assert_eq!(
749            result.unwrap_err(),
750            TestError("Error 1".to_string()) // Should be first error
751        );
752    }
753
754    #[rstest]
755    #[tokio::test]
756    async fn test_succeed_after_fail_is_ignored() {
757        let tracker = AuthTracker::new();
758        let rx = tracker.begin();
759
760        tracker.fail("Auth failed");
761        tracker.succeed(); // This should be no-op
762
763        let result: Result<(), TestError> =
764            tracker.wait_for_result(Duration::from_secs(1), rx).await;
765        assert!(result.is_err()); // Should still be error
766    }
767
768    #[rstest]
769    #[tokio::test]
770    async fn test_fail_after_succeed_is_ignored() {
771        let tracker = AuthTracker::new();
772        let rx = tracker.begin();
773
774        tracker.succeed();
775        tracker.fail("Auth failed"); // This should be no-op
776
777        let result: Result<(), TestError> =
778            tracker.wait_for_result(Duration::from_secs(1), rx).await;
779        assert!(result.is_ok()); // Should still be success
780    }
781
782    /// Simulates a reconnect flow where authentication must complete before resubscription.
783    ///
784    /// This is an integration-style test that verifies:
785    /// 1. On reconnect, authentication starts first
786    /// 2. Subscription logic waits for auth to complete
787    /// 3. Subscriptions only proceed after successful auth
788    #[rstest]
789    #[tokio::test]
790    async fn test_reconnect_flow_waits_for_auth() {
791        let tracker = Arc::new(AuthTracker::new());
792        let subscribed = Arc::new(tokio::sync::Notify::new());
793        let auth_completed = Arc::new(tokio::sync::Notify::new());
794
795        // Simulate reconnect handler
796        let tracker_reconnect = Arc::clone(&tracker);
797        let subscribed_reconnect = Arc::clone(&subscribed);
798        let auth_completed_reconnect = Arc::clone(&auth_completed);
799
800        let reconnect_task = tokio::spawn(async move {
801            // Step 1: Begin authentication
802            let rx = tracker_reconnect.begin();
803
804            // Step 2: Spawn resubscription task that waits for auth
805            let tracker_resub = Arc::clone(&tracker_reconnect);
806            let subscribed_resub = Arc::clone(&subscribed_reconnect);
807            let auth_completed_resub = Arc::clone(&auth_completed_reconnect);
808
809            let resub_task = tokio::spawn(async move {
810                // Wait for auth to complete
811                let result: Result<(), TestError> = tracker_resub
812                    .wait_for_result(Duration::from_secs(5), rx)
813                    .await;
814
815                if result.is_ok() {
816                    auth_completed_resub.notify_one();
817                    // Simulate resubscription
818                    tokio::time::sleep(Duration::from_millis(10)).await;
819                    subscribed_resub.notify_one();
820                }
821            });
822
823            resub_task.await.unwrap();
824        });
825
826        // Simulate server auth response after delay
827        tokio::time::sleep(Duration::from_millis(100)).await;
828        tracker.succeed();
829
830        // Wait for reconnect flow to complete
831        reconnect_task.await.unwrap();
832
833        // Verify auth completed before subscription
834        tokio::select! {
835            () = auth_completed.notified() => {
836                // Good - auth completed
837            }
838            () = tokio::time::sleep(Duration::from_secs(1)) => {
839                panic!("Auth never completed");
840            }
841        }
842
843        // Verify subscription completed
844        tokio::select! {
845            () = subscribed.notified() => {
846                // Good - subscribed
847            }
848            () = tokio::time::sleep(Duration::from_secs(1)) => {
849                panic!("Subscription never completed");
850            }
851        }
852    }
853
854    /// Verifies that failed authentication prevents resubscription in reconnect flow.
855    #[rstest]
856    #[tokio::test]
857    async fn test_reconnect_flow_blocks_on_auth_failure() {
858        let tracker = Arc::new(AuthTracker::new());
859        let subscribed = Arc::new(AtomicBool::new(false));
860
861        let tracker_reconnect = Arc::clone(&tracker);
862        let subscribed_reconnect = Arc::clone(&subscribed);
863
864        let reconnect_task = tokio::spawn(async move {
865            let rx = tracker_reconnect.begin();
866
867            // Spawn resubscription task that waits for auth
868            let tracker_resub = Arc::clone(&tracker_reconnect);
869            let subscribed_resub = Arc::clone(&subscribed_reconnect);
870
871            let resub_task = tokio::spawn(async move {
872                let result: Result<(), TestError> = tracker_resub
873                    .wait_for_result(Duration::from_secs(5), rx)
874                    .await;
875
876                // Only subscribe if auth succeeds
877                if result.is_ok() {
878                    subscribed_resub.store(true, Ordering::Relaxed);
879                }
880            });
881
882            resub_task.await.unwrap();
883        });
884
885        // Simulate server auth failure
886        tokio::time::sleep(Duration::from_millis(50)).await;
887        tracker.fail("Invalid credentials");
888
889        // Wait for reconnect flow to complete
890        reconnect_task.await.unwrap();
891
892        // Verify subscription never happened
893        tokio::time::sleep(Duration::from_millis(100)).await;
894        assert!(!subscribed.load(Ordering::Relaxed));
895    }
896
897    /// Tests state machine transitions exhaustively.
898    #[rstest]
899    #[tokio::test]
900    async fn test_state_machine_transitions() {
901        let tracker = AuthTracker::new();
902
903        // Transition 1: Initial -> Pending (begin)
904        let rx1 = tracker.begin();
905
906        // Transition 2: Pending -> Success (succeed)
907        tracker.succeed();
908        let result1: Result<(), TestError> =
909            tracker.wait_for_result(Duration::from_secs(1), rx1).await;
910        assert!(result1.is_ok());
911
912        // Transition 3: Success -> Pending (begin again)
913        let rx2 = tracker.begin();
914
915        // Transition 4: Pending -> Failure (fail)
916        tracker.fail("Error");
917        let result2: Result<(), TestError> =
918            tracker.wait_for_result(Duration::from_secs(1), rx2).await;
919        assert!(result2.is_err());
920
921        // Transition 5: Failure -> Pending (begin again)
922        let rx3 = tracker.begin();
923
924        // Transition 6: Pending -> Timeout
925        let result3: Result<(), TestError> = tracker
926            .wait_for_result(Duration::from_millis(50), rx3)
927            .await;
928        assert_eq!(
929            result3.unwrap_err(),
930            TestError("Authentication timed out".to_string())
931        );
932
933        // Transition 7: Timeout -> Pending (begin again)
934        let rx4 = tracker.begin();
935
936        // Transition 8: Pending -> Superseded (begin interrupts)
937        let rx5 = tracker.begin();
938        let result4: Result<(), TestError> =
939            tracker.wait_for_result(Duration::from_secs(1), rx4).await;
940        assert_eq!(
941            result4.unwrap_err(),
942            TestError("Authentication attempt superseded".to_string())
943        );
944
945        // Final success to clean up
946        tracker.succeed();
947        let result5: Result<(), TestError> =
948            tracker.wait_for_result(Duration::from_secs(1), rx5).await;
949        assert!(result5.is_ok());
950    }
951
952    /// Verifies no memory leaks from orphaned senders.
953    #[rstest]
954    #[tokio::test]
955    async fn test_no_sender_leaks() {
956        let tracker = AuthTracker::new();
957
958        for _ in 0..100 {
959            let rx = tracker.begin();
960            let _result: Result<(), TestError> =
961                tracker.wait_for_result(Duration::from_millis(1), rx).await;
962        }
963
964        let rx = tracker.begin();
965        tracker.succeed();
966        let result: Result<(), TestError> =
967            tracker.wait_for_result(Duration::from_secs(1), rx).await;
968        assert!(result.is_ok());
969    }
970
971    /// Tests concurrent success/fail calls don't cause panics.
972    #[rstest]
973    #[tokio::test]
974    async fn test_concurrent_succeed_fail_calls() {
975        let tracker = Arc::new(AuthTracker::new());
976        let rx = tracker.begin();
977
978        let mut handles = vec![];
979
980        // Spawn many tasks trying to succeed
981        for _ in 0..50 {
982            let tracker_clone = Arc::clone(&tracker);
983            handles.push(tokio::spawn(async move {
984                tracker_clone.succeed();
985            }));
986        }
987
988        // Spawn many tasks trying to fail
989        for _ in 0..50 {
990            let tracker_clone = Arc::clone(&tracker);
991            handles.push(tokio::spawn(async move {
992                tracker_clone.fail("Error");
993            }));
994        }
995
996        // Wait for all tasks
997        for handle in handles {
998            handle.await.unwrap();
999        }
1000
1001        // Should get either success or failure, but not panic
1002        let result: Result<(), TestError> =
1003            tracker.wait_for_result(Duration::from_secs(1), rx).await;
1004        // Don't care which outcome, just that it doesn't panic
1005        let _ = result;
1006    }
1007
1008    #[rstest]
1009    fn test_is_authenticated_initial_state() {
1010        let tracker = AuthTracker::new();
1011        assert!(!tracker.is_authenticated());
1012    }
1013
1014    #[rstest]
1015    #[tokio::test]
1016    async fn test_is_authenticated_after_succeed() {
1017        let tracker = AuthTracker::new();
1018        assert!(!tracker.is_authenticated());
1019
1020        let _rx = tracker.begin();
1021        assert!(!tracker.is_authenticated());
1022
1023        tracker.succeed();
1024        assert!(tracker.is_authenticated());
1025    }
1026
1027    #[rstest]
1028    #[tokio::test]
1029    async fn test_is_authenticated_after_fail() {
1030        let tracker = AuthTracker::new();
1031        let _rx = tracker.begin();
1032        tracker.fail("error");
1033        assert!(!tracker.is_authenticated());
1034    }
1035
1036    #[rstest]
1037    #[tokio::test]
1038    async fn test_invalidate_clears_auth_state() {
1039        let tracker = AuthTracker::new();
1040        let _rx = tracker.begin();
1041        tracker.succeed();
1042        assert!(tracker.is_authenticated());
1043
1044        tracker.invalidate();
1045        assert!(!tracker.is_authenticated());
1046    }
1047
1048    #[rstest]
1049    #[case(true)]
1050    #[case(false)]
1051    #[tokio::test]
1052    async fn test_invalidate_preserves_terminal_failure(#[case] invalidate_first: bool) {
1053        let tracker = AuthTracker::new();
1054        let receiver = tracker.begin();
1055
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                    clear_receiver = true;
1438                }
1439                Some(ExpectedAuthResult::Failed(expected)) => {
1440                    match receiver.try_recv() {
1441                        Ok(Err(actual)) => prop_assert_eq!(
1442                            actual,
1443                            *expected,
1444                            "receiver {} should fail at step {}",
1445                            receiver_index,
1446                            step
1447                        ),
1448                        actual => prop_assert!(
1449                            false,
1450                            "receiver {} should fail at step {}, was {:?}",
1451                            receiver_index,
1452                            step,
1453                            actual
1454                        ),
1455                    }
1456                    clear_receiver = true;
1457                }
1458                None => {
1459                    prop_assert_eq!(
1460                        receiver.try_recv(),
1461                        Err(tokio::sync::oneshot::error::TryRecvError::Empty),
1462                        "receiver {} should stay pending at step {}",
1463                        receiver_index,
1464                        step
1465                    );
1466                }
1467            }
1468
1469            if clear_receiver {
1470                *receiver_slot = None;
1471            }
1472        }
1473
1474        Ok(())
1475    }
1476
1477    proptest! {
1478        #![proptest_config(ProptestConfig::with_cases(256))]
1479
1480        /// Property: auth attempt traces match the cycle model for success,
1481        /// failure, superseded stale receivers, invalidation, and auth waits.
1482        #[rstest]
1483        fn test_auth_tracker_trace_matches_cycle_model(
1484            ops in proptest::collection::vec(auth_trace_op_strategy(), 1..80)
1485        ) {
1486            let runtime = tokio::runtime::Builder::new_current_thread()
1487                .enable_time()
1488                .build()
1489                .unwrap();
1490            let tracker = AuthTracker::new();
1491            let mut model = AuthTraceModel::new();
1492            let mut receivers: Vec<Option<AuthResultReceiver>> = Vec::new();
1493
1494            for (step, op) in ops.iter().enumerate() {
1495                match op {
1496                    AuthTraceOp::Begin => {
1497                        receivers.push(Some(tracker.begin()));
1498                        model.begin();
1499                    }
1500                    AuthTraceOp::Succeed => {
1501                        tracker.succeed();
1502                        model.succeed();
1503                    }
1504                    AuthTraceOp::Fail => {
1505                        tracker.fail(AUTH_FAILED);
1506                        model.fail();
1507                    }
1508                    AuthTraceOp::Invalidate => {
1509                        tracker.invalidate();
1510                        model.invalidate();
1511                    }
1512                    AuthTraceOp::WaitForAuthenticated => {
1513                        let actual = runtime
1514                            .block_on(tracker.wait_for_authenticated(Duration::from_millis(0)));
1515                        let expected = model.state == AuthState::Authenticated;
1516                        prop_assert_eq!(
1517                            actual,
1518                            expected,
1519                            "wait_for_authenticated mismatch at step {}, op {:?}",
1520                            step,
1521                            op
1522                        );
1523                    }
1524                }
1525
1526                prop_assert_eq!(
1527                    tracker.auth_state(),
1528                    model.state,
1529                    "auth state mismatch at step {}, op {:?}",
1530                    step,
1531                    op
1532                );
1533                assert_auth_receivers_match_model(&mut receivers, &model, step)?;
1534            }
1535        }
1536
1537        /// Verifies that any sequence of begin/succeed/fail/invalidate calls
1538        /// leaves the tracker in a consistent state where `is_authenticated`
1539        /// agrees with the last state-setting call.
1540        #[rstest]
1541        fn test_state_consistency_after_random_operations(
1542            ops in proptest::collection::vec(0u8..4, 1..50)
1543        ) {
1544            let tracker = AuthTracker::new();
1545            let mut expected_auth = false;
1546
1547            for op in &ops {
1548                match op {
1549                    0 => {
1550                        let _rx = tracker.begin();
1551                        expected_auth = false;
1552                    }
1553                    1 => {
1554                        tracker.succeed();
1555                        expected_auth = true;
1556                    }
1557                    2 => {
1558                        tracker.fail("test");
1559                        expected_auth = false;
1560                    }
1561                    3 => {
1562                        tracker.invalidate();
1563                        expected_auth = false;
1564                    }
1565                    _ => unreachable!(),
1566                }
1567            }
1568
1569            prop_assert_eq!(tracker.is_authenticated(), expected_auth);
1570        }
1571
1572        /// Verifies that begin() always clears the failed flag regardless of
1573        /// prior state, so a new auth attempt starts clean.
1574        #[rstest]
1575        fn test_begin_always_clears_failed(
1576            prior_ops in proptest::collection::vec(0u8..4, 0..20)
1577        ) {
1578            let tracker = AuthTracker::new();
1579
1580            for op in &prior_ops {
1581                match op {
1582                    0 => { let _rx = tracker.begin(); }
1583                    1 => tracker.succeed(),
1584                    2 => tracker.fail("test"),
1585                    3 => tracker.invalidate(),
1586                    _ => unreachable!(),
1587                }
1588            }
1589
1590            let _rx = tracker.begin();
1591            // After begin(), state is Unauthenticated
1592            prop_assert_eq!(tracker.auth_state(), AuthState::Unauthenticated);
1593        }
1594
1595        /// Verifies that succeed() always transitions to Authenticated,
1596        /// regardless of prior state.
1597        #[rstest]
1598        fn test_succeed_always_sets_authenticated(
1599            prior_ops in proptest::collection::vec(0u8..4, 0..20)
1600        ) {
1601            let tracker = AuthTracker::new();
1602
1603            for op in &prior_ops {
1604                match op {
1605                    0 => { let _rx = tracker.begin(); }
1606                    1 => tracker.succeed(),
1607                    2 => tracker.fail("test"),
1608                    3 => tracker.invalidate(),
1609                    _ => unreachable!(),
1610                }
1611            }
1612
1613            tracker.succeed();
1614            prop_assert_eq!(tracker.auth_state(), AuthState::Authenticated);
1615        }
1616    }
1617
1618    /// Verifies that `wait_for_authenticated` returns within a bounded time
1619    /// when `succeed()` or `fail()` is called, regardless of the timeout value.
1620    #[rstest]
1621    #[tokio::test]
1622    async fn test_wait_responds_within_bounded_time() {
1623        for auth_result in [true, false] {
1624            let tracker = Arc::new(AuthTracker::new());
1625            let _rx = tracker.begin();
1626
1627            let tracker_clone = Arc::clone(&tracker);
1628
1629            tokio::spawn(async move {
1630                tokio::time::sleep(Duration::from_millis(30)).await;
1631
1632                if auth_result {
1633                    tracker_clone.succeed();
1634                } else {
1635                    tracker_clone.fail("rejected");
1636                }
1637            });
1638
1639            let start = tokio::time::Instant::now();
1640            let result = tracker
1641                .wait_for_authenticated(Duration::from_secs(10))
1642                .await;
1643            let elapsed = start.elapsed();
1644
1645            assert_eq!(result, auth_result);
1646            assert!(
1647                elapsed < Duration::from_millis(500),
1648                "wait_for_authenticated took {elapsed:?} for auth_result={auth_result}"
1649            );
1650        }
1651    }
1652}
1653
1654#[cfg(all(test, feature = "simulation", madsim))]
1655mod simulation_tests {
1656    use std::time::Duration;
1657
1658    use super::*;
1659
1660    #[madsim::test]
1661    async fn test_wait_for_result_succeeds_without_tokio_reactor() {
1662        let tracker = AuthTracker::new();
1663        let rx = tracker.begin();
1664        tracker.succeed();
1665        let result: Result<(), String> = tracker.wait_for_result(Duration::from_secs(1), rx).await;
1666        assert_eq!(result, Ok(()));
1667    }
1668
1669    #[madsim::test]
1670    async fn test_wait_for_result_times_out_on_virtual_clock() {
1671        let tracker = AuthTracker::new();
1672        let rx = tracker.begin();
1673        let result: Result<(), String> =
1674            tracker.wait_for_result(Duration::from_millis(10), rx).await;
1675        assert_eq!(result, Err("Authentication timed out".to_string()));
1676    }
1677
1678    #[madsim::test]
1679    async fn test_wait_for_authenticated_succeeds_without_tokio_reactor() {
1680        let tracker = AuthTracker::new();
1681        let pending = tracker.clone();
1682
1683        madsim::task::spawn(async move {
1684            madsim::time::sleep(Duration::from_millis(1)).await;
1685            pending.succeed();
1686        });
1687        assert!(tracker.wait_for_authenticated(Duration::from_secs(1)).await);
1688    }
1689}
nautilus_network/websocket/auth.rs

nautilus_network/websocket/
auth.rs
