rustlavel_http/
middleware.rs1use crate::handler::{BoxFuture, Handler};
8use crate::request::Request;
9use crate::response::Response;
10use std::future::Future;
11use std::sync::Arc;
12
13pub trait Middleware: Send + Sync + 'static {
14 fn handle(&self, request: Request, next: Next) -> BoxFuture<Response>;
15}
16
17impl<F, Fut> Middleware for F
18where
19 F: Fn(Request, Next) -> Fut + Send + Sync + 'static,
20 Fut: Future<Output = Response> + Send + 'static,
21{
22 fn handle(&self, request: Request, next: Next) -> BoxFuture<Response> {
23 Box::pin(self(request, next))
24 }
25}
26
27pub struct Next {
29 stack: Arc<Vec<Arc<dyn Middleware>>>,
30 index: usize,
31 endpoint: Arc<dyn Handler>,
32}
33
34impl Next {
35 pub(crate) fn new(stack: Arc<Vec<Arc<dyn Middleware>>>, endpoint: Arc<dyn Handler>) -> Self {
36 Next { stack, index: 0, endpoint }
37 }
38
39 pub fn run(self, request: Request) -> BoxFuture<Response> {
41 match self.stack.get(self.index).cloned() {
42 Some(middleware) => {
43 let next = Next {
44 stack: Arc::clone(&self.stack),
45 index: self.index + 1,
46 endpoint: Arc::clone(&self.endpoint),
47 };
48 middleware.handle(request, next)
49 }
50 None => self.endpoint.call(request),
51 }
52 }
53}
54
55#[cfg(test)]
56mod tests {
57 use super::*;
58 use crate::method::Method;
59 use crate::status::Status;
60
61 fn pipeline(stack: Vec<Arc<dyn Middleware>>, endpoint: Arc<dyn Handler>) -> Next {
62 Next::new(Arc::new(stack), endpoint)
63 }
64
65 #[tokio::test]
66 async fn middleware_runs_around_the_handler() {
67 let tag: Arc<dyn Middleware> = Arc::new(|request: Request, next: Next| async move {
68 let response = next.run(request).await;
69 response.with_header("x-tag", "seen")
70 });
71
72 let endpoint: Arc<dyn Handler> =
73 Arc::new(|_req: Request| async { Response::text("handled") });
74
75 let response = pipeline(vec![tag], endpoint).run(Request::new(Method::Get, "/")).await;
76
77 assert_eq!(response.body_string(), "handled");
78 assert_eq!(response.headers.get("x-tag"), Some("seen"));
79 }
80
81 #[tokio::test]
82 async fn a_middleware_can_short_circuit() {
83 let guard: Arc<dyn Middleware> = Arc::new(|_req: Request, _next: Next| async {
84 Response::new(Status::UNAUTHORIZED).with_text("denied")
85 });
86
87 let endpoint: Arc<dyn Handler> =
88 Arc::new(|_req: Request| async {
89 panic!("the handler must not run");
90 #[allow(unreachable_code)]
91 Response::ok()
92 });
93
94 let response = pipeline(vec![guard], endpoint).run(Request::new(Method::Get, "/")).await;
95
96 assert_eq!(response.status, Status::UNAUTHORIZED);
97 }
98
99 #[tokio::test]
100 async fn middleware_runs_in_registration_order() {
101 let order = Arc::new(std::sync::Mutex::new(Vec::new()));
102
103 let make = |label: &'static str, log: Arc<std::sync::Mutex<Vec<&'static str>>>| {
104 let middleware: Arc<dyn Middleware> = Arc::new(move |request: Request, next: Next| {
105 let log = Arc::clone(&log);
106 async move {
107 log.lock().unwrap().push(label);
108 next.run(request).await
109 }
110 });
111 middleware
112 };
113
114 let stack = vec![make("first", Arc::clone(&order)), make("second", Arc::clone(&order))];
115 let endpoint: Arc<dyn Handler> = Arc::new(|_req: Request| async { Response::ok() });
116
117 pipeline(stack, endpoint).run(Request::new(Method::Get, "/")).await;
118
119 assert_eq!(*order.lock().unwrap(), ["first", "second"]);
120 }
121}