yuurei/src/web/route.rs

//! Basic router
use super::*;
use hyper::{
    Method,
};
use std::{
    fmt,
    marker::Send,
    iter,
};
use tokio::{
    sync::{
	mpsc,
	broadcast,
	notify,
    },
};
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
    {
	""
    }
}

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()
    }
}

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

impl Router
{
    /// 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 + 'static>(&mut self, method: Option<Method>, uri: Uri) -> (Index, mpsc::Receiver<String>)
    {
	let (tx, rx) = mpsc::channel(config::get_or_default().dos_max);
	(self.routes.insert((method, OpaqueDebug::new(Box::new(uri)), tx)), rx)
    }

    /// Dispatch the URI location across this router, sending to all that match it.
    ///
    /// # 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>, timeout: impl Future) -> Result<usize, (usize, Vec<Index>)>
    {
	let string = uri.as_ref();
	let (tcx, trx) = broadcast::channel(1);
	tokio::pin!(timeout);
	let begin_to = tokio::sync::Barrier::new(self.routes.len()+1);
	let timeout = async {
	    timeout.await;
	    begin_to.wait().await;
	    tcx.send(());
	};
	let mut timeouts = iter::once(trx)
	    .chain(iter::repeat_with(|| tcx.subscribe()));
	let output = async {
	    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 mut timeout = timeouts.next().unwrap();
						 Some(async move {
						     match tokio::select!{
							 _ = timeout.recv() => {
							     None
							 }
							 s = sender.send(string.to_owned()) => {
							     Some(s)
							 }
						     } {
							 Some(Err(er)) => {
							     warn!("{:?}: Dispatch failed on hooked route for {}", i, er.0);
							     Err(i)
							 },
							 Some(_) => Ok(()),
							 None => {
							     warn!("{:?}: Dispatch timed out on hooked route", i);
							     Err(i)
							 },
						     }
						 })
					     } else {
						 None
					     }
					 },
				     }
				 })).await.into_iter()
		.filter_map(|res| {
		    if res.is_ok() {
			success+=1;
		    }
		    res.err()
		}).collect();
	    (success, vec)
	};
	tokio::pin!(output);
	let (_, (success, vec)) = future::join(timeout, output).await;
	if vec.len() > 0 {
	    Err((success, vec))
	} else {
	    Ok(success)
	}
    }

    /// 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);
	}
    }
}