nautilus_network/websocket/

auth.rs

// -------------------------------------------------------------------------------------------------
//  Copyright (C) 2015-2026 Nautech Systems Pty Ltd. All rights reserved.
//  https://nautechsystems.io
//
//  Licensed under the GNU Lesser General Public License Version 3.0 (the "License");
//  You may not use this file except in compliance with the License.
//  You may obtain a copy of the License at https://www.gnu.org/licenses/lgpl-3.0.en.html
//
//  Unless required by applicable law or agreed to in writing, software
//  distributed under the License is distributed on an "AS IS" BASIS,
//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//  See the License for the specific language governing permissions and
//  limitations under the License.
// -------------------------------------------------------------------------------------------------

//! Adapter authentication state independent of the WebSocket transport state.
//!
//! [`AuthTracker`] separates a specific authentication attempt from the shared session state.
//! [`AuthTracker::begin`] returns a oneshot receiver for the attempt and fails any earlier pending
//! attempt as superseded. [`AuthTracker::succeed`] and [`AuthTracker::fail`] resolve the active
//! attempt and wake state waiters. [`AuthTracker::invalidate`] returns an authenticated session to
//! unauthenticated without resolving a pending attempt or clearing terminal failure.
//!
//! # Client integration
//!
//! Registering a tracker with the client invalidates it on reconnectable connection loss and fails
//! it on terminal shutdown. When authentication-gated replay is enabled, ordinary buffered sends
//! wait for `Authenticated` and are discarded on `Failed`. The adapter remains responsible for
//! sending authentication, interpreting the response, and ordering resubscription.

use std::{
    pin::pin,
    sync::{
        Arc,
        atomic::{AtomicU8, Ordering},
    },
    time::Duration,
};

use parking_lot::Mutex;

use crate::dst;

pub type AuthResultSender = tokio::sync::oneshot::Sender<Result<(), String>>;
pub type AuthResultReceiver = tokio::sync::oneshot::Receiver<Result<(), String>>;

/// Authentication state for a WebSocket session.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
#[repr(u8)]
pub enum AuthState {
    /// Not authenticated (initial state, after begin, or after invalidating authenticated state).
    #[default]
    Unauthenticated = 0,
    /// Successfully authenticated (after succeed).
    Authenticated = 1,
    /// Authentication failed or became impossible (after fail).
    Failed = 2,
}

impl AuthState {
    #[inline]
    #[must_use]
    #[expect(
        clippy::match_same_arms,
        reason = "explicit variant listing is clearer than collapsing 0 with wildcard"
    )]
    fn from_u8(value: u8) -> Self {
        match value {
            0 => Self::Unauthenticated,
            1 => Self::Authenticated,
            2 => Self::Failed,
            _ => Self::Unauthenticated,
        }
    }

    #[inline]
    #[must_use]
    const fn as_u8(self) -> u8 {
        self as u8
    }
}

/// Tracks authentication state for WebSocket connections.
///
/// Each authentication attempt receives a dedicated oneshot channel that resolves when the server
/// responds.
///
/// # State management
///
/// The tracker maintains three states:
///
/// - [`AuthState::Unauthenticated`]: The initial state, the state after [`Self::begin`], and the
///   result of [`Self::invalidate`] from [`AuthState::Authenticated`].
/// - [`AuthState::Authenticated`]: The state after [`Self::succeed`].
/// - [`AuthState::Failed`]: The state after [`Self::fail`]. Authentication waiters return early in
///   this state.
///
/// # Superseding behavior
///
/// If a new authentication attempt begins while another remains pending, the old attempt is
/// cancelled with an error. This prevents responses from an earlier attempt from racing with a
/// later attempt during rapid reconnections.
///
/// # Thread safety
///
/// Clones share the pending attempt and session state. All operations are thread-safe and can run
/// concurrently from multiple tasks.
#[derive(Clone, Debug)]
pub struct AuthTracker {
    tx: Arc<Mutex<Option<AuthResultSender>>>,
    state: Arc<AtomicU8>,
    state_notify: Arc<tokio::sync::Notify>,
}

impl AuthTracker {
    /// Creates a new authentication tracker.
    #[must_use]
    pub fn new() -> Self {
        Self {
            tx: Arc::new(Mutex::new(None)),
            state: Arc::new(AtomicU8::new(AuthState::Unauthenticated.as_u8())),
            state_notify: Arc::new(tokio::sync::Notify::new()),
        }
    }

    /// Returns the current authentication state.
    #[must_use]
    pub fn auth_state(&self) -> AuthState {
        AuthState::from_u8(self.state.load(Ordering::Acquire))
    }

    /// Returns whether the client is currently authenticated.
    #[must_use]
    pub fn is_authenticated(&self) -> bool {
        self.auth_state() == AuthState::Authenticated
    }

    /// Clears authenticated state without affecting pending auth attempts.
    ///
    /// Call this when a live connection drops and reconnect may authenticate
    /// again, so operations requiring authentication are properly guarded. A
    /// terminal [`AuthState::Failed`] state remains failed.
    pub fn invalidate(&self) {
        if self
            .state
            .try_update(Ordering::AcqRel, Ordering::Acquire, |state| {
                (AuthState::from_u8(state) == AuthState::Authenticated)
                    .then_some(AuthState::Unauthenticated.as_u8())
            })
            .is_ok()
        {
            self.state_notify.notify_waiters();
        }
    }

