Skip to main content

supabase_jwt/
jwks.rs

1//! # JWKS (JSON Web Key Set) Smart Cache Management Module
2//!
3//! A JWKS caching system for Supabase Auth with graceful fallback.
4//!
5//! ## Features
6//! - **Smart Caching**: 24-hour cache with 7-day fallback
7//! - **Graceful Fallback**: Uses expired cache during network failures
8//! - **Concurrency Safe**: Prevents duplicate fetching
9//!
10//! ## Caching Strategy
11//! 1. Uses valid cache within 24 hours
12//! 2. Refreshes from remote when expired
13//! 3. Falls back to stale cache (up to 7 days) on network failure
14
15use crate::error::AuthError;
16use serde::{Deserialize, Serialize};
17use std::sync::{Arc, LazyLock};
18use tokio::sync::{Mutex, RwLock};
19
20/// Global HTTP client instance with a connection pool.
21static HTTP_CLIENT: LazyLock<reqwest::Client> = LazyLock::new(|| {
22    reqwest::Client::builder()
23        .timeout(std::time::Duration::from_secs(5)) // Set a 5-second timeout.
24        .pool_max_idle_per_host(10) // Increase the connection pool size.
25        .pool_idle_timeout(std::time::Duration::from_secs(30))
26        .build()
27        .expect("Failed to create HTTP client")
28});
29
30/// Represents a JSON Web Key (JWK).
31#[derive(Debug, Serialize, Deserialize, Clone)]
32pub struct Jwk {
33    /// Key ID.
34    pub kid: String,
35    /// Key type (e.g., "EC").
36    pub kty: String,
37    /// Algorithm (e.g., "ES256").
38    pub alg: Option<String>,
39    /// Key usage (e.g., "sig").
40    #[serde(rename = "use")]
41    pub key_use: Option<String>,
42    /// Key operations.
43    pub key_ops: Option<Vec<String>>,
44    /// Curve (for EC keys).
45    pub crv: Option<String>,
46    /// X coordinate (for EC keys).
47    pub x: Option<String>,
48    /// Y coordinate (for EC keys).
49    pub y: Option<String>,
50    /// RSA modulus (for RSA keys).
51    pub n: Option<String>,
52    /// RSA exponent (for RSA keys).
53    pub e: Option<String>,
54    /// Whether the key is extractable (Supabase-specific field).
55    pub ext: Option<bool>,
56}
57
58/// Represents the response from a JWKS endpoint.
59#[derive(Debug, Serialize, Deserialize, Clone)]
60pub struct JwksResponse {
61    /// A list of JSON Web Keys.
62    pub keys: Vec<Jwk>,
63}
64
65/// Manages caching and retrieval of JWKS data from Supabase.
66const JWKS_CACHE_DURATION: u64 = 24 * 3600; // 24-hour normal cache.
67const JWKS_CACHE_MAX_AGE: u64 = 7 * 24 * 3600; // 7-day maximum cache for graceful fallback.
68
69#[derive(Debug, Clone)]
70pub struct JwksCache {
71    /// The cached JWKS data.
72    cache: Arc<RwLock<Option<JwksResponse>>>,
73    /// The expiration timestamp for the normal cache.
74    expires_at: Arc<RwLock<Option<u64>>>,
75    /// The timestamp when the cache was created, for calculating max age.
76    cached_at: Arc<RwLock<Option<u64>>>,
77    /// The JWKS endpoint URL.
78    jwks_url: String,
79    /// A mutex to prevent concurrent fetches.
80    fetch_mutex: Arc<Mutex<()>>,
81}
82
83impl JwksCache {
84    /// Creates a new `JwksCache` instance.
85    ///
86    /// # Arguments
87    ///
88    /// * `jwks_url` - The Supabase JWKS endpoint URL
89    pub fn new(jwks_url: &str) -> Self {
90        // Basic validation: ensure HTTPS is used.
91        if !jwks_url.starts_with("https://") {
92            tracing::warn!("JWKS URL should use HTTPS: {}", jwks_url);
93        }
94
95        Self {
96            cache: Arc::new(RwLock::new(None)),
97            expires_at: Arc::new(RwLock::new(None)),
98            cached_at: Arc::new(RwLock::new(None)),
99            jwks_url: jwks_url.to_string(),
100            fetch_mutex: Arc::new(Mutex::new(())),
101        }
102    }
103
104    /// Retrieves JWKS data with graceful fallback.
105    ///
106    /// Uses valid cache, refreshes if expired, or falls back to stale cache on failure.
107    pub async fn get_jwks(&self) -> Result<JwksResponse, AuthError> {
108        self.get_jwks_with_fallback().await
109    }
110
111    /// Implements the graceful fallback logic for retrieving JWKS data.
112    async fn get_jwks_with_fallback(&self) -> Result<JwksResponse, AuthError> {
113        // 1. Try to use the valid cache (within 24 hours).
114        if let Some(cached) = self.get_cached_jwks().await {
115            tracing::debug!("Using valid cached JWKS data");
116            return Ok(cached);
117        }
118
119        // 2. If the cache is expired, try to refresh it.
120        match self.fetch_fresh_jwks().await {
121            Ok(jwks) => {
122                tracing::info!("Successfully refreshed JWKS cache");
123                Ok(jwks)
124            }
125            Err(e) => {
126                tracing::warn!("Failed to refresh JWKS, attempting fallback: {:?}", e);
127                // 3. If refreshing fails, use the stale cache (up to 7 days old).
128                self.get_stale_cache().await
129            }
130        }
131    }
132
133    /// Retrieves valid cached data (within 24 hours).
134    async fn get_cached_jwks(&self) -> Option<JwksResponse> {
135        let now = std::time::SystemTime::now()
136            .duration_since(std::time::UNIX_EPOCH)
137            .unwrap()
138            .as_secs();
139
140        let expires_at = *self.expires_at.read().await;
141        if let Some(expires) = expires_at {
142            if now < expires {
143                return self.cache.read().await.clone();
144            }
145        }
146        None
147    }
148
149    /// Retrieves stale cache as fallback (up to 7 days old).
150    async fn get_stale_cache(&self) -> Result<JwksResponse, AuthError> {
151        let now = std::time::SystemTime::now()
152            .duration_since(std::time::UNIX_EPOCH)
153            .unwrap()
154            .as_secs();
155
156        let cached_at = *self.cached_at.read().await;
157        if let Some(cache_time) = cached_at {
158            if now - cache_time <= JWKS_CACHE_MAX_AGE {
159                if let Some(cached) = self.cache.read().await.clone() {
160                    tracing::warn!(
161                        "Using stale JWKS cache as fallback (age: {} hours)",
162                        (now - cache_time) / 3600
163                    );
164                    return Ok(cached);
165                }
166            }
167        }
168
169        let error_msg = "No valid JWKS cache available and network fetch failed";
170        tracing::error!("{}", error_msg);
171        Err(AuthError::JwksError(error_msg.to_string()))
172    }
173
174    /// Fetches fresh JWKS data from the remote endpoint.
175    async fn fetch_fresh_jwks(&self) -> Result<JwksResponse, AuthError> {
176        // Use a mutex to prevent concurrent fetches.
177        let _fetch_guard = self.fetch_mutex.lock().await;
178
179        // Double-check the cache in case it was updated while waiting for the lock.
180        if let Some(cached) = self.get_cached_jwks().await {
181            tracing::debug!("JWKS cache was updated while waiting for lock");
182            return Ok(cached);
183        }
184
185        let now = std::time::SystemTime::now()
186            .duration_since(std::time::UNIX_EPOCH)
187            .unwrap()
188            .as_secs();
189
190        // Fetch fresh data if the cache is expired or non-existent.
191        tracing::info!("Fetching fresh JWKS from: {}", self.jwks_url);
192
193        let response = HTTP_CLIENT.get(&self.jwks_url).send().await.map_err(|e| {
194            let error_msg = format!("Failed to fetch JWKS: {e:?}");
195            tracing::error!("{}", error_msg);
196            AuthError::JwksError(error_msg)
197        })?;
198
199        if !response.status().is_success() {
200            let error_msg = format!("JWKS endpoint returned status: {}", response.status());
201            tracing::error!("{}", error_msg);
202            return Err(AuthError::JwksError(error_msg));
203        }
204
205        let jwks: JwksResponse = response.json().await.map_err(|e| {
206            let error_msg = format!("Failed to parse JWKS response: {e:?}");
207            tracing::error!("{}", error_msg);
208            AuthError::JwksError(error_msg)
209        })?;
210
211        // Basic validation: ensure there are keys.
212        if jwks.keys.is_empty() {
213            let error_msg = "JWKS response contains no keys";
214            tracing::error!("{}", error_msg);
215            return Err(AuthError::JwksError(error_msg.to_string()));
216        }
217
218        // Update the cache and timestamps.
219        *self.cache.write().await = Some(jwks.clone());
220        *self.expires_at.write().await = Some(now + JWKS_CACHE_DURATION);
221        *self.cached_at.write().await = Some(now);
222
223        tracing::info!(
224            "JWKS cache updated, expires at: {} (cached at: {})",
225            now + JWKS_CACHE_DURATION,
226            now
227        );
228        Ok(jwks)
229    }
230
231    /// Finds a specific key by its Key ID.
232    ///
233    /// # Arguments
234    ///
235    /// * `kid` - The Key ID
236    pub async fn find_key(&self, kid: &str) -> Result<Jwk, AuthError> {
237        let jwks = self.get_jwks().await?;
238
239        jwks.keys
240            .iter() // Use `iter()` to avoid consuming the `Vec`.
241            .find(|key| key.kid == kid)
242            .cloned() // Clone only the found `Jwk`.
243            .ok_or_else(|| {
244                tracing::warn!("Key with kid '{}' not found in JWKS", kid);
245                AuthError::NoMatchingKey
246            })
247    }
248}
249
250// Tests have been moved to the unified `tests` module.
251// See: `src/tests/jwks_tests.rs`