Skip to main content

tachyon_web/routing/
middleware.rs

1use crate::http::response::Body;
2use crate::routing::handler::BoxedHandler;
3use hyper::{Request, Response};
4use std::sync::Arc;
5
6/// The continuation for the next handler or middleware in the chain.
7///
8/// Middleware functions take `Next` and call `next.run(req).await` to execute
9/// the remaining pipeline.
10pub struct Next<S> {
11    pub(crate) handler: BoxedHandler<S>,
12    pub(crate) state: Arc<S>,
13}
14
15impl<S> std::fmt::Debug for Next<S> {
16    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
17        f.debug_struct("Next").finish_non_exhaustive()
18    }
19}
20
21impl<S: Send + Sync + 'static> Next<S> {
22    /// Executes the next handler in the pipeline.
23    #[inline]
24    pub async fn run(self, req: Request<Body>) -> Response<Body> {
25        let Self { handler, state } = self;
26        handler(req, state).await
27    }
28
29    /// Access the shared application state from within a middleware.
30    #[inline]
31    #[must_use]
32    pub fn state(&self) -> &S {
33        &self.state
34    }
35}
36
37/// A boxed middleware closure.
38pub type BoxedMiddleware<S> =
39    Arc<dyn Fn(Request<Body>, Next<S>) -> crate::routing::handler::BoxedFuture + Send + Sync>;
40
41/// Position of the middleware in the execution chain.
42#[derive(Debug, Clone, Copy, PartialEq, Eq)]
43pub enum MiddlewarePosition {
44    /// Execute this middleware first (outermost layer).
45    First,
46    /// Execute this middleware last (innermost layer, right before the route handler).
47    Last,
48}
49
50/// A wrapper around a route handler and its associated middlewares.
51#[derive(Clone)]
52pub struct MethodHandler<S> {
53    pub(crate) raw: BoxedHandler<S>,
54    pub(crate) middlewares: Vec<BoxedMiddleware<S>>,
55    pub(crate) compiled: Option<BoxedHandler<S>>,
56}
57
58impl<S> std::fmt::Debug for MethodHandler<S> {
59    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
60        f.debug_struct("MethodHandler")
61            .field("middlewares_count", &self.middlewares.len())
62            .field("compiled", &self.compiled.is_some())
63            .finish_non_exhaustive()
64    }
65}
66
67impl<S: Send + Sync + 'static> MethodHandler<S> {
68    /// Create a new `MethodHandler`.
69    pub fn new(raw: BoxedHandler<S>) -> Self {
70        Self {
71            raw,
72            middlewares: Vec::new(),
73            compiled: None,
74        }
75    }
76
77    /// Compile the middleware chain into a single boxed handler.
78    pub fn compile_in_place(&mut self) {
79        if self.compiled.is_none() {
80            let mut handler = self.raw.clone();
81            for mw in self.middlewares.iter().rev() {
82                let h = handler.clone();
83                let mw_clone = mw.clone();
84                handler = Arc::new(move |req, state| {
85                    let next = Next {
86                        handler: h.clone(),
87                        state,
88                    };
89                    mw_clone(req, next)
90                });
91            }
92            self.compiled = Some(handler);
93        }
94    }
95
96    /// Execute the handler chain.
97    pub async fn call(&self, req: Request<Body>, state: Arc<S>) -> Response<Body> {
98        if let Some(compiled) = &self.compiled {
99            compiled(req, state).await
100        } else {
101            let mut handler = self.raw.clone();
102            for mw in self.middlewares.iter().rev() {
103                let h = handler.clone();
104                let mw_clone = mw.clone();
105                handler = Arc::new(move |req, state| {
106                    let next = Next {
107                        handler: h.clone(),
108                        state,
109                    };
110                    mw_clone(req, next)
111                });
112            }
113            handler(req, state).await
114        }
115    }
116}
117
118#[cfg(test)]
119mod tests {
120    use super::{MethodHandler, MiddlewarePosition, Next};
121    use crate::http::response::{Body, IntoResponse};
122    use crate::routing::handler::{BoxedFuture, BoxedHandler, ResponseFuture};
123    use hyper::{Request, Response};
124    use std::sync::Arc;
125
126    fn handler_returning(body: &'static str) -> BoxedHandler<()> {
127        Arc::new(move |_req, _state| {
128            ResponseFuture::Boxed(Box::pin(async move { body.into_response() }))
129        })
130    }
131
132    fn tag_header_middleware(
133        name: &'static str,
134    ) -> Arc<dyn Fn(Request<Body>, Next<()>) -> BoxedFuture + Send + Sync> {
135        Arc::new(move |req, next| {
136            ResponseFuture::Boxed(Box::pin(async move {
137                let mut resp = next.run(req).await;
138                resp.headers_mut()
139                    .append("x-mw", name.parse().expect("valid header value"));
140                resp
141            }))
142        })
143    }
144
145    #[test]
146    fn middleware_position_variants_are_distinguishable() {
147        assert_ne!(MiddlewarePosition::First, MiddlewarePosition::Last);
148    }
149
150    #[test]
151    fn next_debug_does_not_panic() {
152        let next = Next {
153            handler: handler_returning("unused"),
154            state: Arc::new(()),
155        };
156        assert!(format!("{next:?}").contains("Next"));
157    }
158
159    #[test]
160    fn method_handler_debug_reports_middleware_count_and_compiled_state() {
161        let mut mh = MethodHandler::new(handler_returning("hi"));
162        assert!(format!("{mh:?}").contains("compiled: false"));
163
164        mh.middlewares.push(tag_header_middleware("a"));
165        mh.compile_in_place();
166        let debug = format!("{mh:?}");
167        assert!(debug.contains("middlewares_count: 1"));
168        assert!(debug.contains("compiled: true"));
169    }
170
171    /// `call()` must run correctly even before `compile_in_place()` has ever been called —
172    /// the uncompiled fallback path rebuilds the chain on every call instead of using the
173    /// cached `compiled` handler.
174    #[tokio::test]
175    async fn call_runs_the_middleware_chain_even_when_not_yet_compiled() {
176        let mut mh = MethodHandler::new(handler_returning("body"));
177        mh.middlewares.push(tag_header_middleware("outer"));
178        mh.middlewares.push(tag_header_middleware("inner"));
179        assert!(mh.compiled.is_none());
180
181        let req = Request::builder().body(Body::empty()).unwrap();
182        let resp = mh.call(req, Arc::new(())).await;
183
184        let tags: Vec<&str> = resp
185            .headers()
186            .get_all("x-mw")
187            .iter()
188            .map(|v| v.to_str().unwrap())
189            .collect();
190        // Middlewares run in registration order (each wraps the next), so the first
191        // registered ("outer") is the outermost — its header gets appended last.
192        assert_eq!(tags, vec!["inner", "outer"]);
193    }
194
195    #[tokio::test]
196    async fn call_uses_the_cached_compiled_handler_once_compiled() {
197        let mut mh = MethodHandler::new(handler_returning("body"));
198        mh.middlewares.push(tag_header_middleware("only"));
199        mh.compile_in_place();
200        assert!(mh.compiled.is_some());
201
202        let req = Request::builder().body(Body::empty()).unwrap();
203        let resp: Response<Body> = mh.call(req, Arc::new(())).await;
204        assert_eq!(
205            resp.headers().get("x-mw").unwrap().to_str().unwrap(),
206            "only"
207        );
208    }
209}