spiral-auth 0.1.7

use std::fmt::{Debug, Formatter};
use std::sync::Arc;
use std::time::Duration;

use tokio::sync::{Mutex, RwLock};
use tokio::task::JoinHandle;

use crate::{CredentialResult, Token, TokenSource};

const DEFAULT_REFRESH_BEFORE: Duration = Duration::from_secs(60);
const DEFAULT_RETRY_DELAY: Duration = Duration::from_secs(5);
const MIN_USABLE_LIFETIME: Duration = Duration::from_secs(30);

struct ServiceTokenState {
    source: TokenSource,
    cached: RwLock<Option<Token>>,
    refresh: Mutex<()>,
    refresh_before: Duration,
    retry_delay: Duration,
}

/// Keeps a short-lived service token warm and returns only usable tokens.
///
/// Construction starts a background refresh. Concurrent refreshes are
/// single-flighted, and a failed early refresh leaves the last usable token in
/// place. [`token`](Self::token) returns only a token that remains usable
/// for at least 30 seconds.
pub struct RefreshingAuthTokenProvider {
    state: Arc<ServiceTokenState>,
    background: JoinHandle<()>,
}

impl RefreshingAuthTokenProvider {
    /// Wrap a token source with background warming and early refresh.
    pub fn new(source: TokenSource) -> Self {
        Self::with_timing(source, DEFAULT_REFRESH_BEFORE, DEFAULT_RETRY_DELAY)
    }

    fn with_timing(source: TokenSource, refresh_before: Duration, retry_delay: Duration) -> Self {
        let state = Arc::new(ServiceTokenState {
            source,
            cached: RwLock::new(None),
            refresh: Mutex::new(()),
            refresh_before,
            retry_delay,
        });
        let background_state = Arc::clone(&state);
        let background = tokio::spawn(async move {
            background_state.refresh_loop().await;
        });
        Self { state, background }
    }
}

impl Debug for RefreshingAuthTokenProvider {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("RefreshingAuthTokenProvider")
            .field("source", &self.state.source)
            .finish_non_exhaustive()
    }
}

impl Drop for RefreshingAuthTokenProvider {
    fn drop(&mut self) {
        self.background.abort();
    }
}

impl RefreshingAuthTokenProvider {
    /// A usable token, or `None` when this provider has no credentials to offer.
    pub async fn token(&self) -> CredentialResult<Option<Token>> {
        self.state
            .token_with_lifetime(MIN_USABLE_LIFETIME)
            .await
            .map(Some)
    }

    /// A token this provider already holds, without network requests or prompts.
    pub async fn cached_token(&self) -> CredentialResult<Option<Token>> {
        self.state.cached_with_lifetime(MIN_USABLE_LIFETIME).await
    }
}

impl_token_source!(RefreshingAuthTokenProvider, cached);

impl ServiceTokenState {
    async fn refresh_loop(self: Arc<Self>) {
        loop {
            if let Err(error) = self.token_with_lifetime(self.refresh_before).await {
                tracing::warn!(?error, "failed to refresh service token");
            }

            tokio::time::sleep(self.next_refresh_delay().await).await;
        }
    }

    async fn next_refresh_delay(&self) -> Duration {
        let cached = self.cached.read().await;
        let Some(token) = cached.as_ref() else {
            return self.retry_delay;
        };
        match token.remaining_lifetime() {
            Ok(remaining) if remaining > self.refresh_before => remaining - self.refresh_before,
            _ => self.retry_delay,
        }
    }

    async fn token_with_lifetime(&self, minimum_lifetime: Duration) -> CredentialResult<Token> {
        if let Some(token) = self.cached_with_lifetime(minimum_lifetime).await? {
            return Ok(token);
        }

        let _refresh = self.refresh.lock().await;
        if let Some(token) = self.cached_with_lifetime(minimum_lifetime).await? {
            return Ok(token);
        }

        let Some(token) = self.source.token().await? else {
            bail!("service token source returned no token");
        };
        let remaining = token
            .remaining_lifetime()
            .map_err(|error| invalid!("service token is not a valid JWT: {error}"))?;
        if remaining < MIN_USABLE_LIFETIME {
            bail!("service token expires too soon");
        }
        *self.cached.write().await = Some(token.clone());
        Ok(token)
    }

