salvo-jwt-auth 1.1.0

// port from https://github.com/fergus-hou/oidc_jwt_validator/blob/master/src/cache.rs

use std::{
    collections::HashMap,
    sync::{
        atomic::{AtomicBool, AtomicU64, Ordering},
        Arc, Mutex,
    },
    time::{Duration, Instant},
};

use jsonwebtoken::jwk::JwkSet;
use jsonwebtoken::Validation;
use salvo_core::http::header::HeaderValue;

use super::{current_time, decode_jwk, DecodingInfo, JwkSetFetch};

/// Determines settings about updating the cached JWKS data.
/// The JWKS will be lazily revalidated every time validator validates a token.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct CachePolicy {
    /// Time in Seconds to refresh the JWKS from the OIDC Provider
    /// Default/Minimum value: 1 Second
    pub max_age: Duration,
    /// The amount of time the stale JWKS data can be used while it is being revalidated.
    pub stale_while_revalidate: Option<Duration>,
    /// The amount of time the stale JWKS data should be valid for if we are unable to re-validate it from the URL.
    /// Defaults to 60 seconds. An explicit zero disables stale-on-error use.
    pub stale_if_error: Option<Duration>,
}

impl CachePolicy {
    /// Create a new cache policy from the header value of the Cache-Control header
    #[must_use]
    pub fn from_header_val(value: Option<&HeaderValue>) -> Self {
        Self::from_header_values(value)
    }

    /// Creates a policy from all Cache-Control field values in their received order.
    ///
    /// Repeated field lines are combined with commas before parsing, preserving
    /// quoted extension arguments and opaque bytes. No values yields the default policy.
    #[must_use]
    pub fn from_header_values<'a>(values: impl IntoIterator<Item = &'a HeaderValue>) -> Self {
        // Initialize the default config of polling every second
        let mut config = Self::default();
        let mut values = values.into_iter();
        if let Some(first) = values.next() {
            if let Some(second) = values.next() {
                let mut combined = first.as_bytes().to_vec();
                for value in std::iter::once(second).chain(values) {
                    combined.push(b',');
                    combined.extend_from_slice(value.as_bytes());
                }
                config.parse_bytes(&combined);
            } else {
                config.parse_bytes(first.as_bytes());
            }
        }
        config
    }

    fn parse_bytes(&mut self, value: &[u8]) {
        // Upper bound for any cache directive (30 days). Without it, a malicious or
        // misconfigured IdP returning a huge `max-age` could overflow the downstream
        // `fetched + max_age` arithmetic (debug panic / release wrap) and pin a stale
        // JWKS indefinitely.
        const MAX_CACHE_AGE_SECS: u64 = 30 * 24 * 60 * 60;
        // RFC 9111 section 5.2: commas in quoted extension arguments are not
        // directive boundaries, and directive names are case insensitive. Work
        // on bytes so a legal opaque extension does not discard the whole field.
        let mut quoted = false;
        let mut escaped = false;
        for token in value.split(|byte| {
            if escaped {
                escaped = false;
                return false;
            }
            if quoted && *byte == b'\\' {
                escaped = true;
            } else if *byte == b'"' {
                quoted = !quoted;
            }
            *byte == b',' && !quoted
        }) {
            let Some(equal) = token.iter().position(|byte| *byte == b'=') else {
                continue;
            };
            let key = token[..equal].trim_ascii();
            let Some(secs) = delta_seconds(token[equal + 1..].trim_ascii()) else {
                continue;
            };
            // Modify the default config based on the values that matter.
            // Any values here would be more permissive than the default behavior.
            if key.eq_ignore_ascii_case(b"max-age") {
                // Keep at least 1s so `max-age=0` cannot make every request
                // revalidate and amplify load back onto the IdP.
                self.max_age = Duration::from_secs(secs.clamp(1, MAX_CACHE_AGE_SECS));
            } else if key.eq_ignore_ascii_case(b"stale-while-revalidate") {
                self.stale_while_revalidate =
                    Some(Duration::from_secs(secs.min(MAX_CACHE_AGE_SECS)));
            } else if key.eq_ignore_ascii_case(b"stale-if-error") {
                self.stale_if_error = Some(Duration::from_secs(secs.min(MAX_CACHE_AGE_SECS)));
            }
        }
    }
}

