Skip to main content

kellnr_common/
token_cache.rs

1use std::time::Duration;
2
3use moka::future::Cache;
4
5#[derive(Clone, Debug, PartialEq, Eq)]
6pub struct CachedTokenData {
7    pub user: String,
8    pub is_admin: bool,
9    pub is_read_only: bool,
10}
11
12pub struct TokenCacheManager {
13    cache: Option<Cache<String, CachedTokenData>>,
14}
15
16impl TokenCacheManager {
17    /// Creates a new `TokenCacheManager`.
18    ///
19    /// # Arguments
20    /// * `enabled` - Whether caching is enabled
21    /// * `ttl_seconds` - Time-to-live for cached tokens. Lower values provide
22    ///   better security (revoked tokens expire faster) but increase database load.
23    /// * `max_capacity` - Maximum number of tokens to cache
24    pub fn new(enabled: bool, ttl_seconds: u64, max_capacity: u64) -> Self {
25        let cache = if enabled {
26            Some(
27                Cache::builder()
28                    .max_capacity(max_capacity)
29                    .time_to_live(Duration::from_secs(ttl_seconds))
30                    .build(),
31            )
32        } else {
33            None
34        };
35
36        Self { cache }
37    }
38
39    pub async fn get(&self, token: &str) -> Option<CachedTokenData> {
40        match &self.cache {
41            Some(cache) => cache.get(token).await,
42            None => None,
43        }
44    }
45
46    pub async fn insert(&self, token: String, data: CachedTokenData) {
47        if let Some(cache) = &self.cache {
48            cache.insert(token, data).await;
49        }
50    }
51
52    pub fn invalidate_all(&self) {
53        if let Some(cache) = &self.cache {
54            cache.invalidate_all();
55        }
56    }
57
58    /// Invalidate a single token from the cache.
59    pub async fn invalidate(&self, token: &str) {
60        if let Some(cache) = &self.cache {
61            cache.invalidate(token).await;
62        }
63    }
64
65    pub fn is_enabled(&self) -> bool {
66        self.cache.is_some()
67    }
68}
69
70#[cfg(test)]
71mod tests {
72    use super::*;
73
74    #[tokio::test]
75    async fn test_cache_disabled() {
76        let cache = TokenCacheManager::new(false, 60, 100);
77        assert!(!cache.is_enabled());
78
79        let data = CachedTokenData {
80            user: "test_user".to_string(),
81            is_admin: false,
82            is_read_only: false,
83        };
84
85        cache.insert("token123".to_string(), data).await;
86        assert!(cache.get("token123").await.is_none());
87    }
88
89    #[tokio::test]
90    async fn test_cache_enabled() {
91        let cache = TokenCacheManager::new(true, 60, 100);
92        assert!(cache.is_enabled());
93
94        let data = CachedTokenData {
95            user: "test_user".to_string(),
96            is_admin: true,
97            is_read_only: false,
98        };
99
100        cache.insert("token123".to_string(), data.clone()).await;
101
102        let retrieved = cache.get("token123").await;
103        assert!(retrieved.is_some());
104        let retrieved = retrieved.unwrap();
105        assert_eq!(retrieved.user, "test_user");
106        assert!(retrieved.is_admin);
107        assert!(!retrieved.is_read_only);
108    }
109
110    #[tokio::test]
111    async fn test_cache_miss() {
112        let cache = TokenCacheManager::new(true, 60, 100);
113        assert!(cache.get("nonexistent").await.is_none());
114    }
115
116    #[tokio::test]
117    async fn test_invalidate_all() {
118        let cache = TokenCacheManager::new(true, 60, 100);
119
120        let data = CachedTokenData {
121            user: "test_user".to_string(),
122            is_admin: false,
123            is_read_only: true,
124        };
125
126        cache.insert("token1".to_string(), data.clone()).await;
127        cache.insert("token2".to_string(), data).await;
128
129        assert!(cache.get("token1").await.is_some());
130        assert!(cache.get("token2").await.is_some());
131
132        cache.invalidate_all();
133
134        // Note: moka's invalidate_all is async internally and may not be immediately visible
135        // In production use, entries will be invalidated lazily
136    }
137}