use crate::call::Call;
use crate::response::Response;
use async_trait::async_trait;
use std::collections::VecDeque;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum Phase {
Setup,
Monitoring,
Plugins,
Call,
Fallback,
}
pub type Endpoint =
Arc<dyn Fn(Call) -> Pin<Box<dyn Future<Output = Response> + Send + 'static>> + Send + Sync>;
#[async_trait]
pub trait Middleware: Send + Sync + 'static {
async fn handle(&self, call: Call, next: Next) -> Response;
}
pub struct Next {
chain: VecDeque<Arc<dyn Middleware>>,
endpoint: Endpoint,
}
impl Next {
pub(crate) fn new(chain: VecDeque<Arc<dyn Middleware>>, endpoint: Endpoint) -> Self {
Self { chain, endpoint }
}
pub async fn run(mut self, call: Call) -> Response {
match self.chain.pop_front() {
Some(mw) => mw.handle(call, self).await,
None => (self.endpoint)(call).await,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::response::IntoResponse;
#[test]
fn phases_are_ordered() {
assert!(Phase::Setup < Phase::Monitoring);
assert!(Phase::Monitoring < Phase::Plugins);
assert!(Phase::Plugins < Phase::Call);
assert!(Phase::Call < Phase::Fallback);
}
use bytes::Bytes;
use http::header::HeaderName;
use http::{HeaderMap, HeaderValue, Method, StatusCode, Uri};
fn sample_call() -> Call {
Call::new(
Method::GET,
"/".parse::<Uri>().unwrap(),
HeaderMap::new(),
Bytes::new(),
)
}
fn endpoint() -> Endpoint {
Arc::new(|_call: Call| Box::pin(async { "inner".into_response() }) as _)
}
struct AddHeader;
#[async_trait]
impl Middleware for AddHeader {
async fn handle(&self, call: Call, next: Next) -> Response {
let mut res = next.run(call).await;
res.headers.insert(
HeaderName::from_static("x-mw"),
HeaderValue::from_static("1"),
);
res
}
}
struct ShortCircuit;
#[async_trait]
impl Middleware for ShortCircuit {
async fn handle(&self, _call: Call, _next: Next) -> Response {
Response::new(StatusCode::FORBIDDEN)
}
}
#[tokio::test]
async fn runs_endpoint_when_chain_empty() {
let next = Next::new(VecDeque::new(), endpoint());
let res = next.run(sample_call()).await;
assert_eq!(res.body, Bytes::from("inner"));
}
#[tokio::test]
async fn middleware_post_processes_response() {
let mut chain: VecDeque<Arc<dyn Middleware>> = VecDeque::new();
chain.push_back(Arc::new(AddHeader));
let res = Next::new(chain, endpoint()).run(sample_call()).await;
assert_eq!(res.headers.get("x-mw").unwrap(), "1");
assert_eq!(res.body, Bytes::from("inner"));
}
#[tokio::test]
async fn middleware_can_short_circuit() {
let mut chain: VecDeque<Arc<dyn Middleware>> = VecDeque::new();
chain.push_back(Arc::new(ShortCircuit));
chain.push_back(Arc::new(AddHeader)); let res = Next::new(chain, endpoint()).run(sample_call()).await;
assert_eq!(res.status, StatusCode::FORBIDDEN);
assert!(res.headers.get("x-mw").is_none());
}
}