fn delta_seconds(value: &[u8]) -> Option<u64> {
    let quoted = value.first() == Some(&b'"');
    let value = if quoted {
        value.strip_prefix(b"\"")?.strip_suffix(b"\"")?
    } else {
        value
    };
    let mut bytes = value.iter().copied();
    let mut result = 0u64;
    let mut digits = 0;
    while let Some(byte) = bytes.next() {
        let byte = if quoted && byte == b'\\' {
            bytes.next()?
        } else {
            byte
        };
        if !byte.is_ascii_digit() {
            return None;
        }
        result = result
            .checked_mul(10)?
            .checked_add(u64::from(byte - b'0'))?;
        digits += 1;
    }
    (digits != 0).then_some(result)
}

impl Default for CachePolicy {
    fn default() -> Self {
        Self {
            max_age: Duration::from_secs(1),
            stale_while_revalidate: Some(Duration::from_secs(1)),
            stale_if_error: Some(Duration::from_secs(60)),
        }
    }
}

/// The update action of cache.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum UpdateAction {
    /// We checked the JWKS uri and it was the same as the last time we refreshed it so no action
    /// was taken
    NoUpdate,
    /// We checked the JWKS uri and it was different so we updated our local cache
    JwksUpdate,
    /// The JWKS Uri responded with a different cache-control header
    CacheUpdate(CachePolicy),
    /// The JWKS Uri responded with a different cache-control header and the JWKS was updated
    JwksAndCacheUpdate(CachePolicy),
}

/// Helper struct for determining when our cache needs to be re-validated
/// Utilizes atomics to prevent write-locking as much as possible
#[derive(Debug)]
pub struct CacheState {
    last_update: AtomicU64,
    is_revalidating: AtomicBool,
    is_error: AtomicBool,
    last_revalidation: Mutex<Option<Instant>>,
}

impl CacheState {
    /// Create a new `CacheState`
    #[must_use]
    pub fn new() -> Self {
        Self {
            last_update: AtomicU64::new(current_time()),
            is_revalidating: AtomicBool::new(false),
            is_error: AtomicBool::new(false),
            last_revalidation: Mutex::new(None),
        }
    }
    /// Check is the cache is error
    pub fn is_error(&self) -> bool {
        self.is_error.load(Ordering::SeqCst)
    }
    /// Set the cache is error
    pub fn set_is_error(&self, value: bool) {
        self.is_error.store(value, Ordering::SeqCst);
    }

    /// Get the cache last updated timestamp
    pub fn last_update(&self) -> u64 {
        self.last_update.load(Ordering::SeqCst)
    }
    /// Set the cache last updated timestamp
    pub fn set_last_update(&self, timestamp: u64) {
        self.last_update.store(timestamp, Ordering::SeqCst);
    }

    /// Check if the cache is revalidating
    pub fn is_revalidating(&self) -> bool {
        self.is_revalidating.load(Ordering::Acquire)
    }
    /// Try to mark the cache as revalidating.
    pub fn begin_revalidation(&self) -> bool {
        self.is_revalidating
            .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
            .is_ok()
    }

    pub(super) fn begin_revalidation_after(&self, interval: Duration) -> bool {
        let mut last = self
            .last_revalidation
            .lock()
            .unwrap_or_else(|poisoned| poisoned.into_inner());
        if last.is_some_and(|last| last.elapsed() < interval) || !self.begin_revalidation() {
            return false;
        }
        *last = Some(Instant::now());
        true
    }

