Skip to main content

arete_auth/
verifier.rs

1use crate::claims::AuthContext;
2use crate::error::VerifyError;
3use crate::keys::VerifyingKey;
4use crate::sealed::ClaimSealKeys;
5use crate::token::{JwksVerifier, TokenVerifier};
6use std::sync::Arc;
7use std::time::{Duration, Instant};
8use tokio::sync::RwLock;
9
10/// Cached JWKS with expiration
11#[derive(Clone)]
12struct CachedJwks {
13    verifier: JwksVerifier,
14    fetched_at: Instant,
15}
16
17/// Async verifier with JWKS caching support
18pub struct AsyncVerifier {
19    inner: VerifierInner,
20    jwks_url: Option<String>,
21    cache_duration: Duration,
22    cached_jwks: Arc<RwLock<Option<CachedJwks>>>,
23    /// Issuer for JWKS-based verification
24    issuer: String,
25    /// Audience for JWKS-based verification
26    audiences: crate::AudienceSet,
27    require_origin: bool,
28    claim_seal_keys: Option<ClaimSealKeys>,
29}
30
31enum VerifierInner {
32    Static(TokenVerifier),
33    Jwks(JwksVerifier),
34}
35
36impl AsyncVerifier {
37    /// Create a verifier with a static key
38    pub fn with_static_key(
39        key: VerifyingKey,
40        issuer: impl Into<String>,
41        audience: impl Into<String>,
42    ) -> Self {
43        let issuer_str = issuer.into();
44        let audience_str = audience.into();
45        Self {
46            inner: VerifierInner::Static(TokenVerifier::new(
47                key,
48                issuer_str.clone(),
49                audience_str.clone(),
50            )),
51            jwks_url: None,
52            cache_duration: Duration::from_secs(3600), // 1 hour default
53            cached_jwks: Arc::new(RwLock::new(None)),
54            issuer: issuer_str,
55            audiences: crate::AudienceSet::single(audience_str),
56            require_origin: false,
57            claim_seal_keys: None,
58        }
59    }
60
61    /// Create a verifier with JWKS
62    pub fn with_jwks(
63        jwks: crate::token::Jwks,
64        issuer: impl Into<String>,
65        audience: impl Into<String>,
66    ) -> Self {
67        let issuer_str = issuer.into();
68        let audience_str = audience.into();
69        Self {
70            inner: VerifierInner::Jwks(JwksVerifier::new(
71                jwks,
72                issuer_str.clone(),
73                audience_str.clone(),
74            )),
75            jwks_url: None,
76            cache_duration: Duration::from_secs(3600),
77            cached_jwks: Arc::new(RwLock::new(None)),
78            issuer: issuer_str,
79            audiences: crate::AudienceSet::single(audience_str),
80            require_origin: false,
81            claim_seal_keys: None,
82        }
83    }
84
85    /// Create a verifier that fetches JWKS from a URL
86    #[cfg(feature = "jwks")]
87    pub fn with_jwks_url(
88        url: impl Into<String>,
89        issuer: impl Into<String>,
90        audience: impl Into<String>,
91    ) -> Self {
92        let issuer_str = issuer.into();
93        let audience_str = audience.into();
94        Self {
95            inner: VerifierInner::Static(TokenVerifier::new(
96                VerifyingKey::from_bytes(&[0u8; 32]).expect("zero key should be valid"),
97                issuer_str.clone(),
98                audience_str.clone(),
99            )),
100            jwks_url: Some(url.into()),
101            issuer: issuer_str,
102            audiences: crate::AudienceSet::single(audience_str),
103            cache_duration: Duration::from_secs(3600),
104            cached_jwks: Arc::new(RwLock::new(None)),
105            require_origin: false,
106            claim_seal_keys: None,
107        }
108    }
109
110    /// Require origin validation on verified tokens.
111    pub fn with_origin_validation(mut self) -> Self {
112        self.require_origin = true;
113        self.inner = match self.inner {
114            VerifierInner::Static(verifier) => {
115                VerifierInner::Static(verifier.with_origin_validation())
116            }
117            VerifierInner::Jwks(verifier) => VerifierInner::Jwks(verifier.with_origin_validation()),
118        };
119        self.cached_jwks = Arc::new(RwLock::new(None));
120        self
121    }
122
123    /// Open sealed claims with these keys. Without them, a token that carries
124    /// sealed claims is refused.
125    pub fn with_claim_seal_keys(mut self, keys: ClaimSealKeys) -> Self {
126        self.inner = match self.inner {
127            VerifierInner::Static(verifier) => {
128                VerifierInner::Static(verifier.with_claim_seal_keys(keys.clone()))
129            }
130            VerifierInner::Jwks(verifier) => {
131                VerifierInner::Jwks(verifier.with_claim_seal_keys(keys.clone()))
132            }
133        };
134        self.claim_seal_keys = Some(keys);
135        // A verifier cached before this call would still open with the old keys.
136        self.cached_jwks = Arc::new(RwLock::new(None));
137        self
138    }
139
140    /// Set cache duration for JWKS
141    pub fn with_cache_duration(mut self, duration: Duration) -> Self {
142        self.cache_duration = duration;
143        self
144    }
145
146    /// Verify a token with automatic JWKS fetching and caching
147    #[cfg(feature = "jwks")]
148    pub async fn verify(
149        &self,
150        token: &str,
151        expected_origin: Option<&str>,
152        expected_client_ip: Option<&str>,
153    ) -> Result<AuthContext, VerifyError> {
154        // If using static JWKS or static key, use directly
155        match &self.inner {
156            VerifierInner::Static(verifier) => {
157                verifier.verify(token, expected_origin, expected_client_ip)
158            }
159            VerifierInner::Jwks(verifier) => {
160                verifier.verify(token, expected_origin, expected_client_ip)
161            }
162        }
163    }
164
165    /// Verify a token (non-JWKS version)
166    #[cfg(not(feature = "jwks"))]
167    pub fn verify(
168        &self,
169        token: &str,
170        expected_origin: Option<&str>,
171        expected_client_ip: Option<&str>,
172    ) -> Result<AuthContext, VerifyError> {
173        match &self.inner {
174            VerifierInner::Static(verifier) => {
175                verifier.verify(token, expected_origin, expected_client_ip)
176            }
177            VerifierInner::Jwks(verifier) => {
178                verifier.verify(token, expected_origin, expected_client_ip)
179            }
180        }
181    }
182
183    /// Refresh JWKS cache from the configured URL
184    #[cfg(feature = "jwks")]
185    pub async fn refresh_cache(&self) -> Result<(), VerifyError> {
186        if let Some(ref jwks_url) = self.jwks_url {
187            // Fetch JWKS from URL
188            let jwks = crate::token::JwksVerifier::fetch_jwks(jwks_url)
189                .await
190                .map_err(|e| VerifyError::InvalidFormat(format!("Failed to fetch JWKS: {}", e)))?;
191
192            // Create new verifier with fetched JWKS
193            let verifier =
194                JwksVerifier::with_audience_set(jwks, &self.issuer, self.audiences.clone());
195            let verifier = if self.require_origin {
196                verifier.with_origin_validation()
197            } else {
198                verifier
199            };
200            let verifier = match &self.claim_seal_keys {
201                Some(keys) => verifier.with_claim_seal_keys(keys.clone()),
202                None => verifier,
203            };
204
205            // Update cache
206            let mut cached = self.cached_jwks.write().await;
207            *cached = Some(CachedJwks {
208                verifier,
209                fetched_at: Instant::now(),
210            });
211        }
212        Ok(())
213    }
214
215    /// Get cached verifier if available and not expired
216    async fn get_cached_verifier(&self) -> Option<JwksVerifier> {
217        let cached = self.cached_jwks.read().await;
218        if let Some(ref cached_jwks) = *cached {
219            if cached_jwks.fetched_at.elapsed() < self.cache_duration {
220                return Some(cached_jwks.verifier.clone());
221            }
222        }
223        None
224    }
225
226    /// Verify a token with automatic JWKS caching
227    #[cfg(feature = "jwks")]
228    pub async fn verify_with_cache(
229        &self,
230        token: &str,
231        expected_origin: Option<&str>,
232        expected_client_ip: Option<&str>,
233    ) -> Result<AuthContext, VerifyError> {
234        // Try cached verifier first
235        if let Some(verifier) = self.get_cached_verifier().await {
236            match verifier.verify(token, expected_origin, expected_client_ip) {
237                Ok(ctx) => return Ok(ctx),
238                Err(VerifyError::KeyNotFound(_)) => {
239                    // Key not found in cache, refresh and retry
240                }
241                Err(e) => return Err(e),
242            }
243        }
244
245        // Refresh cache and try again
246        self.refresh_cache().await?;
247
248        if let Some(verifier) = self.get_cached_verifier().await {
249            verifier.verify(token, expected_origin, expected_client_ip)
250        } else if self.jwks_url.is_some() {
251            Err(VerifyError::InvalidFormat(
252                "JWKS cache unavailable after refresh".to_string(),
253            ))
254        } else {
255            // Fallback to inner verifier if no cache available
256            match &self.inner {
257                VerifierInner::Static(verifier) => {
258                    verifier.verify(token, expected_origin, expected_client_ip)
259                }
260                VerifierInner::Jwks(verifier) => {
261                    verifier.verify(token, expected_origin, expected_client_ip)
262                }
263            }
264        }
265    }
266}
267
268/// Simple synchronous verifier for use in non-async contexts
269pub struct SimpleVerifier {
270    inner: TokenVerifier,
271}
272
273impl SimpleVerifier {
274    /// Create a new simple verifier
275    pub fn new(key: VerifyingKey, issuer: impl Into<String>, audience: impl Into<String>) -> Self {
276        Self {
277            inner: TokenVerifier::new(key, issuer, audience),
278        }
279    }
280
281    /// Verify a token synchronously
282    pub fn verify(
283        &self,
284        token: &str,
285        expected_origin: Option<&str>,
286        expected_client_ip: Option<&str>,
287    ) -> Result<AuthContext, VerifyError> {
288        self.inner
289            .verify(token, expected_origin, expected_client_ip)
290    }
291}
292
293#[cfg(test)]
294mod tests {
295    use super::*;
296    use crate::claims::{KeyClass, SessionClaims};
297    use crate::keys::SigningKey;
298    use crate::token::TokenSigner;
299    use base64::Engine;
300
301    #[cfg(feature = "jwks")]
302    use tokio::io::{AsyncReadExt, AsyncWriteExt};
303
304    #[tokio::test]
305    async fn test_async_verifier_with_static_key() {
306        let signing_key = SigningKey::generate();
307        let verifying_key = signing_key.verifying_key();
308
309        let signer = TokenSigner::new(signing_key, "test-issuer");
310        let verifier =
311            AsyncVerifier::with_static_key(verifying_key, "test-issuer", "test-audience");
312
313        let claims = SessionClaims::builder("test-issuer", "test-subject", "test-audience")
314            .with_scope("read")
315            .with_metering_key("meter-123")
316            .with_key_class(KeyClass::Publishable)
317            .build();
318
319        let token = signer.sign(claims).unwrap();
320        let context = verifier.verify(&token, None, None).await.unwrap();
321
322        assert_eq!(context.subject, "test-subject");
323    }
324
325    #[test]
326    fn test_simple_verifier() {
327        let signing_key = SigningKey::generate();
328        let verifying_key = signing_key.verifying_key();
329
330        let signer = TokenSigner::new(signing_key, "test-issuer");
331        let verifier = SimpleVerifier::new(verifying_key, "test-issuer", "test-audience");
332
333        let claims = SessionClaims::builder("test-issuer", "test-subject", "test-audience")
334            .with_scope("read")
335            .with_metering_key("meter-123")
336            .with_key_class(KeyClass::Publishable)
337            .build();
338
339        let token = signer.sign(claims).unwrap();
340        let context = verifier.verify(&token, None, None).unwrap();
341
342        assert_eq!(context.subject, "test-subject");
343        assert_eq!(context.metering_key, "meter-123");
344    }
345
346    #[cfg(feature = "jwks")]
347    #[test]
348    fn test_verify_with_cache_returns_explicit_error_when_cache_stays_empty() {
349        tokio::runtime::Runtime::new().unwrap().block_on(async {
350            let signing_key = SigningKey::generate();
351            let verifying_key = signing_key.verifying_key();
352            let signer = TokenSigner::new(signing_key, "test-issuer");
353
354            let jwks = serde_json::json!({
355                "keys": [{
356                    "kty": "OKP",
357                    "use": "sig",
358                    "kid": verifying_key.key_id(),
359                    "x": base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(verifying_key.to_bytes()),
360                }]
361            })
362            .to_string();
363
364            let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
365            let addr = listener.local_addr().unwrap();
366            let response_body = jwks.clone();
367            tokio::spawn(async move {
368                let (mut socket, _) = listener.accept().await.unwrap();
369                let mut buffer = [0u8; 1024];
370                let _ = socket.read(&mut buffer).await;
371
372                let response = format!(
373                    "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
374                    response_body.len(),
375                    response_body
376                );
377                socket.write_all(response.as_bytes()).await.unwrap();
378            });
379
380            let verifier = AsyncVerifier::with_jwks_url(
381                format!("http://{addr}/jwks"),
382                "test-issuer",
383                "test-audience",
384            )
385            .with_cache_duration(Duration::ZERO);
386
387            let claims = SessionClaims::builder("test-issuer", "test-subject", "test-audience")
388                .with_scope("read")
389                .with_metering_key("meter-123")
390                .with_key_class(KeyClass::Publishable)
391                .build();
392            let token = signer.sign(claims).unwrap();
393
394            let result = verifier.verify_with_cache(&token, None, None).await;
395            assert!(matches!(
396                result,
397                Err(VerifyError::InvalidFormat(ref msg)) if msg == "JWKS cache unavailable after refresh"
398            ));
399        });
400    }
401}