jwk_simple/jwks/cache.rs
1//! Caching traits and wrappers for key stores.
2//!
3//! This module provides a [`KeyCache`] trait for caching key sets,
4//! and a [`CachedKeyStore`] wrapper that combines any cache with any key store.
5//!
6//! For a ready-to-use Moka implementation, enable the `moka` feature
7//! and use `MokaKeyCache`.
8
9#[cfg(all(feature = "moka", not(target_arch = "wasm32")))]
10pub mod moka;
11
12use crate::error::Result;
13use crate::jwk::Key;
14
15use super::{KeySet, KeyStore};
16
17/// A trait for caching key sets used by [`CachedKeyStore`].
18///
19/// This trait is primarily an extension point for library integrators who need a
20/// custom cache backend (in-memory, KV store, persistent storage, etc.). Most users
21/// should use [`CachedKeyStore`] with a provided cache implementation, such as
22/// `MokaKeyCache`.
23///
24/// `CachedKeyStore` does not coordinate concurrent refreshes. Cache backends used in
25/// concurrent contexts should provide their own internal synchronization and atomicity
26/// guarantees for `get`/`set`/`clear` operations.
27///
28/// The cache stores the entire [`KeySet`] as a single unit, which matches how JWKS
29/// endpoints work (they always return the full set). This avoids the N+1 fetch problem
30/// that per-key caching would cause.
31#[cfg_attr(target_arch = "wasm32", async_trait::async_trait(?Send))]
32#[cfg_attr(not(target_arch = "wasm32"), async_trait::async_trait)]
33pub trait KeyCache {
34 /// Gets the cached key set.
35 ///
36 /// Returns `None` if the cache is empty or has expired.
37 async fn get(&self) -> Result<Option<KeySet>>;
38
39 /// Stores a key set in the cache.
40 async fn set(&self, keyset: KeySet) -> Result<()>;
41
42 /// Clears the cache.
43 async fn clear(&self) -> Result<()>;
44}
45
46/// A caching wrapper for any [`KeyStore`] implementation.
47///
48/// This wrapper uses the cache-aside pattern: it first checks the cache for the key set,
49/// and only fetches from the underlying store on a cache miss. Retrieved key sets
50/// are then stored in the cache for future requests.
51///
52/// This is the recommended entry point for cached JWKS retrieval.
53///
54/// When looking up a key by ID, if the cached key set doesn't contain the requested key,
55/// the store refetches from the underlying source. This handles key rotation gracefully:
56/// newly added keys are discovered automatically without waiting for cache expiration.
57///
58/// # Async and Send Bounds
59///
60/// On native targets, the `KeyStore` implementation for `CachedKeyStore`
61/// requires `Send + Sync` for both cache and source store for thread-safe usage.
62/// On `wasm32`, these bounds are relaxed to match the `?Send` async model.
63///
64/// # Examples
65///
66/// ```ignore
67/// use jwk_simple::jwks::{CachedKeyStore, HttpKeyStore, KeyStore, MokaKeyCache};
68/// use std::time::Duration;
69///
70/// // Create a cached remote key store
71/// let cache = MokaKeyCache::new(Duration::from_secs(300));
72/// let remote = HttpKeyStore::new("https://example.com/.well-known/jwks.json")?;
73/// let cached = CachedKeyStore::new(cache, remote);
74///
75/// // First call fetches from remote, caches the key set
76/// let key = cached.get_key("kid").await?;
77///
78/// // Subsequent calls use the cache
79/// let key = cached.get_key("kid").await?;
80/// ```
81#[derive(Debug)]
82pub struct CachedKeyStore<C: KeyCache, S: KeyStore> {
83 cache: C,
84 store: S,
85}
86
87impl<C: KeyCache, S: KeyStore> CachedKeyStore<C, S> {
88 /// Creates a cache wrapper around the provided key store.
89 pub fn new(cache: C, store: S) -> Self {
90 Self { cache, store }
91 }
92
93 /// Returns the configured cache backend for inspection or cache management.
94 pub fn cache(&self) -> &C {
95 &self.cache
96 }
97
98 /// Returns the underlying key store wrapped by this cache layer.
99 pub fn store(&self) -> &S {
100 &self.store
101 }
102}
103
104#[cfg(not(target_arch = "wasm32"))]
105#[async_trait::async_trait]
106impl<C, S> KeyStore for CachedKeyStore<C, S>
107where
108 C: KeyCache + Send + Sync,
109 S: KeyStore + Send + Sync,
110{
111 async fn get_keyset(&self) -> Result<KeySet> {
112 // Try cache first
113 if let Some(keyset) = self.cache.get().await? {
114 return Ok(keyset);
115 }
116
117 // No cache hit: fetch from source and cache
118 let keyset = self.store.get_keyset().await?;
119 self.cache.set(keyset.clone()).await?;
120
121 Ok(keyset)
122 }
123
124 async fn get_key(&self, kid: &str) -> Result<Option<Key>> {
125 // Try cache first
126 if let Some(keyset) = self.cache.get().await?
127 && let Some(key) = keyset.get_by_kid(kid)
128 {
129 return Ok(Some(key.clone()));
130 }
131
132 // Key not in cached set: refetch
133 let keyset = self.store.get_keyset().await?;
134 self.cache.set(keyset.clone()).await?;
135
136 Ok(keyset.get_by_kid(kid).cloned())
137 }
138}
139
140#[cfg(target_arch = "wasm32")]
141#[async_trait::async_trait(?Send)]
142impl<C, S> KeyStore for CachedKeyStore<C, S>
143where
144 C: KeyCache,
145 S: KeyStore,
146{
147 async fn get_keyset(&self) -> Result<KeySet> {
148 // Try cache first
149 if let Some(keyset) = self.cache.get().await? {
150 return Ok(keyset);
151 }
152
153 // No cache hit: fetch from source and cache
154 let keyset = self.store.get_keyset().await?;
155 self.cache.set(keyset.clone()).await?;
156
157 Ok(keyset)
158 }
159
160 async fn get_key(&self, kid: &str) -> Result<Option<Key>> {
161 // Try cache first
162 if let Some(keyset) = self.cache.get().await?
163 && let Some(key) = keyset.get_by_kid(kid)
164 {
165 return Ok(Some(key.clone()));
166 }
167
168 // Key not in cached set: refetch
169 let keyset = self.store.get_keyset().await?;
170 self.cache.set(keyset.clone()).await?;
171
172 Ok(keyset.get_by_kid(kid).cloned())
173 }
174}