highlevel_api/auth/

refresh.rs

use std::collections::HashMap;
use std::sync::Arc;

use tokio::sync::Mutex;

use super::token::TokenData;
use crate::error::Result;

/// Check if a token is expired and refresh it if needed.
///
/// Calls `refresh` only when `token.is_expired()` returns true. The caller is
/// responsible for persisting the new token.
pub async fn auto_refresh_token<F, Fut>(token: &mut TokenData, refresh: F) -> Result<()>
where
    F: FnOnce() -> Fut,
    Fut: std::future::Future<Output = Result<TokenData>>,
{
    if !token.is_expired() {
        return Ok(());
    }
    let new_token = refresh().await?;
    *token = new_token;
    Ok(())
}

/// Single-flight token refresher keyed by an arbitrary string (typically
/// location ID or company ID). Ensures that concurrent callers for the same
/// key share one refresh request, preventing refresh-token rotation conflicts.
#[derive(Clone)]
pub struct TokenRefresher {
    locks: Arc<Mutex<HashMap<String, Arc<Mutex<()>>>>>,
}

impl TokenRefresher {
    pub fn new() -> Self {
        Self {
            locks: Arc::new(Mutex::new(HashMap::new())),
        }
    }

    /// Acquire a per-key lock and refresh the token if expired.
    ///
    /// `key` is the location/company identifier. `token` is the current
    /// `TokenData` (will be updated in place if refreshed). `refresh` is an
    /// async callable that performs the refresh and returns a new `TokenData`.
    pub async fn ensure_valid<F, Fut>(
        &self,
        key: &str,
        token: &mut TokenData,
        refresh: F,
    ) -> Result<()>
    where
        F: FnOnce() -> Fut + Send,
        Fut: std::future::Future<Output = Result<TokenData>> + Send,
    {
        // Fast path: not expired, skip locking entirely.
        if !token.is_expired() {
            return Ok(());
        }

        // Get or create the per-key mutex.
        let lock = {
            let mut map = self.locks.lock().await;
            map.entry(key.to_string())
                .or_insert_with(|| Arc::new(Mutex::new(())))
                .clone()
        };

        let _guard = lock.lock().await;

        // Double-check after acquiring the lock: another thread may have
        // already refreshed.
        if !token.is_expired() {
            return Ok(());
        }

        let new_token = refresh().await?;
        *token = new_token;
        Ok(())
    }
}

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

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn unused_token_does_not_refresh() {
        let mut token = TokenData {
            access_token: "tok".into(),
            refresh_token: None,
            token_type: "Bearer".into(),
            expires_in: 3600,
            scope: None,
            user_type: None,
            company_id: None,
            location_id: None,
            issued_at: None,
        };

        let call_count = Arc::new(std::sync::atomic::AtomicU32::new(0));
        let replacement = token.clone();
        let cc = call_count.clone();
        auto_refresh_token(&mut token, move || {
            cc.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
            let t = replacement.clone();
            async { Ok(t) }
        })
        .await
        .unwrap();

        assert_eq!(call_count.load(std::sync::atomic::Ordering::SeqCst), 0);
    }

    #[tokio::test]
    async fn expired_token_triggers_refresh() {
        let mut token = TokenData {
            access_token: "tok".into(),
            refresh_token: None,
            token_type: "Bearer".into(),
            expires_in: 3600,
            scope: None,
            user_type: None,
            company_id: None,
            location_id: None,
            issued_at: Some(1),
        };

        let call_count = Arc::new(std::sync::atomic::AtomicU32::new(0));
        let replacement = TokenData {
            access_token: "new_tok".into(),
            refresh_token: None,
            token_type: "Bearer".into(),
            expires_in: 3600,
            scope: None,
            user_type: None,
            company_id: None,
            location_id: None,
            issued_at: None,
        };
        let cc = call_count.clone();
        auto_refresh_token(&mut token, move || {
            cc.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
            let t = replacement.clone();
            async { Ok(t) }
        })
        .await
        .unwrap();

        assert_eq!(call_count.load(std::sync::atomic::Ordering::SeqCst), 1);
        assert_eq!(token.access_token, "new_tok");
    }

    #[tokio::test]
    async fn single_flight_refreshes_once() {
        let refresher = TokenRefresher::new();

        // Shared token state across all concurrent callers.
        let shared = Arc::new(tokio::sync::Mutex::new(TokenData {
            access_token: "tok".into(),
            refresh_token: None,
            token_type: "Bearer".into(),
            expires_in: 3600,
            scope: None,
            user_type: None,
            company_id: None,
            location_id: None,
            issued_at: Some(1),
        }));

        let call_count = Arc::new(std::sync::atomic::AtomicU32::new(0));

        let mut tasks = vec![];
        for _ in 0..5 {
            let r = refresher.clone();
            let cc = call_count.clone();
            let shared = shared.clone();
            tasks.push(tokio::spawn(async move {
                let mut token = shared.lock().await;
                r.ensure_valid("loc_1", &mut *token, || {
                    cc.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
                    async move {
                        Ok(TokenData {
                            access_token: "refreshed".into(),
                            refresh_token: None,
                            token_type: "Bearer".into(),
                            expires_in: 3600,
                            scope: None,
                            user_type: None,
                            company_id: None,
                            location_id: None,
                            issued_at: None,
                        })
                    }
                })
                .await
                .unwrap();
                token.access_token.clone()
            }));
        }

        for task in tasks {
            assert_eq!(task.await.unwrap(), "refreshed");
        }
        assert_eq!(call_count.load(std::sync::atomic::Ordering::SeqCst), 1);
    }
}