    /// Begins a new authentication attempt.
    ///
    /// Returns a receiver that will be notified when authentication completes.
    /// If a previous authentication attempt is still pending, it will be cancelled
    /// with an error message indicating it was superseded.
    ///
    /// Transitions to `Unauthenticated` since a new attempt invalidates any
    /// previous status.
    #[allow(
        clippy::must_use_candidate,
        reason = "callers use this for side effects"
    )]
    pub fn begin(&self) -> AuthResultReceiver {
        let (sender, receiver) = tokio::sync::oneshot::channel();
        self.state
            .store(AuthState::Unauthenticated.as_u8(), Ordering::Release);

        let mut guard = self.tx.lock();
        if let Some(old) = guard.take() {
            log::warn!("New authentication request superseding previous pending request");
            let _ = old.send(Err("Authentication attempt superseded".to_string()));
        } else {
            log::debug!("Starting new authentication request");
        }
        *guard = Some(sender);

        receiver
    }

    /// Marks the current authentication attempt as successful.
    ///
    /// Transitions to `Authenticated` and notifies any waiting receiver
    /// with `Ok(())`. This should be called when the server sends a successful
    /// authentication response.
    ///
    /// The state is always updated even if no receiver is waiting (e.g., after
    /// a timeout), since the server has confirmed authentication.
    pub fn succeed(&self) {
        self.state
            .store(AuthState::Authenticated.as_u8(), Ordering::Release);
        self.state_notify.notify_waiters();

        if let Some(sender) = self.tx.lock().take() {
            let _ = sender.send(Ok(()));
        }
    }

    /// Marks the current authentication attempt as failed.
    ///
    /// Transitions to `Failed` and notifies any waiting receiver
    /// with `Err(message)`. This should be called when the server sends an
    /// authentication error response, or on terminal client shutdown.
    ///
    /// The state is always updated even if no receiver is waiting, since the
    /// server has rejected authentication or future auth is impossible.
    pub fn fail(&self, error: impl Into<String>) {
        self.state
            .store(AuthState::Failed.as_u8(), Ordering::Release);
        self.state_notify.notify_waiters();
        let message = error.into();

        if let Some(sender) = self.tx.lock().take() {
            let _ = sender.send(Err(message));
        }
    }

    /// Waits for the authentication result with a timeout.
    ///
    /// Returns `Ok(())` if authentication succeeds, or an error if it fails,
    /// times out, or the channel is closed.
    ///
    /// # Type Parameters
    ///
    /// - `E`: Error type that implements `From<String>` for error message conversion
    ///
    /// # Errors
    ///
    /// Returns an error in the following cases:
    /// - Authentication fails (server rejects credentials)
    /// - Authentication times out (no response within timeout duration)
    /// - Authentication channel closes unexpectedly
    /// - Authentication attempt is superseded by a new attempt
    pub async fn wait_for_result<E>(
        &self,
        timeout: Duration,
        receiver: AuthResultReceiver,
    ) -> Result<(), E>
    where
        E: From<String>,
    {
        match dst::time::timeout(timeout, receiver).await {
            Ok(Ok(Ok(()))) => Ok(()),
            Ok(Ok(Err(msg))) => Err(E::from(msg)),
            Ok(Err(_)) => Err(E::from("Authentication channel closed".to_string())),
            Err(_) => {
                // Don't clear the sender: a concurrent begin() may have replaced it,
                // and guard.take() would cancel the newer sender. The next begin()
                // call cleans up any stale sender.
                Err(E::from("Authentication timed out".to_string()))
            }
        }
    }

    /// Waits for the tracker to enter the authenticated state.
    ///
    /// Returns `true` if authenticated within the timeout, `false` if the timeout
    /// expires or authentication explicitly fails. Uses event-driven notification
    /// from `succeed()` / `fail()` / `invalidate()` to avoid polling.
    ///
    /// Returns early with `false` when `fail()` is called (e.g., the exchange
    /// rejects credentials), so callers are not blocked for the full timeout
    /// on a definitive auth rejection.
    ///
    /// This is intended for callers on a separate task who need to gate operations
    /// on authentication state (e.g., order sends that must wait for re-authentication
    /// after a WebSocket reconnection).
    pub async fn wait_for_authenticated(&self, timeout: Duration) -> bool {
        if self.is_authenticated() {
            return true;
        }

        dst::time::timeout(timeout, async {
            loop {
                // Enable before the state check: an unpolled Notified is unregistered and misses notifies
                let mut notified = pin!(self.state_notify.notified());
                notified.as_mut().enable();

                match self.auth_state() {
                    AuthState::Authenticated => return true,
                    AuthState::Failed => return false,
                    AuthState::Unauthenticated => notified.await,
                }
            }
        })
        .await
        .unwrap_or(false)
    }
}

impl Default for AuthTracker {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
#[cfg(not(all(feature = "simulation", madsim)))]
mod tests {
    use std::{
        sync::atomic::{AtomicBool, Ordering},
        time::Duration,
    };

    use rstest::rstest;

    use super::*;

    #[derive(Debug, PartialEq)]
    struct TestError(String);

    impl From<String> for TestError {
        fn from(msg: String) -> Self {
            Self(msg)
        }
    }

    #[rstest]
    #[tokio::test]
    async fn test_successful_authentication() {
        let tracker = AuthTracker::new();
        let rx = tracker.begin();

        tracker.succeed();

        let result: Result<(), TestError> =
            tracker.wait_for_result(Duration::from_secs(1), rx).await;

        assert!(result.is_ok());
    }

    #[rstest]
    #[tokio::test]
    async fn test_failed_authentication() {
        let tracker = AuthTracker::new();
        let rx = tracker.begin();

        tracker.fail("Invalid credentials");

        let result: Result<(), TestError> =
            tracker.wait_for_result(Duration::from_secs(1), rx).await;

        assert_eq!(
            result.unwrap_err(),
            TestError("Invalid credentials".to_string())
        );
    }

    #[rstest]
    #[tokio::test]
    async fn test_authentication_timeout() {
        let tracker = AuthTracker::new();
        let rx = tracker.begin();

        // Don't call succeed or fail - let it timeout

        let result: Result<(), TestError> =
            tracker.wait_for_result(Duration::from_millis(50), rx).await;

        assert_eq!(
            result.unwrap_err(),
            TestError("Authentication timed out".to_string())
        );
    }

