datse/src/server/web/state.rs

//! 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<Uuid, (auth::AuthRequest, delay_queue::Key)>,
    timeouts: DelayQueue<Uuid>,
}

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<Self::Output> {
	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<auth::AuthRequest, AuthCacheError>
    {
	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<ServerState>,

    auth_tokens: RwLock<AuthContainer>,
    //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
    }
}