olai_http/

token.rs

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// 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.

use std::future::Future;
use std::time::{Duration, Instant};

use base64::Engine;
use chrono::{DateTime, Utc};
use serde::Deserialize;
use tokio::sync::Mutex;

/// A temporary authentication token with an associated expiry
#[derive(Debug, Clone)]
pub struct TemporaryToken<T> {
    /// The temporary credential
    pub token: T,
    /// The instant at which this credential is no longer valid
    /// None means the credential does not expire
    pub expiry: Option<Instant>,
}

/// The instant `expires` falls at, or now if it has already passed.
pub(crate) fn expiry_from(expires: DateTime<Utc>) -> Instant {
    Instant::now() + (expires - Utc::now()).to_std().unwrap_or_default()
}

/// The `exp` claim of `token` in seconds since the Unix epoch, if `token` is a JWT
/// with one.
///
/// The signature is not verified: the expiry only decides when to refresh.
pub(crate) fn jwt_exp_secs(token: &str) -> Option<u64> {
    #[derive(Deserialize)]
    struct Claims {
        exp: u64,
    }
    let payload = token.split('.').nth(1)?;
    let claims = base64::engine::general_purpose::URL_SAFE_NO_PAD
        .decode(payload.trim_end_matches('='))
        .ok()?;
    serde_json::from_slice::<Claims>(&claims)
        .ok()
        .map(|claims| claims.exp)
}

/// Thread-safe cache for a [`TemporaryToken`] that proactively refreshes before
/// the token expires.
///
/// ## Min-TTL strategy
///
/// A cached token is considered *still valid* only when its remaining lifetime
/// exceeds `min_ttl` (default **5 minutes**).  Once the remaining lifetime
/// drops below that threshold the next call to [`get_or_insert_with`] will
/// trigger a synchronous re-fetch — while the lock is held — before returning
/// the new token.  This ensures callers always receive a token with a
/// comfortable margin before it expires, avoiding races where a token is used
/// just as it becomes invalid.
///
/// An additional `fetch_backoff` guard (default **100 ms**) prevents a
/// thundering-herd of re-fetch attempts in the window where the token is below
/// `min_ttl` but has not yet expired: if the previous fetch happened within
/// `fetch_backoff` and the token has not actually expired yet, the stale-but-
/// not-yet-expired token is returned immediately.
///
/// [`get_or_insert_with`]: TokenCache::get_or_insert_with
#[derive(Debug)]
pub struct TokenCache<T> {
    cache: Mutex<Option<(TemporaryToken<T>, Instant)>>,
    min_ttl: Duration,
    fetch_backoff: Duration,
}

impl<T> Default for TokenCache<T> {
    fn default() -> Self {
        Self {
            cache: Default::default(),
            min_ttl: Duration::from_secs(300),
            // How long to wait before re-attempting a token fetch after receiving one that
            // is still within the min-ttl
            fetch_backoff: Duration::from_millis(100),
        }
    }
}

impl<T: Clone + Send> TokenCache<T> {
    /// Override the minimum remaining TTL for a cached token to be used
    pub(crate) fn with_min_ttl(self, min_ttl: Duration) -> Self {
        Self { min_ttl, ..self }
    }

