salvo_jwt_auth/oidc/
cache.rs

1// port from https://github.com/fergus-hou/oidc_jwt_validator/blob/master/src/cache.rs
2
3use std::{
4    collections::HashMap,
5    sync::{
6        atomic::{AtomicBool, AtomicU64, Ordering},
7        Arc,
8    },
9    time::{Duration, Instant},
10};
11
12use jsonwebtoken::jwk::JwkSet;
13use jsonwebtoken::Validation;
14use salvo_core::http::header::HeaderValue;
15
16use super::{current_time, decode_jwk, DecodingInfo, JwkSetFetch};
17
18/// Determines settings about updating the cached JWKS data.
19/// The JWKS will be lazily revalidated every time validator validates a token.
20#[derive(Debug, Clone, Copy, PartialEq, Eq)]
21pub struct CachePolicy {
22    /// Time in Seconds to refresh the JWKS from the OIDC Provider
23    /// Default/Minimum value: 1 Second
24    pub max_age: Duration,
25    /// The amount of time a s
26    pub stale_while_revalidate: Option<Duration>,
27    /// The amount of time the stale JWKS data should be valid for if we are unable to re-validate it from the URL.
28    /// Minimum Value: 60 Seconds
29    pub stale_if_error: Option<Duration>,
30}
31
32impl CachePolicy {
33    /// Create a new cache policy from the header value of the Cache-Control header
34    #[must_use] pub fn from_header_val(value: Option<&HeaderValue>) -> Self {
35        // Initialize the default config of polling every second
36        let mut config = Self::default();
37
38        if let Some(value) = value {
39            if let Ok(value) = value.to_str() {
40                config.parse_str(value);
41            }
42        }
43        config
44    }
45
46    fn parse_str(&mut self, value: &str) {
47        // Iterate over every token in the header value
48        for token in value.split(',') {
49            // split them into whitespace trimmed pairs
50            let (key, val) = {
51                let mut split = token.split('=').map(str::trim);
52                (split.next(), split.next())
53            };
54            //Modify the default config based on the values that matter
55            //Any values here would be more permissive than the default behavior
56            match (key, val) {
57                (Some("max-age"), Some(val)) => {
58                    if let Ok(secs) = val.parse::<u64>() {
59                        self.max_age = Duration::from_secs(secs);
60                    }
61                }
62                (Some("stale-while-revalidate"), Some(val)) => {
63                    if let Ok(secs) = val.parse::<u64>() {
64                        self.stale_while_revalidate = Some(Duration::from_secs(secs));
65                    }
66                }
67                (Some("stale-if-error"), Some(val)) => {
68                    if let Ok(secs) = val.parse::<u64>() {
69                        self.stale_if_error = Some(Duration::from_secs(secs));
70                    }
71                }
72                _ => {},
73            };
74        }
75    }
76}
77
78impl Default for CachePolicy {
79    fn default() -> Self {
80        Self {
81            max_age: Duration::from_secs(1),
82            stale_while_revalidate: Some(Duration::from_secs(1)),
83            stale_if_error: Some(Duration::from_secs(60)),
84        }
85    }
86}
87
88/// The update action of cache.
89#[derive(Debug, Clone, Copy, PartialEq, Eq)]
90pub enum UpdateAction {
91    /// We checked the JWKS uri and it was the same as the last time we refreshed it so no action was taken
92    NoUpdate,
93    /// We checked the JWKS uri and it was different so we updated our local cache
94    JwksUpdate,
95    /// The JWKS Uri responded with a different cache-control header
96    CacheUpdate(CachePolicy),
97    /// The JWKS Uri responded with a different cache-control header and the JWKS was updated
98    JwksAndCacheUpdate(CachePolicy),
99}
100
101/// Helper struct for determining when our cache needs to be re-validated
102/// Utilizes atomics to prevent write-locking as much as possible
103#[derive(Debug)]
104pub struct CacheState {
105    last_update: AtomicU64,
106    is_revalidating: AtomicBool,
107    is_error: AtomicBool,
108}
109
110impl CacheState {
111    /// Create a new `CacheState`
112    #[must_use] pub fn new() -> Self {
113        Self {
114            last_update: AtomicU64::new(current_time()),
115            is_revalidating: AtomicBool::new(false),
116            is_error: AtomicBool::new(false),
117        }
118    }
119    /// Check is the cache is error
120    pub fn is_error(&self) -> bool {
121        self.is_error.load(Ordering::SeqCst)
122    }
123    /// Set the cache is error
124    pub fn set_is_error(&self, value: bool) {
125        self.is_error.store(value, Ordering::SeqCst);
126    }
127
128    /// Get the cache last updated timestamp
129    pub fn last_update(&self) -> u64 {
130        self.last_update.load(Ordering::SeqCst)
131    }
132    /// Set the cache last updated timestamp
133    pub fn set_last_update(&self, timestamp: u64) {
134        self.last_update.store(timestamp, Ordering::SeqCst);
135    }
136
137    /// Check if the cache is revalidating
138    pub fn is_revalidating(&self) -> bool {
139        self.is_revalidating.load(Ordering::SeqCst)
140    }
141    /// Set the cache is revalidating
142    pub fn set_is_revalidating(&self, value: bool) {
143        self.is_revalidating.store(value, Ordering::SeqCst);
144    }
145}
146
147impl Default for CacheState {
148    fn default() -> Self {
149        Self::new()
150    }
151}
152
153/// Helper Struct for storing
154#[derive(Debug)]
155pub struct JwkSetStore {
156    /// The current JWKS
157    pub jwks: JwkSet,
158    decoding_map: HashMap<String, Arc<DecodingInfo>>,
159    /// The cache policy for this store
160    pub cache_policy: CachePolicy,
161    validation: Validation,
162}
163
164impl JwkSetStore {
165    /// Create a new `JwkSetStore`
166    #[must_use] pub fn new(jwks: JwkSet, cache_policy: CachePolicy, validation: Validation) -> Self {
167        Self {
168            jwks,
169            decoding_map: HashMap::new(),
170            cache_policy,
171            validation,
172        }
173    }
174
175    fn update_jwks(&mut self, new_jwks: JwkSet) {
176        self.jwks = new_jwks;
177        let keys = self
178            .jwks
179            .keys
180            .iter()
181            .filter_map(|i| decode_jwk(i, &self.validation).ok());
182        // Clear our cache of decoding keys
183        self.decoding_map.clear();
184        // Load the keys back into our hashmap cache.
185        for key in keys {
186            self.decoding_map.insert(key.0, Arc::new(key.1));
187        }
188    }
189
190    /// Get the DecodingInfo for a given kid
191    #[must_use] pub fn get_key(&self, kid: &str) -> Option<Arc<DecodingInfo>> {
192        self.decoding_map.get(kid).cloned()
193    }
194
195    pub(crate) fn update_fetch(&mut self, fetch: JwkSetFetch) -> UpdateAction {
196        tracing::debug!("Decoding JWKS");
197        let time = Instant::now();
198        let new_jwks = fetch.jwks;
199        // If we didn't parse out a cache policy from the last request
200        // Assume that it's the same as the last
201        let cache_policy = fetch.cache_policy.unwrap_or(self.cache_policy);
202        let result = match (self.jwks == new_jwks, self.cache_policy == cache_policy) {
203            // Everything is the same
204            (true, true) => {
205                tracing::debug!("JWKS Content has not changed since last update");
206                UpdateAction::NoUpdate
207            }
208            // The JWKS changed but the cache policy hasn't
209            (false, true) => {
210                tracing::info!("JWKS Content has changed since last update");
211                self.update_jwks(new_jwks);
212                UpdateAction::JwksUpdate
213            }
214            // The cache policy changed, but the JWKS hasn't
215            (true, false) => {
216                self.cache_policy = cache_policy;
217                UpdateAction::CacheUpdate(cache_policy)
218            }
219            // Both the cache and the JWKS have changed
220            (false, false) => {
221                tracing::info!("cache-control header and JWKS content has changed since last update");
222                self.update_jwks(new_jwks);
223                self.cache_policy = cache_policy;
224                UpdateAction::JwksAndCacheUpdate(cache_policy)
225            }
226        };
227        let elapsed = time.elapsed();
228        tracing::debug!("Decoded and parsed JWKS in {:#?}", elapsed);
229        result
230    }
231}
232
233#[cfg(test)]
234mod tests {
235    use super::*;
236
237    #[test]
238    fn test_cache_policy_default() {
239        let policy = CachePolicy::default();
240        assert_eq!(policy.max_age, Duration::from_secs(1));
241        assert_eq!(policy.stale_while_revalidate, Some(Duration::from_secs(1)));
242        assert_eq!(policy.stale_if_error, Some(Duration::from_secs(60)));
243    }
244
245    #[test]
246    fn test_cache_policy_from_header_val_max_age() {
247        let header_value = HeaderValue::from_static("max-age=3600");
248        let policy = CachePolicy::from_header_val(Some(&header_value));
249        assert_eq!(policy.max_age, Duration::from_secs(3600));
250        assert_eq!(policy.stale_while_revalidate, Some(Duration::from_secs(1)));
251        assert_eq!(policy.stale_if_error, Some(Duration::from_secs(60)));
252    }
253
254    #[test]
255    fn test_cache_policy_from_header_val_stale_while_revalidate() {
256        let header_value = HeaderValue::from_static("max-age=3600, stale-while-revalidate=60");
257        let policy = CachePolicy::from_header_val(Some(&header_value));
258        assert_eq!(policy.max_age, Duration::from_secs(3600));
259        assert_eq!(policy.stale_while_revalidate, Some(Duration::from_secs(60)));
260        assert_eq!(policy.stale_if_error, Some(Duration::from_secs(60)));
261    }
262
263    #[test]
264    fn test_cache_policy_from_header_val_stale_if_error() {
265        let header_value = HeaderValue::from_static("max-age=3600, stale-if-error=120");
266        let policy = CachePolicy::from_header_val(Some(&header_value));
267        assert_eq!(policy.max_age, Duration::from_secs(3600));
268        assert_eq!(policy.stale_while_revalidate, Some(Duration::from_secs(1)));
269        assert_eq!(policy.stale_if_error, Some(Duration::from_secs(120)));
270    }
271
272    #[test]
273    fn test_cache_policy_from_header_val_all() {
274        let header_value = HeaderValue::from_static("max-age=3600, stale-while-revalidate=60, stale-if-error=120");
275        let policy = CachePolicy::from_header_val(Some(&header_value));
276        assert_eq!(policy.max_age, Duration::from_secs(3600));
277        assert_eq!(policy.stale_while_revalidate, Some(Duration::from_secs(60)));
278        assert_eq!(policy.stale_if_error, Some(Duration::from_secs(120)));
279    }
280
281    #[test]
282    fn test_cache_policy_from_header_val_none() {
283        let policy = CachePolicy::from_header_val(None);
284        assert_eq!(policy, CachePolicy::default());
285    }
286
287    #[test]
288    fn test_cache_policy_from_header_val_invalid() {
289        let header_value = HeaderValue::from_static("invalid-directive");
290        let policy = CachePolicy::from_header_val(Some(&header_value));
291        assert_eq!(policy, CachePolicy::default());
292    }
293}
salvo_jwt_auth/oidc/cache.rs

salvo_jwt_auth/oidc/
cache.rs