    #[rstest]
    #[tokio::test]
    async fn test_begin_supersedes_previous_sender() {
        let tracker = AuthTracker::new();

        let first = tracker.begin();
        let second = tracker.begin();

        // First receiver should get superseded error
        let result = first.await.expect("oneshot closed unexpectedly");
        assert_eq!(result, Err("Authentication attempt superseded".to_string()));

        // Second attempt should succeed
        tracker.succeed();
        let result: Result<(), TestError> = tracker
            .wait_for_result(Duration::from_secs(1), second)
            .await;

        assert!(result.is_ok());
    }

    #[rstest]
    #[tokio::test]
    async fn test_succeed_without_pending_auth() {
        let tracker = AuthTracker::new();

        tracker.succeed();

        assert_eq!(tracker.auth_state(), AuthState::Authenticated);
    }

    #[rstest]
    #[tokio::test]
    async fn test_fail_without_pending_auth() {
        let tracker = AuthTracker::new();

        tracker.fail("Some error");

        assert_eq!(tracker.auth_state(), AuthState::Failed);
    }

    #[rstest]
    #[tokio::test]
    async fn test_multiple_sequential_authentications() {
        let tracker = AuthTracker::new();

        // First auth succeeds
        let rx1 = tracker.begin();
        tracker.succeed();
        let result1: Result<(), TestError> =
            tracker.wait_for_result(Duration::from_secs(1), rx1).await;
        assert!(result1.is_ok());

        // Second auth fails
        let rx2 = tracker.begin();
        tracker.fail("Credentials expired");
        let result2: Result<(), TestError> =
            tracker.wait_for_result(Duration::from_secs(1), rx2).await;
        assert_eq!(
            result2.unwrap_err(),
            TestError("Credentials expired".to_string())
        );

        // Third auth succeeds
        let rx3 = tracker.begin();
        tracker.succeed();
        let result3: Result<(), TestError> =
            tracker.wait_for_result(Duration::from_secs(1), rx3).await;
        assert!(result3.is_ok());
    }

    #[rstest]
    #[tokio::test]
    async fn test_channel_closed_before_result() {
        let tracker = AuthTracker::new();
        let rx = tracker.begin();

        drop(tracker.tx.lock().take());

        let result: Result<(), TestError> =
            tracker.wait_for_result(Duration::from_secs(1), rx).await;

        assert_eq!(
            result,
            Err(TestError("Authentication channel closed".to_string()))
        );
        assert_eq!(tracker.auth_state(), AuthState::Unauthenticated);
    }

    #[rstest]
    #[tokio::test]
    async fn test_concurrent_auth_attempts() {
        let tracker = Arc::new(AuthTracker::new());
        let mut handles = vec![];

        // Spawn 10 concurrent auth attempts
        for i in 0..10 {
            let tracker_clone = Arc::clone(&tracker);
            let handle = tokio::spawn(async move {
                let rx = tracker_clone.begin();

                // Only the last one should succeed
                if i == 9 {
                    tokio::time::sleep(Duration::from_millis(10)).await;
                    tracker_clone.succeed();
                }

                let result: Result<(), TestError> = tracker_clone
                    .wait_for_result(Duration::from_secs(1), rx)
                    .await;

                (i, result)
            });
            handles.push(handle);
        }

        let mut successes = 0;
        let mut superseded = 0;

        for handle in handles {
            let (i, result) = handle.await.unwrap();
            match result {
                Ok(()) => {
                    // Only task 9 should succeed
                    assert_eq!(i, 9);
                    successes += 1;
                }
                Err(TestError(msg)) if msg.contains("superseded") => {
                    superseded += 1;
                }
                Err(e) => panic!("Unexpected error: {e:?}"),
            }
        }

        assert_eq!(successes, 1);
        assert_eq!(superseded, 9);
    }

    #[rstest]
    fn test_default_trait() {
        let _tracker = AuthTracker::default();
    }

    #[rstest]
    #[tokio::test]
    async fn test_clone_trait() {
        let tracker = AuthTracker::new();
        let cloned = tracker.clone();

        // Verify cloned instance shares state with original (Arc behavior)
        let rx = tracker.begin();
        cloned.succeed(); // Succeed via clone affects original
        let result: Result<(), TestError> =
            tracker.wait_for_result(Duration::from_secs(1), rx).await;
        assert!(result.is_ok());
    }

    #[rstest]
    fn test_debug_trait() {
        let tracker = AuthTracker::new();
        let debug_str = format!("{tracker:?}");
        assert!(debug_str.contains("AuthTracker"));
    }

    #[rstest]
    #[tokio::test]
    async fn test_timeout_clears_sender() {
        let tracker = AuthTracker::new();

        // Start auth that will timeout
        let rx1 = tracker.begin();
        let result1: Result<(), TestError> = tracker
            .wait_for_result(Duration::from_millis(50), rx1)
            .await;
        assert_eq!(
            result1.unwrap_err(),
            TestError("Authentication timed out".to_string())
        );

        // Verify sender was cleared - new auth should work
        let rx2 = tracker.begin();
        tracker.succeed();
        let result2: Result<(), TestError> =
            tracker.wait_for_result(Duration::from_secs(1), rx2).await;
        assert!(result2.is_ok());
    }

    #[rstest]
    #[tokio::test]
    async fn test_fail_clears_sender() {
        let tracker = AuthTracker::new();

        // Auth fails
        let rx1 = tracker.begin();
        tracker.fail("Bad credentials");
        let result1: Result<(), TestError> =
            tracker.wait_for_result(Duration::from_secs(1), rx1).await;
        assert!(result1.is_err());

        // Verify sender was cleared - new auth should work
        let rx2 = tracker.begin();
        tracker.succeed();
        let result2: Result<(), TestError> =
            tracker.wait_for_result(Duration::from_secs(1), rx2).await;
        assert!(result2.is_ok());
    }