    /// Mark the cache as no longer revalidating.
    pub fn finish_revalidation(&self) {
        let mut last = self
            .last_revalidation
            .lock()
            .unwrap_or_else(|poisoned| poisoned.into_inner());
        if last.is_some() {
            // Start the cooldown after the attempt, including failed attempts,
            // so a slow or failing provider cannot be polled without a break.
            *last = Some(Instant::now());
        }
        self.is_revalidating.store(false, Ordering::Release);
    }
}

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

/// Helper Struct for storing
#[derive(Debug)]
pub struct JwkSetStore {
    /// The current JWKS
    pub jwks: JwkSet,
    decoding_map: HashMap<String, Arc<DecodingInfo>>,
    /// The cache policy for this store
    pub cache_policy: CachePolicy,
    validation: Validation,
    allow_symmetric_jwks: bool,
}

impl JwkSetStore {
    /// Create a new `JwkSetStore`
    #[must_use]
    pub fn new(jwks: JwkSet, cache_policy: CachePolicy, validation: Validation) -> Self {
        Self {
            jwks,
            decoding_map: HashMap::new(),
            cache_policy,
            validation,
            allow_symmetric_jwks: false,
        }
    }

    /// Allow symmetric (`kty: "oct"` / HS*) JWKs from this JWKS endpoint.
    ///
    /// **Disabled by default.** OIDC JWKS endpoints normally publish only
    /// asymmetric public keys (RSA / EC / OKP); accepting symmetric keys
    /// widens the trust surface because a leaked or substituted secret
    /// would let an attacker mint tokens the decoder accepts. Enable this
    /// only when the issuer is known to publish HS* keys and is trusted to
    /// keep them secret.
    #[must_use]
    pub fn with_allow_symmetric_jwks(mut self, allow: bool) -> Self {
        self.allow_symmetric_jwks = allow;
        self
    }

    fn update_jwks(&mut self, new_jwks: JwkSet) {
        self.jwks = new_jwks;
        let allow_symmetric = self.allow_symmetric_jwks;
        let keys = self
            .jwks
            .keys
            .iter()
            .filter_map(|i| decode_jwk(i, &self.validation, allow_symmetric).ok());
        // Clear our cache of decoding keys
        self.decoding_map.clear();
        // Load the keys back into our hashmap cache.
        for key in keys {
            self.decoding_map.insert(key.0, Arc::new(key.1));
        }
    }

    /// Get the DecodingInfo for a given kid
    #[must_use]
    pub fn get_key(&self, kid: &str) -> Option<Arc<DecodingInfo>> {
        self.decoding_map.get(kid).cloned()
    }

