yuurei/src/web/route.rs

//! Basic router
use super::*;
use hyper::{
    Method,
};
use std::{
    fmt,
    marker::{
	Send,
	Sync,
    },
    iter,
};
use tokio::{
    sync::{
	mpsc::{
	    self,
	    error::SendTimeoutError,
	},
    },
    time,
};
use futures::{
    future::{
	self,
	Future,
    },
};
use generational_arena::{
    Index,
    Arena,
};

pub trait UriRoute
{
    fn is_match(&self, uri: &str) -> bool;
    #[inline] fn as_string(&self) -> &str
    {
	""
    }

    #[inline] fn type_name(&self) -> &str
    {
	std::any::type_name::<Self>()
    }

    #[inline] fn mutate_uri(&self, uri: String) -> String
    {
	uri
    }
}

impl UriRoute for str
{
    #[inline] fn is_match(&self, uri: &str) -> bool {
	self.eq(uri)
    }
    #[inline] fn as_string(&self) -> &str {
	self
    }
}

impl<T: AsRef<str>> UriRoute for T
{
    #[inline] fn is_match(&self, uri: &str) -> bool {
	self.as_ref().eq(uri)
    }
    #[inline] fn as_string(&self) -> &str {
	self.as_ref()
    }
}

impl UriRoute for regex::Regex
{
    #[inline] fn is_match(&self, uri: &str) -> bool {
	self.test(uri)
    }
    #[inline] fn as_string(&self) -> &str {
	self.as_str()
    }
}

/// A router for all under a prefix
#[derive(Debug, Clone, PartialEq, Hash)]
pub struct PrefixRouter(String);

impl PrefixRouter
{
    /// Create a new instance with this string
    pub fn new(string: impl Into<String>) -> Self
    {
	Self(string.into())
    }
}

impl UriRoute for PrefixRouter
{
    #[inline] fn is_match(&self, uri: &str) -> bool {
	uri.starts_with(self.0.as_str())
    }
    #[inline] fn as_string(&self) -> &str {
	self.0.as_str()
    }

    fn mutate_uri(&self, mut uri: String) -> String {
	uri.replace_range(..self.0.len(), "");
	uri
    }
}

impl fmt::Display for PrefixRouter
{
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result
    {
	write!(f, "{}*", self.0)
    }
}


/// Contains a routing table
#[derive(Debug)]
pub struct Router<T: Send>
{
    routes: Arena<(Option<Method>, OpaqueDebug<Box<dyn UriRoute + Send + Sync + 'static>>, mpsc::Sender<(String, T)>)>,
}

impl<T: Send> fmt::Display for Router<T>
{
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result
    {
	write!(f, "Router {{ routes: ")?;
	for (i, (method, route, _)) in self.routes.iter() {
	    writeln!(f, "\t ({:?} => ({:?}, {} ({:?}))),", i, method, route.type_name(), route.as_string())?;
	}
	write!(f, "}}")
    }
}

impl<T: Send + Clone> Router<T>
{
    /// Create an empty routing table
    pub fn new() -> Self
    {
	Self {
	    routes: Arena::new(),
	}
    }

    /// Push a new route into the router.
    ///
    /// # Returns
    /// The hook's new index, and the receiver that `dispatch()` sends to.
    pub fn hook<Uri: UriRoute + Send + Sync + 'static>(&mut self, method: Option<Method>, uri: Uri) -> (Index, mpsc::Receiver<(String, T)>)
    {
	let (tx, rx) = mpsc::channel(config::get_or_default().dos_max);
	(self.routes.insert((method, OpaqueDebug::new(Box::new(uri)), tx)), rx)
    }

    /// Remove all hooks
    pub fn clear(&mut self)
    {
	self.routes.clear();
    }

    /// Dispatch the URI location across this router, sending to all that match it.
    ///
    /// # Timeout
    /// The timeout is waited on the *individual* dispatches. If you want a global timeout, please timeout the future returned by this function instead.
    /// Timed-out dispatches are counted the same as sending errors.
    ///
    /// # Returns
    /// When one or more dispatchers match but faile, `Err` is returned. Inside the `Err` tuple is the amount of successful dispatches, and also a vector containing the indecies of the failed hook sends.
    pub async fn dispatch(&mut self, method: &Method, uri: impl AsRef<str>, nonce: T, timeout: Option<time::Duration>) -> Result<usize, (usize, Vec<Index>)>
    {
	let string = uri.as_ref();
	let mut success=0usize;
	let vec: Vec<_> =
	    future::join_all(self.routes.iter_mut()
			     .filter_map(|(i, (a_method, route, sender))| {
				 match a_method {
				     Some(x) if x != method => None,
 				     _ => {
					 if route.is_match(string) {
					     trace!("{:?} `{}`: -> {}",i, route.as_string(), string);
					     let timeout = timeout.clone();
					     let nonce= nonce.clone();
					     macro_rules! send {
						 () => {
						     {
							 let string = route.mutate_uri(string.to_owned());
							 match timeout { 
							     None => sender.send((string, nonce)).await
								 .map_err(|e| SendTimeoutError::Closed(e.0)),
							     Some(time) => sender.send_timeout((string, nonce), time).await
							 }
						     }
						 }
					     };
					     Some(async move {
						 match send!() {
						     Err(SendTimeoutError::Closed(er)) => {
							 error!("{:?}: Dispatch failed on hooked route for `{}`", i, er.0);
							 Err(i)
						     },
						     Err(SendTimeoutError::Timeout(er)) => {
							 warn!("{:?}: Dispatch timed out on hooked route for `{}`", i, er.0);
							 Err(i)
						     },
						     _ => Ok(()),
						 }
					     })
					 } else {
					     None
					 }
				     },
				 }
			     })).await.into_iter()
	    .filter_map(|res| {
		if res.is_ok() {
		    success+=1;
		}
		res.err()
	    }).collect();
	if vec.len() > 0 {
	    Err((success, vec))
	} else {
	    Ok(success)
	}
    }

    /// Forcefully dispatch `uri` on hook `which`, regardless of method or URI matching.
    ///
    /// # Returns
    /// If `which` is not contained within the table, immediately returns `None`, otherwise returns a future that completes when the dispatch is complete.
    /// Note: This future must be `await`ed for the dispatch to happen.
    pub fn dispatch_force(&mut self, which: Index, uri: String, nonce: T, timeout: Option<time::Duration>) -> Option<impl Future<Output = Result<(), SendTimeoutError<(String, T)>>> + '_>
    {
	self.routes.get_mut(which).map(move |(_,_,send)| {
	    match timeout {
		Some(timeout) => send.send_timeout((uri, nonce), timeout).boxed(),
		None => send.send((uri, nonce)).map(|res| res.map_err(|e| SendTimeoutError::Closed(e.0))).boxed(),
	    }
	})
    }

    /// Attempt to unhook these hooks. If one or more of the provided indecies does not exist in the routing table, it is ignored.
    pub fn unhook<I>(&mut self, items: I)
    where I: IntoIterator<Item = Index>
    {
	for item in items.into_iter() {
	    self.routes.remove(item);
	}
    }
}