    #[rstest]
    #[tokio::test]
    async fn test_succeed_clears_sender() {
        let tracker = AuthTracker::new();

        // Auth succeeds
        let rx1 = tracker.begin();
        tracker.succeed();
        let result1: Result<(), TestError> =
            tracker.wait_for_result(Duration::from_secs(1), rx1).await;
        assert!(result1.is_ok());

        // Verify sender was cleared - new auth should work
        let rx2 = tracker.begin();
        tracker.succeed();
        let result2: Result<(), TestError> =
            tracker.wait_for_result(Duration::from_secs(1), rx2).await;
        assert!(result2.is_ok());
    }

    #[rstest]
    #[tokio::test]
    async fn test_rapid_begin_succeed_cycles() {
        let tracker = AuthTracker::new();

        // Rapidly cycle through auth attempts
        for _ in 0..100 {
            let rx = tracker.begin();
            tracker.succeed();
            let result: Result<(), TestError> =
                tracker.wait_for_result(Duration::from_secs(1), rx).await;
            assert!(result.is_ok());
        }
    }

    #[rstest]
    #[tokio::test]
    async fn test_double_succeed_is_safe() {
        let tracker = AuthTracker::new();
        let rx = tracker.begin();

        // Call succeed twice
        tracker.succeed();
        tracker.succeed(); // Second call should be no-op

        let result: Result<(), TestError> =
            tracker.wait_for_result(Duration::from_secs(1), rx).await;
        assert!(result.is_ok());
    }

    #[rstest]
    #[tokio::test]
    async fn test_double_fail_is_safe() {
        let tracker = AuthTracker::new();
        let rx = tracker.begin();

        // Call fail twice
        tracker.fail("Error 1");
        tracker.fail("Error 2"); // Second call should be no-op

        let result: Result<(), TestError> =
            tracker.wait_for_result(Duration::from_secs(1), rx).await;
        assert_eq!(
            result.unwrap_err(),
            TestError("Error 1".to_string()) // Should be first error
        );
    }

    #[rstest]
    #[tokio::test]
    async fn test_succeed_after_fail_is_ignored() {
        let tracker = AuthTracker::new();
        let rx = tracker.begin();

        tracker.fail("Auth failed");
        tracker.succeed(); // This should be no-op

        let result: Result<(), TestError> =
            tracker.wait_for_result(Duration::from_secs(1), rx).await;
        assert!(result.is_err()); // Should still be error
    }

    #[rstest]
    #[tokio::test]
    async fn test_fail_after_succeed_is_ignored() {
        let tracker = AuthTracker::new();
        let rx = tracker.begin();

        tracker.succeed();
        tracker.fail("Auth failed"); // This should be no-op

        let result: Result<(), TestError> =
            tracker.wait_for_result(Duration::from_secs(1), rx).await;
        assert!(result.is_ok()); // Should still be success
    }

    /// Simulates a reconnect flow where authentication must complete before resubscription.
    ///
    /// This is an integration-style test that verifies:
    /// 1. On reconnect, authentication starts first
    /// 2. Subscription logic waits for auth to complete
    /// 3. Subscriptions only proceed after successful auth
    #[rstest]
    #[tokio::test]
    async fn test_reconnect_flow_waits_for_auth() {
        let tracker = Arc::new(AuthTracker::new());
        let subscribed = Arc::new(tokio::sync::Notify::new());
        let auth_completed = Arc::new(tokio::sync::Notify::new());

        // Simulate reconnect handler
        let tracker_reconnect = Arc::clone(&tracker);
        let subscribed_reconnect = Arc::clone(&subscribed);
        let auth_completed_reconnect = Arc::clone(&auth_completed);

        let reconnect_task = tokio::spawn(async move {
            // Step 1: Begin authentication
            let rx = tracker_reconnect.begin();

            // Step 2: Spawn resubscription task that waits for auth
            let tracker_resub = Arc::clone(&tracker_reconnect);
            let subscribed_resub = Arc::clone(&subscribed_reconnect);
            let auth_completed_resub = Arc::clone(&auth_completed_reconnect);

            let resub_task = tokio::spawn(async move {
                // Wait for auth to complete
                let result: Result<(), TestError> = tracker_resub
                    .wait_for_result(Duration::from_secs(5), rx)
                    .await;

                if result.is_ok() {
                    auth_completed_resub.notify_one();
                    // Simulate resubscription
                    tokio::time::sleep(Duration::from_millis(10)).await;
                    subscribed_resub.notify_one();
                }
            });

            resub_task.await.unwrap();
        });

        // Simulate server auth response after delay
        tokio::time::sleep(Duration::from_millis(100)).await;
        tracker.succeed();

        // Wait for reconnect flow to complete
        reconnect_task.await.unwrap();

        // Verify auth completed before subscription
        tokio::select! {
            () = auth_completed.notified() => {
                // Good - auth completed
            }
            () = tokio::time::sleep(Duration::from_secs(1)) => {
                panic!("Auth never completed");
            }
        }

        // Verify subscription completed
        tokio::select! {
            () = subscribed.notified() => {
                // Good - subscribed
            }
            () = tokio::time::sleep(Duration::from_secs(1)) => {
                panic!("Subscription never completed");
            }
        }
    }