    pub(crate) fn update_fetch(&mut self, fetch: JwkSetFetch) -> UpdateAction {
        tracing::debug!("Decoding JWKS");
        let time = Instant::now();
        let new_jwks = fetch.jwks;
        // If we didn't parse out a cache policy from the last request
        // Assume that it is the same as the last
        let cache_policy = fetch.cache_policy.unwrap_or(self.cache_policy);
        let result = match (self.jwks == new_jwks, self.cache_policy == cache_policy) {
            // Everything is the same
            (true, true) => {
                tracing::debug!("JWKS Content has not changed since last update");
                UpdateAction::NoUpdate
            }
            // The JWKS changed but the cache policy hasn't
            (false, true) => {
                tracing::info!("JWKS Content has changed since last update");
                self.update_jwks(new_jwks);
                UpdateAction::JwksUpdate
            }
            // The cache policy changed, but the JWKS hasn't
            (true, false) => {
                self.cache_policy = cache_policy;
                UpdateAction::CacheUpdate(cache_policy)
            }
            // Both the cache and the JWKS have changed
            (false, false) => {
                tracing::info!("cache-control header and JWKS content has changed since last update");
                self.update_jwks(new_jwks);
                self.cache_policy = cache_policy;
                UpdateAction::JwksAndCacheUpdate(cache_policy)
            }
        };
        let elapsed = time.elapsed();
        tracing::debug!("Decoded and parsed JWKS in {:#?}", elapsed);
        result
    }
}

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

    #[test]
    fn test_cache_policy_default() {
        let policy = CachePolicy::default();
        assert_eq!(policy.max_age, Duration::from_secs(1));
        assert_eq!(policy.stale_while_revalidate, Some(Duration::from_secs(1)));
        assert_eq!(policy.stale_if_error, Some(Duration::from_secs(60)));
    }

    #[test]
    fn test_cache_policy_from_header_val_max_age() {
        let header_value = HeaderValue::from_static("max-age=3600");
        let policy = CachePolicy::from_header_val(Some(&header_value));
        assert_eq!(policy.max_age, Duration::from_secs(3600));
        assert_eq!(policy.stale_while_revalidate, Some(Duration::from_secs(1)));
        assert_eq!(policy.stale_if_error, Some(Duration::from_secs(60)));
    }

    #[test]
    fn test_cache_policy_from_header_val_stale_while_revalidate() {
        let header_value = HeaderValue::from_static("max-age=3600, stale-while-revalidate=60");
        let policy = CachePolicy::from_header_val(Some(&header_value));
        assert_eq!(policy.max_age, Duration::from_secs(3600));
        assert_eq!(policy.stale_while_revalidate, Some(Duration::from_secs(60)));
        assert_eq!(policy.stale_if_error, Some(Duration::from_secs(60)));
    }

    #[test]
    fn test_cache_policy_from_header_val_stale_if_error() {
        let header_value = HeaderValue::from_static("max-age=3600, stale-if-error=120");
        let policy = CachePolicy::from_header_val(Some(&header_value));
        assert_eq!(policy.max_age, Duration::from_secs(3600));
        assert_eq!(policy.stale_while_revalidate, Some(Duration::from_secs(1)));
        assert_eq!(policy.stale_if_error, Some(Duration::from_secs(120)));
    }

    #[test]
    fn test_cache_policy_from_header_val_all() {
        let header_value = HeaderValue::from_static("max-age=3600, stale-while-revalidate=60, stale-if-error=120");
        let policy = CachePolicy::from_header_val(Some(&header_value));
        assert_eq!(policy.max_age, Duration::from_secs(3600));
        assert_eq!(policy.stale_while_revalidate, Some(Duration::from_secs(60)));
        assert_eq!(policy.stale_if_error, Some(Duration::from_secs(120)));
    }

    #[test]
    fn test_cache_policy_from_header_val_none() {
        let policy = CachePolicy::from_header_val(None);
        assert_eq!(policy, CachePolicy::default());
    }

    #[test]
    fn test_cache_policy_from_header_val_invalid() {
        let header_value = HeaderValue::from_static("invalid-directive");
        let policy = CachePolicy::from_header_val(Some(&header_value));
        assert_eq!(policy, CachePolicy::default());
    }

    #[test]
    fn test_cache_state_revalidation_claim_is_atomic() {
        let state = CacheState::new();

        assert!(state.begin_revalidation());
        assert!(state.is_revalidating());
        assert!(!state.begin_revalidation());

        state.finish_revalidation();
        assert!(!state.is_revalidating());
        assert!(state.begin_revalidation());
    }

    #[test]
    fn revalidation_cooldown_is_shared_and_expires() {
        let state = Arc::new(CacheState::new());
        let clone = Arc::clone(&state);
        let interval = Duration::from_secs(1);
        assert!(state.begin_revalidation_after(interval));
        assert!(!clone.begin_revalidation_after(interval));
        state.finish_revalidation();
        assert!(!clone.begin_revalidation_after(interval));

        *state.last_revalidation.lock().unwrap() = Some(Instant::now() - interval);
        assert!(clone.begin_revalidation_after(interval));
    }
}