use std::sync::Arc;
use rainier_container::Container;
use rainier_support::{Error, Result};
use crate::pipeline::Middleware;
type Builder = Arc<dyn Fn(&Container) -> Result<Arc<dyn Middleware>> + Send + Sync>;
enum Stage {
Ready(Arc<dyn Middleware>),
Deferred(&'static str, Builder),
}
impl Clone for Stage {
fn clone(&self) -> Self {
match self {
Stage::Ready(middleware) => Stage::Ready(Arc::clone(middleware)),
Stage::Deferred(name, build) => Stage::Deferred(name, Arc::clone(build)),
}
}
}
#[derive(Default, Clone)]
pub struct MiddlewareStack {
stages: Vec<Stage>,
}
impl MiddlewareStack {
pub fn new() -> Self {
Self::default()
}
pub fn with(mut self, middleware: impl Middleware) -> Self {
self.stages.push(Stage::Ready(Arc::new(middleware)));
self
}
pub fn with_arc(mut self, middleware: Arc<dyn Middleware>) -> Self {
self.stages.push(Stage::Ready(middleware));
self
}
pub fn with_stack(mut self, other: MiddlewareStack) -> Self {
self.stages.extend(other.stages);
self
}
pub fn deferred<M, F>(mut self, build: F) -> Self
where
M: Middleware,
F: Fn(&Container) -> Result<M> + Send + Sync + 'static,
{
let name = crate::pipeline::short_type_name(std::any::type_name::<M>());
self.stages.push(Stage::Deferred(
name,
Arc::new(move |container| build(container).map(|m| Arc::new(m) as Arc<dyn Middleware>)),
));
self
}
pub fn resolved<S, M, F>(self, build: F) -> Self
where
S: Send + Sync + 'static,
M: Middleware,
F: Fn(Arc<S>) -> M + Send + Sync + 'static,
{
self.deferred(move |container| Ok(build(container.resolve::<S>()?)))
}
pub fn prepend(&mut self, outer: &MiddlewareStack) {
let mut combined = outer.stages.clone();
combined.append(&mut self.stages);
self.stages = combined;
}
pub fn len(&self) -> usize {
self.stages.len()
}
pub fn is_empty(&self) -> bool {
self.stages.is_empty()
}
pub fn resolve(&self, container: &Container) -> Result<Vec<Arc<dyn Middleware>>> {
self.stages
.iter()
.map(|stage| match stage {
Stage::Ready(middleware) => Ok(Arc::clone(middleware)),
Stage::Deferred(name, build) => build(container).map_err(|e| {
Error::internal(format!("could not build the `{name}` middleware: {e}"))
}),
})
.collect()
}
pub fn labels(&self) -> Vec<&'static str> {
self.stages
.iter()
.map(|stage| match stage {
Stage::Ready(middleware) => middleware.name(),
Stage::Deferred(name, _) => name,
})
.collect()
}
}
impl std::fmt::Debug for MiddlewareStack {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_list().entries(self.labels()).finish()
}
}
pub trait IntoMiddlewareStack {
fn into_middleware_stack(self) -> MiddlewareStack;
}
impl<M: Middleware> IntoMiddlewareStack for M {
fn into_middleware_stack(self) -> MiddlewareStack {
MiddlewareStack::new().with(self)
}
}
impl IntoMiddlewareStack for MiddlewareStack {
fn into_middleware_stack(self) -> MiddlewareStack {
self
}
}
impl IntoMiddlewareStack for Arc<dyn Middleware> {
fn into_middleware_stack(self) -> MiddlewareStack {
MiddlewareStack::new().with_arc(self)
}
}
impl IntoMiddlewareStack for Vec<Arc<dyn Middleware>> {
fn into_middleware_stack(self) -> MiddlewareStack {
self.into_iter().fold(MiddlewareStack::new(), MiddlewareStack::with_arc)
}
}
impl<M: Middleware> From<M> for MiddlewareStack {
fn from(middleware: M) -> Self {
MiddlewareStack::new().with(middleware)
}
}
macro_rules! tuple_stacks {
($( ($($name:ident),+) ),+ $(,)?) => {
$(
#[allow(non_snake_case, reason = "the idiomatic spelling for a tuple impl")]
impl<$($name: Middleware),+> IntoMiddlewareStack for ($($name,)+) {
fn into_middleware_stack(self) -> MiddlewareStack {
let ($($name,)+) = self;
MiddlewareStack::new()$(.with($name))+
}
}
)+
};
}
tuple_stacks! {
(A, B),
(A, B, C),
(A, B, C, D),
(A, B, C, D, E),
(A, B, C, D, E, F),
(A, B, C, D, E, F, G),
(A, B, C, D, E, F, G, H),
}
#[cfg(test)]
mod tests {
use super::*;
use crate::pipeline::Next;
use rainier_http::{Request, Response};
macro_rules! tag {
($name:ident) => {
struct $name;
#[async_trait::async_trait]
impl Middleware for $name {
async fn handle(&self, request: Request, next: Next) -> Response {
next.run(request).await
}
fn name(&self) -> &'static str {
stringify!($name)
}
}
};
}
tag!(First);
tag!(Second);
tag!(Third);
fn resolved_names(stack: &MiddlewareStack) -> Vec<&'static str> {
stack.resolve(&Container::new()).unwrap().iter().map(|m| m.name()).collect()
}
#[test]
fn a_stack_keeps_the_order_it_was_built_in() {
let stack = MiddlewareStack::new().with(First).with(Second).with(Third);
assert_eq!(resolved_names(&stack), vec!["First", "Second", "Third"]);
}
#[test]
fn one_middleware_converts_to_a_stack_of_one() {
assert_eq!(resolved_names(&First.into_middleware_stack()), vec!["First"]);
}
#[test]
fn a_tuple_converts_in_order() {
let stack = (First, Second, Third).into_middleware_stack();
assert_eq!(resolved_names(&stack), vec!["First", "Second", "Third"]);
}
#[test]
fn a_stack_converts_to_itself() {
let stack = MiddlewareStack::new().with(First);
assert_eq!(resolved_names(&stack.into_middleware_stack()), vec!["First"]);
}
#[test]
fn nesting_a_stack_flattens_it_in_place() {
let inner = MiddlewareStack::new().with(Second).with(Third);
let outer = MiddlewareStack::new().with(First).with_stack(inner);
assert_eq!(resolved_names(&outer), vec!["First", "Second", "Third"]);
}
#[test]
fn prepending_puts_the_group_outside_the_routes_own() {
let mut route = MiddlewareStack::new().with(Third);
route.prepend(&MiddlewareStack::new().with(First).with(Second));
assert_eq!(resolved_names(&route), vec!["First", "Second", "Third"]);
}
#[test]
fn a_deferred_stage_is_built_at_resolution_not_declaration() {
use std::sync::atomic::{AtomicUsize, Ordering};
static BUILT: AtomicUsize = AtomicUsize::new(0);
let stack = MiddlewareStack::new().deferred(|_| {
BUILT.fetch_add(1, Ordering::SeqCst);
Ok(First)
});
assert_eq!(BUILT.load(Ordering::SeqCst), 0, "declaring must not build");
assert_eq!(resolved_names(&stack), vec!["First"]);
assert_eq!(BUILT.load(Ordering::SeqCst), 1);
}
#[test]
fn a_deferred_stage_that_fails_names_its_type() {
let stack = MiddlewareStack::new()
.deferred(|_| Err::<First, _>(Error::internal("AuthManager is not bound")));
let err = stack.resolve(&Container::new()).err().expect("resolution should fail");
assert!(err.message().contains("`First`"), "{}", err.message());
assert!(err.message().contains("AuthManager is not bound"), "{}", err.message());
}
#[test]
fn resolved_builds_from_a_binding() {
struct Guard(&'static str);
let container = Container::new();
container.instance(Guard("bound"));
let stack = MiddlewareStack::new().resolved(|guard: Arc<Guard>| {
assert_eq!(guard.0, "bound");
First
});
let built = stack.resolve(&container).unwrap();
assert_eq!(built.iter().map(|m| m.name()).collect::<Vec<_>>(), vec!["First"]);
}
#[test]
fn resolved_reports_the_missing_binding_by_type() {
struct NeverBound;
let stack = MiddlewareStack::new().resolved(|_: Arc<NeverBound>| First);
let err = stack.resolve(&Container::new()).err().expect("should fail");
assert!(err.message().contains("`First`"), "{}", err.message());
assert!(err.message().contains("NeverBound"), "{}", err.message());
}
#[test]
fn a_deferred_stage_labels_itself_before_it_is_built() {
let stack = MiddlewareStack::new().with(First).deferred(|_| Ok(Second));
assert_eq!(stack.labels(), vec!["First", "Second"]);
assert_eq!(format!("{stack:?}"), r#"["First", "Second"]"#);
}
#[test]
fn a_shared_instance_is_not_cloned_per_route() {
let shared: Arc<dyn Middleware> = Arc::new(First);
let a = MiddlewareStack::new().with_arc(Arc::clone(&shared));
let b = MiddlewareStack::new().with_arc(Arc::clone(&shared));
let container = Container::new();
assert!(Arc::ptr_eq(
&a.resolve(&container).unwrap()[0],
&b.resolve(&container).unwrap()[0]
));
}
}