    /// Verifies that failed authentication prevents resubscription in reconnect flow.
    #[rstest]
    #[tokio::test]
    async fn test_reconnect_flow_blocks_on_auth_failure() {
        let tracker = Arc::new(AuthTracker::new());
        let subscribed = Arc::new(AtomicBool::new(false));

        let tracker_reconnect = Arc::clone(&tracker);
        let subscribed_reconnect = Arc::clone(&subscribed);

        let reconnect_task = tokio::spawn(async move {
            let rx = tracker_reconnect.begin();

            // Spawn resubscription task that waits for auth
            let tracker_resub = Arc::clone(&tracker_reconnect);
            let subscribed_resub = Arc::clone(&subscribed_reconnect);

            let resub_task = tokio::spawn(async move {
                let result: Result<(), TestError> = tracker_resub
                    .wait_for_result(Duration::from_secs(5), rx)
                    .await;

                // Only subscribe if auth succeeds
                if result.is_ok() {
                    subscribed_resub.store(true, Ordering::Relaxed);
                }
            });

            resub_task.await.unwrap();
        });

        // Simulate server auth failure
        tokio::time::sleep(Duration::from_millis(50)).await;
        tracker.fail("Invalid credentials");

        // Wait for reconnect flow to complete
        reconnect_task.await.unwrap();

        // Verify subscription never happened
        tokio::time::sleep(Duration::from_millis(100)).await;
        assert!(!subscribed.load(Ordering::Relaxed));
    }

    /// Tests state machine transitions exhaustively.
    #[rstest]
    #[tokio::test]
    async fn test_state_machine_transitions() {
        let tracker = AuthTracker::new();

        // Transition 1: Initial -> Pending (begin)
        let rx1 = tracker.begin();

        // Transition 2: Pending -> Success (succeed)
        tracker.succeed();
        let result1: Result<(), TestError> =
            tracker.wait_for_result(Duration::from_secs(1), rx1).await;
        assert!(result1.is_ok());

        // Transition 3: Success -> Pending (begin again)
        let rx2 = tracker.begin();

        // Transition 4: Pending -> Failure (fail)
        tracker.fail("Error");
        let result2: Result<(), TestError> =
            tracker.wait_for_result(Duration::from_secs(1), rx2).await;
        assert!(result2.is_err());

        // Transition 5: Failure -> Pending (begin again)
        let rx3 = tracker.begin();

        // Transition 6: Pending -> Timeout
        let result3: Result<(), TestError> = tracker
            .wait_for_result(Duration::from_millis(50), rx3)
            .await;
        assert_eq!(
            result3.unwrap_err(),
            TestError("Authentication timed out".to_string())
        );

        // Transition 7: Timeout -> Pending (begin again)
        let rx4 = tracker.begin();

        // Transition 8: Pending -> Superseded (begin interrupts)
        let rx5 = tracker.begin();
        let result4: Result<(), TestError> =
            tracker.wait_for_result(Duration::from_secs(1), rx4).await;
        assert_eq!(
            result4.unwrap_err(),
            TestError("Authentication attempt superseded".to_string())
        );

        // Final success to clean up
        tracker.succeed();
        let result5: Result<(), TestError> =
            tracker.wait_for_result(Duration::from_secs(1), rx5).await;
        assert!(result5.is_ok());
    }

    /// Verifies no memory leaks from orphaned senders.
    #[rstest]
    #[tokio::test]
    async fn test_no_sender_leaks() {
        let tracker = AuthTracker::new();

        for _ in 0..100 {
            let rx = tracker.begin();
            let _result: Result<(), TestError> =
                tracker.wait_for_result(Duration::from_millis(1), rx).await;
        }

        let rx = tracker.begin();
        tracker.succeed();
        let result: Result<(), TestError> =
            tracker.wait_for_result(Duration::from_secs(1), rx).await;
        assert!(result.is_ok());
    }

    /// Tests concurrent success/fail calls don't cause panics.
    #[rstest]
    #[tokio::test]
    async fn test_concurrent_succeed_fail_calls() {
        let tracker = Arc::new(AuthTracker::new());
        let rx = tracker.begin();

        let mut handles = vec![];

        // Spawn many tasks trying to succeed
        for _ in 0..50 {
            let tracker_clone = Arc::clone(&tracker);
            handles.push(tokio::spawn(async move {
                tracker_clone.succeed();
            }));
        }

        // Spawn many tasks trying to fail
        for _ in 0..50 {
            let tracker_clone = Arc::clone(&tracker);
            handles.push(tokio::spawn(async move {
                tracker_clone.fail("Error");
            }));
        }

        // Wait for all tasks
        for handle in handles {
            handle.await.unwrap();
        }

        // Should get either success or failure, but not panic
        let result: Result<(), TestError> =
            tracker.wait_for_result(Duration::from_secs(1), rx).await;
        // Don't care which outcome, just that it doesn't panic
        let _ = result;
    }

    #[rstest]
    fn test_is_authenticated_initial_state() {
        let tracker = AuthTracker::new();
        assert!(!tracker.is_authenticated());
    }

    #[rstest]
    #[tokio::test]
    async fn test_is_authenticated_after_succeed() {
        let tracker = AuthTracker::new();
        assert!(!tracker.is_authenticated());

        let _rx = tracker.begin();
        assert!(!tracker.is_authenticated());

        tracker.succeed();
        assert!(tracker.is_authenticated());
    }

    #[rstest]
    #[tokio::test]
    async fn test_is_authenticated_after_fail() {
        let tracker = AuthTracker::new();
        let _rx = tracker.begin();
        tracker.fail("error");
        assert!(!tracker.is_authenticated());
    }

    #[rstest]
    #[tokio::test]
    async fn test_invalidate_clears_auth_state() {
        let tracker = AuthTracker::new();
        let _rx = tracker.begin();
        tracker.succeed();
        assert!(tracker.is_authenticated());

        tracker.invalidate();
        assert!(!tracker.is_authenticated());
    }

    #[rstest]
    #[case(true)]
    #[case(false)]
    #[tokio::test]
    async fn test_invalidate_preserves_terminal_failure(#[case] invalidate_first: bool) {
        let tracker = AuthTracker::new();
        let receiver = tracker.begin();

        if invalidate_first {
            tracker.invalidate();
            tracker.fail("terminal");
        } else {
            tracker.fail("terminal");
            tracker.invalidate();
        }

        let result: Result<(), TestError> = tracker
            .wait_for_result(Duration::from_secs(1), receiver)
            .await;

        assert_eq!(tracker.auth_state(), AuthState::Failed);
        assert_eq!(result.unwrap_err(), TestError("terminal".to_string()));
        assert!(
            !tracker
                .wait_for_authenticated(Duration::from_millis(10))
                .await
        );
    }

