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#[derive(Clone)]
12struct CachedJwks {
13 verifier: JwksVerifier,
14 fetched_at: Instant,
15}
16
17pub struct AsyncVerifier {
19 inner: VerifierInner,
20 jwks_url: Option<String>,
21 cache_duration: Duration,
22 cached_jwks: Arc<RwLock<Option<CachedJwks>>>,
23 issuer: String,
25 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 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), 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 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 #[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 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 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 self.cached_jwks = Arc::new(RwLock::new(None));
137 self
138 }
139
140 pub fn with_cache_duration(mut self, duration: Duration) -> Self {
142 self.cache_duration = duration;
143 self
144 }
145
146 #[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 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 #[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 #[cfg(feature = "jwks")]
185 pub async fn refresh_cache(&self) -> Result<(), VerifyError> {
186 if let Some(ref jwks_url) = self.jwks_url {
187 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 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 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 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 #[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 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 }
241 Err(e) => return Err(e),
242 }
243 }
244
245 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 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
268pub struct SimpleVerifier {
270 inner: TokenVerifier,
271}
272
273impl SimpleVerifier {
274 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 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}