    pub async fn get_or_insert_with<F, Fut, E>(&self, f: F) -> Result<T, E>
    where
        F: FnOnce() -> Fut + Send,
        Fut: Future<Output = Result<TemporaryToken<T>, E>> + Send,
    {
        let now = Instant::now();
        let mut locked = self.cache.lock().await;

        if let Some((cached, fetched_at)) = locked.as_ref() {
            match cached.expiry {
                Some(ttl) => {
                    let remaining = ttl.checked_duration_since(now);
                    if remaining.unwrap_or_default() > self.min_ttl ||
                        // if we've recently attempted to fetch this token and it's not actually
                        // expired, we'll wait to re-fetch it and return the cached one
                        (fetched_at.elapsed() < self.fetch_backoff && remaining.is_some())
                    {
                        return Ok(cached.token.clone());
                    }
                }
                None => return Ok(cached.token.clone()),
            }
        }

        let cached = f().await?;
        let token = cached.token.clone();
        *locked = Some((cached, Instant::now()));

        Ok(token)
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use std::sync::Arc;
    use std::sync::atomic::{AtomicU32, Ordering};
    use std::time::{Duration, Instant};

    // Helper function to create a token with a specific expiry duration from now
    fn create_token(expiry_duration: Option<Duration>) -> TemporaryToken<String> {
        TemporaryToken {
            token: "test_token".to_string(),
            expiry: expiry_duration.map(|d| Instant::now() + d),
        }
    }

    #[tokio::test]
    async fn test_expired_token_is_refreshed() {
        let cache = TokenCache::default();
        static COUNTER: AtomicU32 = AtomicU32::new(0);

        async fn get_token() -> Result<TemporaryToken<String>, String> {
            COUNTER.fetch_add(1, Ordering::SeqCst);
            Ok::<_, String>(create_token(Some(Duration::from_secs(0))))
        }

        // Should fetch initial token
        let _ = cache.get_or_insert_with(get_token).await.unwrap();
        assert_eq!(COUNTER.load(Ordering::SeqCst), 1);

        tokio::time::sleep(Duration::from_millis(2)).await;

        // Token is expired, so should fetch again
        let _ = cache.get_or_insert_with(get_token).await.unwrap();
        assert_eq!(COUNTER.load(Ordering::SeqCst), 2);
    }

    #[tokio::test]
    async fn test_min_ttl_causes_refresh() {
        let cache = TokenCache {
            cache: Default::default(),
            min_ttl: Duration::from_secs(1),
            fetch_backoff: Duration::from_millis(1),
        };

        static COUNTER: AtomicU32 = AtomicU32::new(0);

        async fn get_token() -> Result<TemporaryToken<String>, String> {
            COUNTER.fetch_add(1, Ordering::SeqCst);
            Ok::<_, String>(create_token(Some(Duration::from_millis(100))))
        }

        // Initial fetch
        let _ = cache.get_or_insert_with(get_token).await.unwrap();
        assert_eq!(COUNTER.load(Ordering::SeqCst), 1);

        // Should not fetch again since not expired and within fetch_backoff
        let _ = cache.get_or_insert_with(get_token).await.unwrap();
        assert_eq!(COUNTER.load(Ordering::SeqCst), 1);

        tokio::time::sleep(Duration::from_millis(2)).await;

        // Should fetch, since we've passed fetch_backoff
        let _ = cache.get_or_insert_with(get_token).await.unwrap();
        assert_eq!(COUNTER.load(Ordering::SeqCst), 2);
    }

    #[tokio::test]
    async fn test_valid_token_within_ttl_is_reused() {
        // Default min_ttl is 5 minutes; a token valid for an hour is comfortably
        // within its window and must be served from cache, never re-fetched.
        let cache = TokenCache::default();
        static COUNTER: AtomicU32 = AtomicU32::new(0);

        async fn get_token() -> Result<TemporaryToken<String>, String> {
            COUNTER.fetch_add(1, Ordering::SeqCst);
            Ok::<_, String>(create_token(Some(Duration::from_secs(3600))))
        }

        // First call performs the fetch.
        let _ = cache.get_or_insert_with(get_token).await.unwrap();
        assert_eq!(COUNTER.load(Ordering::SeqCst), 1);

        // Several subsequent calls all hit the cache.
        for _ in 0..5 {
            let _ = cache.get_or_insert_with(get_token).await.unwrap();
        }
        assert_eq!(COUNTER.load(Ordering::SeqCst), 1);
    }

    #[tokio::test]
    async fn test_non_expiring_token_is_reused() {
        // A token with `expiry == None` never expires and must always be served
        // from cache after the initial fetch.
        let cache = TokenCache::default();
        static COUNTER: AtomicU32 = AtomicU32::new(0);

        async fn get_token() -> Result<TemporaryToken<String>, String> {
            COUNTER.fetch_add(1, Ordering::SeqCst);
            Ok::<_, String>(create_token(None))
        }

        let _ = cache.get_or_insert_with(get_token).await.unwrap();
        let _ = cache.get_or_insert_with(get_token).await.unwrap();
        assert_eq!(COUNTER.load(Ordering::SeqCst), 1);
    }

    #[tokio::test]
    async fn test_concurrent_callers_single_flight() {
        // `get_or_insert_with` holds the cache mutex across the fetch future, so
        // concurrent callers serialize: the first performs the (slow) fetch and
        // stores a token comfortably within `min_ttl`, and the rest observe the
        // freshly-cached token and return it without re-fetching.
        let cache = Arc::new(TokenCache::default());
        static COUNTER: AtomicU32 = AtomicU32::new(0);

        async fn get_token() -> Result<TemporaryToken<String>, String> {
            COUNTER.fetch_add(1, Ordering::SeqCst);
            // Simulate a slow network fetch so callers genuinely overlap.
            tokio::time::sleep(Duration::from_millis(20)).await;
            Ok::<_, String>(create_token(Some(Duration::from_secs(3600))))
        }

        let mut handles = Vec::new();
        for _ in 0..10 {
            let cache = Arc::clone(&cache);
            handles.push(tokio::spawn(async move {
                cache.get_or_insert_with(get_token).await.unwrap()
            }));
        }
        for h in handles {
            h.await.unwrap();
        }

        // Exactly one fetch despite ten concurrent callers.
        assert_eq!(COUNTER.load(Ordering::SeqCst), 1);
    }
}