    #[rstest]
    #[tokio::test]
    async fn test_begin_clears_auth_state() {
        let tracker = AuthTracker::new();
        let _rx1 = tracker.begin();
        tracker.succeed();
        assert!(tracker.is_authenticated());

        let _rx2 = tracker.begin();
        assert!(!tracker.is_authenticated());
    }

    #[rstest]
    fn test_is_authenticated_shared_across_clones() {
        let tracker = AuthTracker::new();
        let cloned = tracker.clone();

        let _rx = tracker.begin();
        tracker.succeed();

        assert!(cloned.is_authenticated());
    }

    #[rstest]
    fn test_invalidate_shared_across_clones() {
        let tracker = AuthTracker::new();
        let cloned = tracker.clone();

        let _rx = tracker.begin();
        tracker.succeed();
        assert!(tracker.is_authenticated());

        cloned.invalidate();
        assert!(!tracker.is_authenticated());
    }

    #[rstest]
    fn test_succeed_without_begin_still_updates_auth_state() {
        let tracker = AuthTracker::new();
        assert!(!tracker.is_authenticated());

        // State updates even without begin() to handle late responses after timeout
        tracker.succeed();
        assert!(tracker.is_authenticated());
    }

    #[rstest]
    fn test_fail_without_begin_still_updates_auth_state() {
        let tracker = AuthTracker::new();
        tracker.succeed();
        assert!(tracker.is_authenticated());

        // State updates even without begin() to handle late responses
        tracker.fail("error");
        assert!(!tracker.is_authenticated());
    }

    #[rstest]
    #[tokio::test]
    async fn test_auth_state_false_after_timeout_until_late_response() {
        let tracker = AuthTracker::new();
        let rx = tracker.begin();
        assert!(!tracker.is_authenticated());

        let result: Result<(), TestError> =
            tracker.wait_for_result(Duration::from_millis(10), rx).await;

        assert!(result.is_err());
        assert!(!tracker.is_authenticated());

        // Late response after timeout still updates state
        tracker.succeed();
        assert!(tracker.is_authenticated());
    }

    #[rstest]
    #[tokio::test]
    async fn test_wait_for_authenticated_already_authenticated() {
        let tracker = AuthTracker::new();
        let _rx = tracker.begin();
        tracker.succeed();

        assert!(
            tracker
                .wait_for_authenticated(Duration::from_millis(50))
                .await
        );
    }

    #[rstest]
    #[tokio::test]
    async fn test_wait_for_authenticated_succeeds_after_delay() {
        let tracker = AuthTracker::new();
        let _rx = tracker.begin();

        let tracker_clone = tracker.clone();

        tokio::spawn(async move {
            tokio::time::sleep(Duration::from_millis(50)).await;
            tracker_clone.succeed();
        });

        assert!(tracker.wait_for_authenticated(Duration::from_secs(1)).await);
    }

    #[rstest]
    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn test_wait_for_authenticated_no_lost_wakeup_under_race() {
        // Regression: a succeed() landing between the waiter's state check and
        // its first poll of Notified must not be lost. Notified only registers
        // with the Notify once polled or enabled, so without enable() this
        // stalls for the full timeout and returns false on some iterations.
        for _ in 0..200 {
            let tracker = AuthTracker::new();
            let _rx = tracker.begin();

            let succeeder = tracker.clone();
            let handle = std::thread::spawn(move || succeeder.succeed());

            assert!(
                tracker
                    .wait_for_authenticated(Duration::from_millis(500))
                    .await,
                "wakeup lost despite successful authentication"
            );
            handle.join().unwrap();
        }
    }

    #[rstest]
    #[tokio::test]
    async fn test_wait_for_authenticated_returns_false_on_failure() {
        let tracker = AuthTracker::new();
        let _rx = tracker.begin();

        let tracker_clone = tracker.clone();

        tokio::spawn(async move {
            tokio::time::sleep(Duration::from_millis(50)).await;
            tracker_clone.fail("rejected");
        });

        let start = tokio::time::Instant::now();
        let result = tracker.wait_for_authenticated(Duration::from_secs(5)).await;
        let elapsed = start.elapsed();

        assert!(!result);
        assert!(elapsed < Duration::from_secs(1));
    }

    #[rstest]
    #[tokio::test]
    async fn test_wait_for_authenticated_times_out() {
        let tracker = AuthTracker::new();
        let _rx = tracker.begin();

        assert!(
            !tracker
                .wait_for_authenticated(Duration::from_millis(50))
                .await
        );
    }

    #[rstest]
    #[tokio::test]
    async fn test_wait_for_authenticated_begin_clears_failed() {
        let tracker = AuthTracker::new();
        let _rx = tracker.begin();
        tracker.fail("first attempt");

        assert!(
            !tracker
                .wait_for_authenticated(Duration::from_millis(10))
                .await
        );

        // begin() clears the failed flag, allowing a fresh wait
        let _rx = tracker.begin();

        let tracker_clone = tracker.clone();

        tokio::spawn(async move {
            tokio::time::sleep(Duration::from_millis(50)).await;
            tracker_clone.succeed();
        });

        assert!(tracker.wait_for_authenticated(Duration::from_secs(1)).await);
    }

    #[rstest]
    #[tokio::test]
    async fn test_wait_for_authenticated_invalidate_does_not_return_false() {
        let tracker = AuthTracker::new();
        let _rx = tracker.begin();

        let tracker_clone = tracker.clone();

        tokio::spawn(async move {
            // invalidate wakes the loop but should not cause early false return
            tokio::time::sleep(Duration::from_millis(20)).await;
            tracker_clone.invalidate();
            // then succeed shortly after
            tokio::time::sleep(Duration::from_millis(20)).await;
            tracker_clone.succeed();
        });

        assert!(tracker.wait_for_authenticated(Duration::from_secs(1)).await);
    }

