//! Web server state use super::*; use std::{ collections::HashMap, sync::Arc, pin::Pin, task::Context, task::Poll, fmt,error, }; use tokio::{ sync::{ RwLock, RwLockWriteGuard, RwLockReadGuard, }, time::{ self, DelayQueue, delay_queue, }, }; #[derive(Debug)] pub struct AuthContainer { active_req: HashMap, timeouts: DelayQueue, } pub struct AuthPurge<'a, F = fn(auth::AuthRequest)>(&'a mut AuthContainer, F); impl<'a, F> Future for AuthPurge<'a, F> where F: FnMut(auth::AuthRequest) + 'a + Unpin { type Output = Result<(), time::Error>; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let this = self.get_mut(); while let Some(res) = futures::ready!(this.0.timeouts.poll_expired(cx)) { let ent = res?; let _ = this.0.active_req.remove(ent.get_ref()).map(|x| x.0).map(&mut this.1); } Poll::Ready(Ok(())) } } impl AuthContainer { /// Gerate a new empty auth token container fn new() -> Self { Self { active_req: HashMap::new(), timeouts: DelayQueue::new(), } } /// Returns a future that purges expired entries, running the provided closure on them. /// /// The future will yield if: /// * The stire is not empty /// * There are non-expired entries in the store pub fn purge_and_then<'a, F: FnMut(auth::AuthRequest) +Unpin +'a>(&'a mut self, and_then: F) -> AuthPurge<'a, F> { AuthPurge(self, and_then) } /// Returns a future that purges expired entries. See `purge_and_then`. #[inline] pub fn purge(&mut self) -> AuthPurge<'_> { AuthPurge(self, std::mem::drop) } /// Purge all expired entries. #[inline] pub fn purge_now(&mut self) { self.purge().now_or_never(); } /// Purge all expired entries, running the provided closure on them. pub fn purge_now_and_then<'a, F: FnMut(auth::AuthRequest) +Unpin+'a>(&'a mut self, and_then: F) { self.purge_and_then(and_then).now_or_never(); } /// Insert a request into the store, setting it to expire once its ttl is up. pub fn insert_req(&mut self, req: auth::AuthRequest) { self.purge_now(); let k = self.timeouts.insert(req.id, req.ttl()); self.active_req.insert(req.id, (req, k)); } /// Attempt to retrieve a value from the store by its ID. /// /// # Notes /// `AuthCacheError::Timeout` will only be returned if the request we're trying to extract has timed out *and not yet been removed* yet by an earlier, potentially unrelated, call to `handle_req` *or* `insert_req` (or an explicit purge). /// If an error is returned you cannot rely on the accuracy of the error kind. pub fn handle_req(&mut self, id: Uuid) -> Result { let mut timed_out=false; self.purge_now_and_then(|other| if other.id==id { timed_out = true; }); if timed_out { Err(AuthCacheError::Timeout) } else { self.active_req.remove(&id).ok_or(AuthCacheError::Removed).map(|(v, k)| { self.timeouts.remove(&k); v }) } } } #[derive(Debug)] #[non_exhaustive] pub enum AuthCacheError { Removed, Timeout, } impl error::Error for AuthCacheError{} impl fmt::Display for AuthCacheError { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self { Self::Removed => write!(f, "id was not present"), Self::Timeout => write!(f, "id timed out as we read it"), } } } #[derive(Debug)] pub struct State { backend: RwLock, auth_tokens: RwLock, //TODO: user auths, public keys, hashed passwords, etc. settings: Settings, } impl State { /// Create state from a backend state and server settings pub fn new(backend: ServerState, settings: Settings) -> Self { Self { auth_tokens: RwLock::new(AuthContainer::new()), backend: RwLock::new(backend), settings, } } /// The web server settings pub fn cfg(&self) -> &Settings { &self.settings } /// Get a write reference to the auth container pub async fn auth_tokens(&self) -> RwLockWriteGuard<'_, AuthContainer> { self.auth_tokens.write().await } /// Get a read reference to the auth container. /// /// Typically only useful for debugging/logging. pub async fn auth_tokens_ref(&self) -> RwLockReadGuard<'_, AuthContainer> { self.auth_tokens.read().await } }