1use std::collections::HashMap;
7use std::sync::{Arc, MutexGuard, RwLockWriteGuard};
8use std::sync::atomic::{AtomicUsize, Ordering};
9
10use chrono::{Duration, Utc};
11
12use crate::{endpoint::PreGrant, code_grant::accesstoken::BearerToken};
13use super::Time;
14use super::grant::Grant;
15use super::generator::{TagGrant, TaggedAssertion, Assertion};
16
17pub trait Issuer {
24 fn issue(&mut self, grant: Grant) -> Result<IssuedToken, ()>;
26
27 fn refresh(&mut self, _refresh: &str, _grant: Grant) -> Result<RefreshedToken, ()>;
29
30 fn recover_token<'a>(&'a self, _: &'a str) -> Result<Option<Grant>, ()>;
32
33 fn recover_refresh<'a>(&'a self, _: &'a str) -> Result<Option<Grant>, ()>;
35}
36
37#[derive(Clone, Debug)]
39pub struct IssuedToken {
40 pub token: String,
42
43 pub refresh: Option<String>,
45
46 pub until: Time,
51
52 pub token_type: TokenType,
54}
55
56#[non_exhaustive]
66#[derive(Clone, Debug)]
67pub enum TokenType {
68 Bearer,
72}
73
74#[derive(Clone, Debug)]
76pub struct RefreshedToken {
77 pub token: String,
79
80 pub refresh: Option<String>,
84
85 pub until: Time,
90
91 pub token_type: TokenType,
93}
94
95pub struct TokenMap<G: TagGrant = Box<dyn TagGrant + Send + Sync + 'static>> {
101 duration: Option<Duration>,
102 generator: G,
103 usage: u64,
104 access: HashMap<Arc<str>, Arc<Token>>,
105 refresh: HashMap<Arc<str>, Arc<Token>>,
106}
107
108struct Token {
109 access: Arc<str>,
111
112 refresh: Option<Arc<str>>,
114
115 grant: Grant,
117}
118
119impl<G: TagGrant> TokenMap<G> {
120 pub fn new(generator: G) -> Self {
122 Self {
123 duration: None,
124 generator,
125 usage: 0,
126 access: HashMap::new(),
127 refresh: HashMap::new(),
128 }
129 }
130
131 pub fn valid_for(&mut self, duration: Duration) {
133 self.duration = Some(duration);
134 }
135
136 pub fn valid_for_default(&mut self) {
138 self.duration = None;
139 }
140
141 pub fn revoke(&mut self, token: &str) {
148 self.access.remove(token);
149 self.refresh.remove(token);
150 }
151
152 pub fn import_grant(&mut self, token: String, mut grant: Grant) {
157 self.set_duration(&mut grant);
158 let key: Arc<str> = Arc::from(token);
159 let token = Token::from_access(key.clone(), grant);
160 self.access.insert(key, Arc::new(token));
161 }
162
163 fn set_duration(&self, grant: &mut Grant) {
164 if let Some(duration) = &self.duration {
165 grant.until = Utc::now() + *duration;
166 }
167 }
168}
169
170impl Token {
171 fn from_access(access: Arc<str>, grant: Grant) -> Self {
172 Token {
173 access,
174 refresh: None,
175 grant,
176 }
177 }
178
179 fn from_refresh(access: Arc<str>, refresh: Arc<str>, grant: Grant) -> Self {
180 Token {
181 access,
182 refresh: Some(refresh),
183 grant,
184 }
185 }
186}
187
188impl IssuedToken {
189 pub fn without_refresh(token: String, until: Time) -> Self {
223 IssuedToken {
224 token,
225 refresh: None,
226 until,
227 token_type: TokenType::Bearer,
228 }
229 }
230
231 pub fn refreshable(&self) -> bool {
235 self.refresh.is_some()
236 }
237
238 pub fn convert_bearer_token(self, pre_grant: PreGrant) -> BearerToken {
240 BearerToken(self, pre_grant.scope)
241 }
242}
243
244impl<G: TagGrant> Issuer for TokenMap<G> {
245 fn issue(&mut self, mut grant: Grant) -> Result<IssuedToken, ()> {
246 self.set_duration(&mut grant);
247 let next_usage = self.usage.wrapping_add(2);
252
253 let (access, refresh) = {
254 let access = self.generator.tag(self.usage, &grant)?;
255 let refresh = self.generator.tag(self.usage.wrapping_add(1), &grant)?;
256 debug_assert!(
257 access.len() > 0,
258 "An empty access token was generated, this is horribly insecure."
259 );
260 debug_assert!(
261 refresh.len() > 0,
262 "An empty refresh token was generated, this is horribly insecure."
263 );
264 (access, refresh)
265 };
266
267 let until = grant.until;
268 let access_key: Arc<str> = Arc::from(access.clone());
269 let refresh_key: Arc<str> = Arc::from(refresh.clone());
270 let token = Token::from_refresh(access_key.clone(), refresh_key.clone(), grant);
271 let token = Arc::new(token);
272
273 self.access.insert(access_key, token.clone());
274 self.refresh.insert(refresh_key, token);
275 self.usage = next_usage;
276 Ok(IssuedToken {
277 token: access,
278 refresh: Some(refresh),
279 until,
280 token_type: TokenType::Bearer,
281 })
282 }
283
284 fn refresh(&mut self, refresh: &str, mut grant: Grant) -> Result<RefreshedToken, ()> {
285 let (refresh_key, mut token) = self
287 .refresh
288 .remove_entry(refresh)
289 .ok_or(())?;
291
292 assert!(Arc::ptr_eq(token.refresh.as_ref().unwrap(), &refresh_key));
293 self.set_duration(&mut grant);
294 let until = grant.until;
295
296 let tag = self.usage;
297 let new_access = self.generator.tag(tag, &grant)?;
298
299 let tag = tag.wrapping_add(1);
300 let new_refresh = self.generator.tag(tag, &grant)?;
301
302 let new_access_key: Arc<str> = Arc::from(new_access.clone());
303 let new_refresh_key: Arc<str> = Arc::from(new_refresh.clone());
304
305 if let Some(atoken) = self.access.remove(&token.access) {
306 assert!(Arc::ptr_eq(&token, &atoken));
307 }
308
309 {
310 let mut_token = Arc::get_mut(&mut token)
312 .unwrap_or_else(|| unreachable!("Grant data was only shared with access and refresh"));
313 mut_token.access = new_access_key.clone();
315 mut_token.refresh = Some(new_refresh_key.clone());
316 mut_token.grant = grant;
317 }
318
319 self.access.insert(new_access_key, token.clone());
320 self.refresh.insert(new_refresh_key, token);
321
322 self.usage = tag.wrapping_add(1);
323 Ok(RefreshedToken {
324 token: new_access,
325 refresh: Some(new_refresh),
326 until,
327 token_type: TokenType::Bearer,
328 })
329 }
330
331 fn recover_token<'a>(&'a self, token: &'a str) -> Result<Option<Grant>, ()> {
332 Ok(self.access.get(token).map(|token| token.grant.clone()))
333 }
334
335 fn recover_refresh<'a>(&'a self, token: &'a str) -> Result<Option<Grant>, ()> {
336 Ok(self.refresh.get(token).map(|token| token.grant.clone()))
337 }
338}
339
340pub struct TokenSigner {
345 duration: Option<Duration>,
346 signer: Assertion,
347 counter: AtomicUsize,
349 have_refresh: bool,
350}
351
352impl TokenSigner {
353 pub fn new(secret: Assertion) -> TokenSigner {
358 TokenSigner {
359 duration: None,
360 signer: secret,
361 counter: AtomicUsize::new(0),
362 have_refresh: false,
363 }
364 }
365
366 pub fn ephemeral() -> TokenSigner {
372 TokenSigner::new(Assertion::ephemeral())
373 }
374
375 pub fn valid_for(&mut self, duration: Duration) {
383 self.duration = Some(duration);
384 }
385
386 pub fn valid_for_default(&mut self) {
391 self.duration = None;
392 }
393
394 pub fn generate_refresh_tokens(&mut self, refresh: bool) {
401 self.have_refresh = refresh;
402 }
403
404 fn next_counter(&self) -> usize {
406 self.counter.fetch_add(1, Ordering::Relaxed)
410 }
411
412 fn refreshable_token(&self, grant: &Grant) -> Result<IssuedToken, ()> {
413 let first_ctr = self.next_counter() as u64;
414 let second_ctr = self.next_counter() as u64;
415
416 let token = self.as_token().sign(first_ctr, grant)?;
417 let refresh = self.as_refresh().sign(second_ctr, grant)?;
418
419 Ok(IssuedToken {
420 token,
421 refresh: Some(refresh),
422 until: grant.until,
423 token_type: TokenType::Bearer,
424 })
425 }
426
427 fn unrefreshable_token(&self, grant: &Grant) -> Result<IssuedToken, ()> {
428 let counter = self.next_counter() as u64;
429
430 let token = self.as_token().sign(counter, grant)?;
431
432 Ok(IssuedToken::without_refresh(token, grant.until))
433 }
434
435 fn as_token(&self) -> TaggedAssertion {
436 self.signer.tag("token")
437 }
438
439 fn as_refresh(&self) -> TaggedAssertion {
440 self.signer.tag("refresh")
441 }
442}
443
444impl<'s, I: Issuer + ?Sized> Issuer for &'s mut I {
445 fn issue(&mut self, grant: Grant) -> Result<IssuedToken, ()> {
446 (**self).issue(grant)
447 }
448
449 fn refresh(&mut self, token: &str, grant: Grant) -> Result<RefreshedToken, ()> {
450 (**self).refresh(token, grant)
451 }
452
453 fn recover_token<'a>(&'a self, token: &'a str) -> Result<Option<Grant>, ()> {
454 (**self).recover_token(token)
455 }
456
457 fn recover_refresh<'a>(&'a self, token: &'a str) -> Result<Option<Grant>, ()> {
458 (**self).recover_refresh(token)
459 }
460}
461
462impl<I: Issuer + ?Sized> Issuer for Box<I> {
463 fn issue(&mut self, grant: Grant) -> Result<IssuedToken, ()> {
464 (**self).issue(grant)
465 }
466
467 fn refresh(&mut self, token: &str, grant: Grant) -> Result<RefreshedToken, ()> {
468 (**self).refresh(token, grant)
469 }
470
471 fn recover_token<'a>(&'a self, token: &'a str) -> Result<Option<Grant>, ()> {
472 (**self).recover_token(token)
473 }
474
475 fn recover_refresh<'a>(&'a self, token: &'a str) -> Result<Option<Grant>, ()> {
476 (**self).recover_refresh(token)
477 }
478}
479
480impl<'s, I: Issuer + ?Sized> Issuer for MutexGuard<'s, I> {
481 fn issue(&mut self, grant: Grant) -> Result<IssuedToken, ()> {
482 (**self).issue(grant)
483 }
484
485 fn refresh(&mut self, token: &str, grant: Grant) -> Result<RefreshedToken, ()> {
486 (**self).refresh(token, grant)
487 }
488
489 fn recover_token<'a>(&'a self, token: &'a str) -> Result<Option<Grant>, ()> {
490 (**self).recover_token(token)
491 }
492
493 fn recover_refresh<'a>(&'a self, token: &'a str) -> Result<Option<Grant>, ()> {
494 (**self).recover_refresh(token)
495 }
496}
497
498impl<'s, I: Issuer + ?Sized> Issuer for RwLockWriteGuard<'s, I> {
499 fn issue(&mut self, grant: Grant) -> Result<IssuedToken, ()> {
500 (**self).issue(grant)
501 }
502
503 fn refresh(&mut self, token: &str, grant: Grant) -> Result<RefreshedToken, ()> {
504 (**self).refresh(token, grant)
505 }
506
507 fn recover_token<'a>(&'a self, token: &'a str) -> Result<Option<Grant>, ()> {
508 (**self).recover_token(token)
509 }
510
511 fn recover_refresh<'a>(&'a self, token: &'a str) -> Result<Option<Grant>, ()> {
512 (**self).recover_refresh(token)
513 }
514}
515
516impl Issuer for TokenSigner {
517 fn issue(&mut self, grant: Grant) -> Result<IssuedToken, ()> {
518 (&mut &*self).issue(grant)
519 }
520
521 fn refresh(&mut self, _refresh: &str, _grant: Grant) -> Result<RefreshedToken, ()> {
522 Err(())
523 }
524
525 fn recover_token<'a>(&'a self, token: &'a str) -> Result<Option<Grant>, ()> {
526 (&&*self).recover_token(token)
527 }
528
529 fn recover_refresh<'a>(&'a self, token: &'a str) -> Result<Option<Grant>, ()> {
530 (&&*self).recover_refresh(token)
531 }
532}
533
534impl<'a> Issuer for &'a TokenSigner {
535 fn issue(&mut self, mut grant: Grant) -> Result<IssuedToken, ()> {
536 if let Some(duration) = &self.duration {
537 grant.until = Utc::now() + *duration;
538 }
539
540 if self.have_refresh {
541 self.refreshable_token(&grant)
542 } else {
543 self.unrefreshable_token(&grant)
544 }
545 }
546
547 fn refresh(&mut self, _refresh: &str, _grant: Grant) -> Result<RefreshedToken, ()> {
548 Err(())
549 }
550
551 fn recover_token<'t>(&'t self, token: &'t str) -> Result<Option<Grant>, ()> {
552 Ok(self.as_token().extract(token).ok())
553 }
554
555 fn recover_refresh<'t>(&'t self, token: &'t str) -> Result<Option<Grant>, ()> {
556 if !self.have_refresh {
557 return Ok(None);
558 }
559
560 Ok(self.as_refresh().extract(token).ok())
561 }
562}
563
564#[cfg(test)]
565pub mod tests {
567 use super::*;
568 use crate::primitives::grant::Extensions;
569 use crate::primitives::generator::RandomGenerator;
570 use chrono::{Duration, Utc};
571
572 fn grant_template() -> Grant {
573 Grant {
574 client_id: "Client".to_string(),
575 owner_id: "Owner".to_string(),
576 redirect_uri: "https://example.com".parse().unwrap(),
577 scope: "default".parse().unwrap(),
578 until: Utc::now() + Duration::hours(1),
579 extensions: Extensions::new(),
580 }
581 }
582
583 pub fn simple_test_suite(issuer: &mut dyn Issuer) {
591 let request = grant_template();
592
593 let issued = issuer.issue(request.clone()).expect("Issuing failed");
594 let from_token = issuer
595 .recover_token(&issued.token)
596 .expect("Issuer failed during recover")
597 .expect("Issued token appears to be invalid");
598
599 assert_ne!(Some(&issued.token), issued.refresh.as_ref());
600 assert_eq!(from_token.client_id, "Client");
601 assert_eq!(from_token.owner_id, "Owner");
602 assert!(Utc::now() < from_token.until);
603
604 let issued_2 = issuer.issue(request).expect("Issuing failed");
605 assert_ne!(issued.token, issued_2.token);
606 assert_ne!(Some(&issued.token), issued_2.refresh.as_ref());
607 assert_ne!(issued.refresh, issued_2.refresh);
608 assert_ne!(issued.refresh.as_ref(), Some(&issued_2.token));
609 }
610
611 #[test]
612 fn signer_test_suite() {
613 let mut signer = TokenSigner::ephemeral();
614 signer.generate_refresh_tokens(true);
617 simple_test_suite(&mut signer);
618 }
619
620 #[test]
621 fn signer_no_default_refresh() {
622 let mut signer = TokenSigner::ephemeral();
623 let issued = signer.issue(grant_template());
624
625 let token = issued.expect("Issuing without refresh token failed");
626 assert!(!token.refreshable());
627 }
628
629 #[test]
630 fn random_test_suite() {
631 let mut token_map = TokenMap::new(RandomGenerator::new(16));
632 simple_test_suite(&mut token_map);
633 }
634
635 #[test]
636 fn random_has_refresh() {
637 let mut token_map = TokenMap::new(RandomGenerator::new(16));
638 let issued = token_map.issue(grant_template());
639
640 let token = issued.expect("Issuing without refresh token failed");
641 assert!(token.refreshable());
642 }
643
644 #[test]
645 fn random_refresh_rotation() {
646 let mut token_map = TokenMap::new(RandomGenerator::new(16));
647 let issued = token_map.issue(grant_template());
648
649 let token = issued.expect("Issuing with refresh token failed");
650 let refresh = token.refresh.expect("No refresh token returned");
651
652 let refreshed_token = token_map
653 .refresh(&refresh, grant_template())
654 .expect("Failed to refresh access token");
655
656 let new_refresh = refreshed_token.refresh.expect("No new refresh token returned");
657
658 assert!(refresh != new_refresh);
659 }
660
661 #[test]
662 #[should_panic]
663 fn bad_generator() {
664 struct BadGenerator;
665 impl TagGrant for BadGenerator {
666 fn tag(&mut self, _: u64, _: &Grant) -> Result<String, ()> {
667 Ok("YOLO.HowBadCanItBeToRepeatTokens?".into())
668 }
669 }
670 let mut token_map = TokenMap::new(BadGenerator);
671 simple_test_suite(&mut token_map);
672 }
673}