    #[rstest]
    #[tokio::test]
    async fn test_wait_for_authenticated_concurrent_waiters() {
        let tracker = Arc::new(AuthTracker::new());
        let _rx = tracker.begin();

        let mut handles = vec![];

        for _ in 0..10 {
            let t = Arc::clone(&tracker);
            handles.push(tokio::spawn(async move {
                t.wait_for_authenticated(Duration::from_secs(1)).await
            }));
        }

        tokio::time::sleep(Duration::from_millis(50)).await;
        tracker.succeed();

        for handle in handles {
            assert!(handle.await.unwrap());
        }
    }

    #[rstest]
    #[tokio::test]
    async fn test_wait_for_authenticated_not_authenticated_initially() {
        let tracker = AuthTracker::new();

        // Not authenticated, no begin() called, no failed flag set
        // Should time out
        assert!(
            !tracker
                .wait_for_authenticated(Duration::from_millis(50))
                .await
        );
    }
}

#[cfg(test)]
#[cfg(not(all(feature = "simulation", madsim)))]
mod proptest_tests {
    use std::{sync::Arc, time::Duration};

    use proptest::prelude::*;
    use rstest::rstest;

    use super::*;

    const AUTH_FAILED: &str = "model auth failed";
    const AUTH_SUPERSEDED: &str = "Authentication attempt superseded";

    #[derive(Debug, Clone)]
    enum AuthTraceOp {
        Begin,
        Succeed,
        Fail,
        Invalidate,
        WaitForAuthenticated,
    }

    #[derive(Debug, Clone, PartialEq, Eq)]
    enum ExpectedAuthResult {
        Success,
        Failed(&'static str),
    }

    #[derive(Debug)]
    struct AuthTraceModel {
        state: AuthState,
        pending_receiver: Option<usize>,
        expected_results: Vec<Option<ExpectedAuthResult>>,
    }

    impl AuthTraceModel {
        fn new() -> Self {
            Self {
                state: AuthState::Unauthenticated,
                pending_receiver: None,
                expected_results: Vec::new(),
            }
        }

        fn begin(&mut self) {
            if let Some(receiver_index) = self.pending_receiver.take() {
                self.expected_results[receiver_index] =
                    Some(ExpectedAuthResult::Failed(AUTH_SUPERSEDED));
            }

            self.pending_receiver = Some(self.expected_results.len());
            self.expected_results.push(None);
            self.state = AuthState::Unauthenticated;
        }

        fn succeed(&mut self) {
            self.state = AuthState::Authenticated;

            if let Some(receiver_index) = self.pending_receiver.take() {
                self.expected_results[receiver_index] = Some(ExpectedAuthResult::Success);
            }
        }

        fn fail(&mut self) {
            self.state = AuthState::Failed;

            if let Some(receiver_index) = self.pending_receiver.take() {
                self.expected_results[receiver_index] =
                    Some(ExpectedAuthResult::Failed(AUTH_FAILED));
            }
        }

        fn invalidate(&mut self) {
            if self.state == AuthState::Authenticated {
                self.state = AuthState::Unauthenticated;
            }
        }
    }

    fn auth_trace_op_strategy() -> impl Strategy<Value = AuthTraceOp> {
        prop_oneof![
            Just(AuthTraceOp::Begin),
            Just(AuthTraceOp::Succeed),
            Just(AuthTraceOp::Fail),
            Just(AuthTraceOp::Invalidate),
            Just(AuthTraceOp::WaitForAuthenticated),
        ]
    }

    fn assert_auth_receivers_match_model(
        receivers: &mut [Option<AuthResultReceiver>],
        model: &AuthTraceModel,
        step: usize,
    ) -> Result<(), TestCaseError> {
        for (receiver_index, receiver_slot) in receivers.iter_mut().enumerate() {
            let Some(receiver) = receiver_slot.as_mut() else {
                continue;
            };

            let mut clear_receiver = false;

            match &model.expected_results[receiver_index] {
                Some(ExpectedAuthResult::Success) => {
                    match receiver.try_recv() {
                        Ok(Ok(())) => {}
                        actual => prop_assert!(
                            false,
                            "receiver {} should succeed at step {}, was {:?}",
                            receiver_index,
                            step,
                            actual
                        ),
                    }
                    clear_receiver = true;
                }
                Some(ExpectedAuthResult::Failed(expected)) => {
                    match receiver.try_recv() {
                        Ok(Err(actual)) => prop_assert_eq!(
                            actual,
                            *expected,
                            "receiver {} should fail at step {}",
                            receiver_index,
                            step
                        ),
                        actual => prop_assert!(
                            false,
                            "receiver {} should fail at step {}, was {:?}",
                            receiver_index,
                            step,
                            actual
                        ),
                    }
                    clear_receiver = true;
                }
                None => {
                    prop_assert_eq!(
                        receiver.try_recv(),
                        Err(tokio::sync::oneshot::error::TryRecvError::Empty),
                        "receiver {} should stay pending at step {}",
                        receiver_index,
                        step
                    );
                }
            }

            if clear_receiver {
                *receiver_slot = None;
            }
        }

        Ok(())
    }