    async fn cached_with_lifetime(
        &self,
        minimum_lifetime: Duration,
    ) -> CredentialResult<Option<Token>> {
        let cached = self.cached.read().await;
        let Some(token) = cached.as_ref() else {
            return Ok(None);
        };
        let remaining = token
            .remaining_lifetime()
            .map_err(|error| invalid!("cached service token is not a valid JWT: {error}"))?;
        Ok((remaining >= minimum_lifetime).then(|| token.clone()))
    }
}

#[cfg(test)]
mod tests {
    use std::collections::VecDeque;
    use std::sync::atomic::{AtomicUsize, Ordering};

    use serde_json::json;

    use super::*;
    use crate::token::unsigned_jwt;

    fn token(expires_in: Duration) -> Token {
        let now = std::time::SystemTime::now()
            .duration_since(std::time::SystemTime::UNIX_EPOCH)
            .unwrap()
            .as_secs();
        unsigned_jwt(json!({ "exp": now + expires_in.as_secs() }))
    }

    struct SequenceProvider {
        tokens: Mutex<VecDeque<Token>>,
        calls: AtomicUsize,
    }

    impl Debug for SequenceProvider {
        fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
            f.debug_struct("SequenceProvider").finish_non_exhaustive()
        }
    }

    impl SequenceProvider {
        /// A usable token, or `None` when this provider has no credentials to offer.
        pub async fn token(&self) -> CredentialResult<Option<Token>> {
            self.calls.fetch_add(1, Ordering::SeqCst);
            let mut tokens = self.tokens.lock().await;
            Ok(tokens.pop_front().or_else(|| tokens.back().cloned()))
        }
    }

    impl_token_source!(Arc<SequenceProvider>);

    fn source(tokens: impl IntoIterator<Item = Token>) -> Arc<SequenceProvider> {
        Arc::new(SequenceProvider {
            tokens: Mutex::new(tokens.into_iter().collect()),
            calls: AtomicUsize::new(0),
        })
    }

    #[tokio::test]
    async fn concurrent_getters_share_one_refresh() {
        let source = source([token(Duration::from_secs(3_600))]);
        let provider = TokenSource::from(RefreshingAuthTokenProvider::new(source.clone().into()));
        let getters = (0..20).map(|_| {
            let provider = provider.clone();
            tokio::spawn(async move { provider.token().await.unwrap().unwrap() })
        });
        for getter in getters {
            getter.await.unwrap();
        }
        assert_eq!(source.calls.load(Ordering::SeqCst), 1);
    }

    #[tokio::test]
    async fn refreshes_early_in_the_background() {
        let first = token(Duration::from_secs(90));
        let second = token(Duration::from_secs(3_600));
        let source = source([first.clone(), second.clone()]);
        let provider = RefreshingAuthTokenProvider::with_timing(
            source.clone().into(),
            Duration::from_secs(100),
            Duration::from_millis(5),
        );

        tokio::time::timeout(Duration::from_secs(1), async {
            while source.calls.load(Ordering::SeqCst) < 2 {
                tokio::time::sleep(Duration::from_millis(5)).await;
            }
        })
        .await
        .unwrap();

        assert_eq!(provider.token().await.unwrap().unwrap(), second);
    }

    #[tokio::test]
    async fn rejects_a_token_that_is_already_unusable() {
        let source = source([token(Duration::from_secs(5))]);
        let provider = RefreshingAuthTokenProvider::new(source.into());
        let error = provider.token().await.unwrap_err();
        assert!(error.to_string().contains("expires too soon"));
    }
}