    proptest! {
        #![proptest_config(ProptestConfig::with_cases(256))]

        /// Property: auth attempt traces match the cycle model for success,
        /// failure, superseded stale receivers, invalidation, and auth waits.
        #[rstest]
        fn test_auth_tracker_trace_matches_cycle_model(
            ops in proptest::collection::vec(auth_trace_op_strategy(), 1..80)
        ) {
            let runtime = tokio::runtime::Builder::new_current_thread()
                .enable_time()
                .build()
                .unwrap();
            let tracker = AuthTracker::new();
            let mut model = AuthTraceModel::new();
            let mut receivers: Vec<Option<AuthResultReceiver>> = Vec::new();

            for (step, op) in ops.iter().enumerate() {
                match op {
                    AuthTraceOp::Begin => {
                        receivers.push(Some(tracker.begin()));
                        model.begin();
                    }
                    AuthTraceOp::Succeed => {
                        tracker.succeed();
                        model.succeed();
                    }
                    AuthTraceOp::Fail => {
                        tracker.fail(AUTH_FAILED);
                        model.fail();
                    }
                    AuthTraceOp::Invalidate => {
                        tracker.invalidate();
                        model.invalidate();
                    }
                    AuthTraceOp::WaitForAuthenticated => {
                        let actual = runtime
                            .block_on(tracker.wait_for_authenticated(Duration::from_millis(0)));
                        let expected = model.state == AuthState::Authenticated;
                        prop_assert_eq!(
                            actual,
                            expected,
                            "wait_for_authenticated mismatch at step {}, op {:?}",
                            step,
                            op
                        );
                    }
                }

                prop_assert_eq!(
                    tracker.auth_state(),
                    model.state,
                    "auth state mismatch at step {}, op {:?}",
                    step,
                    op
                );
                assert_auth_receivers_match_model(&mut receivers, &model, step)?;
            }
        }

        /// Verifies that any sequence of begin/succeed/fail/invalidate calls
        /// leaves the tracker in a consistent state where `is_authenticated`
        /// agrees with the last state-setting call.
        #[rstest]
        fn test_state_consistency_after_random_operations(
            ops in proptest::collection::vec(0u8..4, 1..50)
        ) {
            let tracker = AuthTracker::new();
            let mut expected_auth = false;

            for op in &ops {
                match op {
                    0 => {
                        let _rx = tracker.begin();
                        expected_auth = false;
                    }
                    1 => {
                        tracker.succeed();
                        expected_auth = true;
                    }
                    2 => {
                        tracker.fail("test");
                        expected_auth = false;
                    }
                    3 => {
                        tracker.invalidate();
                        expected_auth = false;
                    }
                    _ => unreachable!(),
                }
            }

            prop_assert_eq!(tracker.is_authenticated(), expected_auth);
        }

        /// Verifies that begin() always clears the failed flag regardless of
        /// prior state, so a new auth attempt starts clean.
        #[rstest]
        fn test_begin_always_clears_failed(
            prior_ops in proptest::collection::vec(0u8..4, 0..20)
        ) {
            let tracker = AuthTracker::new();

            for op in &prior_ops {
                match op {
                    0 => { let _rx = tracker.begin(); }
                    1 => tracker.succeed(),
                    2 => tracker.fail("test"),
                    3 => tracker.invalidate(),
                    _ => unreachable!(),
                }
            }

            let _rx = tracker.begin();
            // After begin(), state is Unauthenticated
            prop_assert_eq!(tracker.auth_state(), AuthState::Unauthenticated);
        }

        /// Verifies that succeed() always transitions to Authenticated,
        /// regardless of prior state.
        #[rstest]
        fn test_succeed_always_sets_authenticated(
            prior_ops in proptest::collection::vec(0u8..4, 0..20)
        ) {
            let tracker = AuthTracker::new();

            for op in &prior_ops {
                match op {
                    0 => { let _rx = tracker.begin(); }
                    1 => tracker.succeed(),
                    2 => tracker.fail("test"),
                    3 => tracker.invalidate(),
                    _ => unreachable!(),
                }
            }

            tracker.succeed();
            prop_assert_eq!(tracker.auth_state(), AuthState::Authenticated);
        }
    }

    /// Verifies that `wait_for_authenticated` returns within a bounded time
    /// when `succeed()` or `fail()` is called, regardless of the timeout value.
    #[rstest]
    #[tokio::test]
    async fn test_wait_responds_within_bounded_time() {
        for auth_result in [true, false] {
            let tracker = Arc::new(AuthTracker::new());
            let _rx = tracker.begin();

            let tracker_clone = Arc::clone(&tracker);

            tokio::spawn(async move {
                tokio::time::sleep(Duration::from_millis(30)).await;

                if auth_result {
                    tracker_clone.succeed();
                } else {
                    tracker_clone.fail("rejected");
                }
            });

            let start = tokio::time::Instant::now();
            let result = tracker
                .wait_for_authenticated(Duration::from_secs(10))
                .await;
            let elapsed = start.elapsed();

            assert_eq!(result, auth_result);
            assert!(
                elapsed < Duration::from_millis(500),
                "wait_for_authenticated took {elapsed:?} for auth_result={auth_result}"
            );
        }
    }
}

#[cfg(all(test, feature = "simulation", madsim))]
mod simulation_tests {
    use std::time::Duration;

    use super::*;

    #[madsim::test]
    async fn test_wait_for_result_succeeds_without_tokio_reactor() {
        let tracker = AuthTracker::new();
        let rx = tracker.begin();
        tracker.succeed();
        let result: Result<(), String> = tracker.wait_for_result(Duration::from_secs(1), rx).await;
        assert_eq!(result, Ok(()));
    }

    #[madsim::test]
    async fn test_wait_for_result_times_out_on_virtual_clock() {
        let tracker = AuthTracker::new();
        let rx = tracker.begin();
        let result: Result<(), String> =
            tracker.wait_for_result(Duration::from_millis(10), rx).await;
        assert_eq!(result, Err("Authentication timed out".to_string()));
    }

    #[madsim::test]
    async fn test_wait_for_authenticated_succeeds_without_tokio_reactor() {
        let tracker = AuthTracker::new();
        let pending = tracker.clone();

        madsim::task::spawn(async move {
            madsim::time::sleep(Duration::from_millis(1)).await;
            pending.succeed();
        });
        assert!(tracker.wait_for_authenticated(Duration::from_secs(1)).await);
    }
}