Skip to main content

ntex_service/
middleware.rs

1use std::{fmt, marker::PhantomData, rc::Rc};
2
3use crate::dev::{Apply, ApplyCtx};
4use crate::{IntoServiceFactory, Service, ServiceChainFactory, ServiceFactory};
5
6/// Applies middleware to every service produced by a factory.
7pub fn apply<Sf, St, Req, M>(
8    mw: M,
9    factory: impl IntoServiceFactory<Sf, St, Req>,
10) -> ServiceChainFactory<ApplyMiddleware<M, Sf>, St, Req>
11where
12    Sf: ServiceFactory<St, Req>,
13    M: Middleware<Sf::Service, St>,
14{
15    ServiceChainFactory {
16        factory: ApplyMiddleware::new(mw, factory.into_factory()),
17        _t: PhantomData,
18    }
19}
20
21/// Wraps an inner service during service construction.
22///
23/// Middleware can run before and after the inner service, and can modify
24/// requests, responses, or errors.
25///
26/// For example, timeout middleware:
27///
28/// ```rust
29/// use ntex_service::{Ctx, Service};
30/// use ntex::{time::sleep, util::Either, util::select};
31///
32/// pub struct Timeout<S> {
33///     service: S,
34///     timeout: std::time::Duration,
35/// }
36///
37/// pub enum TimeoutError<E> {
38///    Service(E),
39///    Timeout,
40/// }
41///
42/// impl<S, R> Service<(), R> for Timeout<S>
43/// where
44///     S: Service<(), R>,
45/// {
46///     type Res = S::Res;
47///     type Error = TimeoutError<S::Error>;
48///
49///     ntex_service::forward_bind!((), service);
50///
51///     async fn ready(&self, ctx: Ctx<'_, Self>) -> Result<(), Self::Error> {
52///         ctx.ready(&self.service).await.map_err(TimeoutError::Service)
53///     }
54///
55///     async fn call(&self, req: R, ctx: Ctx<'_, Self>) -> Result<Self::Res, Self::Error> {
56///         match select(sleep(self.timeout), ctx.call(&self.service, req)).await {
57///             Either::Left(_) => Err(TimeoutError::Timeout),
58///             Either::Right(res) => res.map_err(TimeoutError::Service),
59///         }
60///     }
61/// }
62/// ```
63///
64/// The timeout service is independent of the wrapped service implementation and
65/// can be applied to any compatible service.
66///
67/// A middleware factory for `Timeout` could look like this:
68///
69/// ```rust
70/// # use ntex_service::Middleware;
71/// # pub struct Timeout<S> {
72/// #     service: S,
73/// #     timeout: std::time::Duration,
74/// # }
75/// pub struct TimeoutMiddleware {
76///     timeout: std::time::Duration,
77/// }
78///
79/// impl<S> Middleware<S, ()> for TimeoutMiddleware
80/// {
81///     type Service = Timeout<S>;
82///
83///     fn create(&self, _: &(), service: S) -> Self::Service {
84///         Timeout {
85///             service,
86///             timeout: self.timeout,
87///         }
88///     }
89/// }
90/// ```
91pub trait Middleware<S, St> {
92    /// Service created by this middleware.
93    type Service;
94
95    /// Creates and returns a new middleware service.
96    fn create(&self, st: &St, service: S) -> Self::Service;
97
98    /// Creates a service factory that instantiates a service and applies
99    /// the current middleware to it.
100    ///
101    /// This is equivalent to `apply(self, factory)`.
102    fn apply_to<Sf, Req>(
103        self,
104        factory: Sf,
105    ) -> ServiceChainFactory<ApplyMiddleware<Self, Sf>, St, Req>
106    where
107        Sf: ServiceFactory<St, Req, Service = S>,
108        Self: Sized,
109        Self::Service: Service<St, Req>,
110    {
111        crate::factory(ApplyMiddleware::new(self, factory))
112    }
113}
114
115impl<M, S, St> Middleware<S, St> for Rc<M>
116where
117    M: Middleware<S, St>,
118{
119    type Service = M::Service;
120
121    fn create(&self, st: &St, service: S) -> M::Service {
122        self.as_ref().create(st, service)
123    }
124}
125
126/// A service factory with middleware applied.
127pub struct ApplyMiddleware<M, Sf>(Rc<(M, Sf)>);
128
129impl<M, Sf> ApplyMiddleware<M, Sf> {
130    /// Creates a new `ApplyMiddleware` factory.
131    pub(crate) fn new(mw: M, sf: Sf) -> Self {
132        Self(Rc::new((mw, sf)))
133    }
134}
135
136impl<M, Sf> Clone for ApplyMiddleware<M, Sf> {
137    fn clone(&self) -> Self {
138        Self(self.0.clone())
139    }
140}
141
142impl<M, Sf> fmt::Debug for ApplyMiddleware<M, Sf>
143where
144    M: fmt::Debug,
145    Sf: fmt::Debug,
146{
147    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
148        f.debug_struct("ApplyMiddleware")
149            .field("factory", &self.0.1)
150            .field("middleware", &self.0.0)
151            .finish()
152    }
153}
154
155impl<M, Sf, St, Req> ServiceFactory<St, Req> for ApplyMiddleware<M, Sf>
156where
157    Sf: ServiceFactory<St, Req>,
158    M: Middleware<Sf::Service, St>,
159    M::Service: Service<St, Req>,
160{
161    type Res = <M::Service as Service<St, Req>>::Res;
162    type Error = <M::Service as Service<St, Req>>::Error;
163
164    type Service = M::Service;
165    type InitError = Sf::InitError;
166
167    #[inline]
168    async fn create(&self, st: &St) -> Result<Self::Service, Self::InitError> {
169        Ok(self.0.0.create(st, self.0.1.create(st).await?))
170    }
171}
172
173/// Middleware that returns the wrapped service unchanged.
174#[derive(Debug, Clone, Copy)]
175pub struct Identity;
176
177impl<S, St> Middleware<S, St> for Identity {
178    type Service = S;
179
180    #[inline]
181    fn create(&self, _: &St, service: S) -> Self::Service {
182        service
183    }
184}
185
186/// Two middleware values applied in sequence.
187///
188/// The inner middleware is applied first, then the outer middleware wraps its
189/// service.
190#[derive(Debug, Clone)]
191pub struct Stack<Inner, Outer> {
192    inner: Inner,
193    outer: Outer,
194}
195
196impl<Inner, Outer> Stack<Inner, Outer> {
197    /// Creates a middleware stack.
198    pub fn new(inner: Inner, outer: Outer) -> Self {
199        Stack { inner, outer }
200    }
201}
202
203impl<S, St, Inner, Outer> Middleware<S, St> for Stack<Inner, Outer>
204where
205    Inner: Middleware<S, St>,
206    Outer: Middleware<Inner::Service, St>,
207{
208    type Service = Outer::Service;
209
210    fn create(&self, st: &St, service: S) -> Self::Service {
211        self.outer.create(st, self.inner.create(st, service))
212    }
213}
214
215#[doc(hidden)]
216/// Creates middleware from an asynchronous function.
217///
218/// The function receives an input request and an [`ApplyCtx`] that can call the
219/// wrapped service.
220pub fn fn_layer<F, S, St, Req, In, Out, Err>(f: F) -> FnMiddleware<F, S, St, Req, In, Out, Err>
221where
222    F: AsyncFn(In, &ApplyCtx<'_, S, St, Req>) -> Result<Out, Err> + Clone,
223    S: Service<St, Req>,
224{
225    FnMiddleware { f, r: PhantomData }
226}
227
228/// Middleware backed by an asynchronous function.
229pub struct FnMiddleware<F, S, St, Req, In, Out, Err> {
230    f: F,
231    r: PhantomData<fn(S, St, Req) -> (In, Out, Err)>,
232}
233
234impl<F, S, St, Req, In, Out, Err> Clone for FnMiddleware<F, S, St, Req, In, Out, Err>
235where
236    F: Clone,
237{
238    fn clone(&self) -> Self {
239        FnMiddleware {
240            f: self.f.clone(),
241            r: PhantomData,
242        }
243    }
244}
245
246impl<F, S, St, Req, In, Out, Err> fmt::Debug for FnMiddleware<F, S, St, Req, In, Out, Err> {
247    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
248        f.debug_struct("FnMiddleware")
249            .field("layer", &std::any::type_name::<F>())
250            .finish()
251    }
252}
253
254impl<F, S, St, Req, In, Out, Err> Middleware<S, St> for FnMiddleware<F, S, St, Req, In, Out, Err>
255where
256    S: Service<St, Req>,
257    F: AsyncFn(In, &ApplyCtx<'_, S, St, Req>) -> Result<Out, Err> + Clone,
258    Err: From<S::Error>,
259{
260    type Service = Apply<S, St, Req, F, In, Out, Err>;
261
262    fn create(&self, _: &St, service: S) -> Self::Service {
263        Apply::new(service, self.f.clone())
264    }
265}
266
267#[cfg(test)]
268#[allow(clippy::redundant_clone)]
269mod tests {
270    use std::{cell::Cell, rc::Rc};
271
272    use super::*;
273    use crate::{Ctx, Pipeline, factory, fn_service};
274
275    #[derive(Debug, Clone)]
276    struct Mw(Rc<Cell<usize>>);
277
278    impl<S, St> Middleware<S, St> for Mw {
279        type Service = Srv<S>;
280
281        fn create(&self, _: &St, service: S) -> Self::Service {
282            self.0.set(self.0.get() + 1);
283            Srv(service, self.0.clone())
284        }
285    }
286
287    #[derive(Debug, Clone)]
288    struct Srv<S>(S, Rc<Cell<usize>>);
289
290    impl<S: Service<(), R>, R> Service<(), R> for Srv<S> {
291        type Res = S::Res;
292        type Error = S::Error;
293
294        crate::forward_bind!((), 0);
295
296        async fn ready(&self, ctx: Ctx<'_, Self>) -> Result<(), Self::Error> {
297            ctx.ready(&self.0).await
298        }
299
300        async fn call(&self, req: R, ctx: Ctx<'_, Self>) -> Result<S::Res, S::Error> {
301            ctx.call(&self.0, req).await
302        }
303
304        async fn shutdown(&self, _: Ctx<'_, Self, ()>) {
305            self.1.set(self.1.get() + 1);
306        }
307    }
308
309    #[ntex::test]
310    async fn middleware() {
311        let cnt_sht = Rc::new(Cell::new(0));
312        let fac = apply(
313            Rc::new(Mw(cnt_sht.clone()).clone()),
314            fn_service(|i: usize| async move { Ok::<_, ()>(i * 2) }),
315        )
316        .clone();
317
318        let srv = Pipeline::new((), fac.create(&()).await.unwrap().clone());
319        let res = srv.call(10).await;
320        assert!(res.is_ok());
321        assert_eq!(res.unwrap(), 20);
322        let _ = format!("{fac:?} {srv:?}");
323
324        assert_eq!(srv.ready().await, Ok(()));
325        srv.shutdown().await;
326        assert_eq!(cnt_sht.get(), 2);
327
328        let fac = factory(fn_service(|i: usize| async move { Ok::<_, ()>(i * 2) }))
329            .apply(Rc::new(Mw(Rc::new(Cell::new(0))).clone()))
330            .clone();
331
332        let srv = Pipeline::new((), fac.create(&()).await.unwrap().clone());
333        let res = srv.call(10).await;
334        assert!(res.is_ok());
335        assert_eq!(res.unwrap(), 20);
336        let _ = format!("{fac:?} {srv:?}");
337
338        assert_eq!(srv.ready().await, Ok(()));
339    }
340
341    #[ntex::test]
342    async fn middleware_apply() {
343        let cnt_sht = Rc::new(Cell::new(0));
344        let fac = Mw(cnt_sht.clone())
345            .apply_to(factory(async |i: usize| Ok::<_, ()>(i * 2)))
346            .boxed();
347
348        let srv = Pipeline::new((), fac.create(&()).await.unwrap());
349        let res = srv.call(10).await;
350        assert!(res.is_ok());
351        assert_eq!(res.unwrap(), 20);
352        let _ = format!("{fac:?} {srv:?}");
353
354        assert_eq!(srv.ready().await, Ok(()));
355        srv.shutdown().await;
356        assert_eq!(cnt_sht.get(), 2);
357    }
358
359    #[ntex::test]
360    async fn middleware_chain() {
361        let cnt_sht = Rc::new(Cell::new(0));
362        let fac = factory(fn_service(async move |i: usize| Ok::<_, ()>(i * 2)))
363            .apply(Mw(cnt_sht.clone()).clone());
364
365        let srv = Pipeline::new((), fac.create(&()).await.unwrap().clone());
366        let res = srv.call(10).await;
367        assert!(res.is_ok());
368        assert_eq!(res.unwrap(), 20);
369        let _ = format!("{fac:?} {srv:?}");
370
371        assert_eq!(srv.ready().await, Ok(()));
372        srv.shutdown().await;
373        assert_eq!(cnt_sht.get(), 2);
374    }
375
376    #[ntex::test]
377    async fn stack() {
378        let cnt_sht = Rc::new(Cell::new(0));
379        let mw = Stack::new(Identity, Mw(cnt_sht.clone()));
380        let _ = format!("{mw:?}");
381
382        let pl = Pipeline::new(
383            (),
384            Middleware::create(
385                &mw,
386                &(),
387                fn_service(|i: usize| async move { Ok::<_, ()>(i * 2) }),
388            ),
389        );
390        let res = pl.call(10).await;
391        assert!(res.is_ok());
392        assert_eq!(res.unwrap(), 20);
393        assert_eq!(pl.ready().await, Ok(()));
394        pl.shutdown().await;
395        assert_eq!(cnt_sht.get(), 2);
396    }
397
398    #[ntex::test]
399    async fn fn_middleware_service() {
400        let cnt_sht = Rc::new(Cell::new(0));
401        let cnt_sht2 = cnt_sht.clone();
402        let mw = fn_layer(async move |req: &'static str, svc| {
403            cnt_sht2.set(cnt_sht2.get() + 1);
404            let result = svc.call(1).await?;
405            Ok::<_, ()>((req, result))
406        })
407        .clone();
408        let _ = format!("{mw:?}");
409
410        let svc = Pipeline::new(
411            (),
412            mw.create(&(), fn_service(async move |i: usize| Ok::<_, ()>(i * 2))),
413        );
414
415        let res = svc.call("test").await;
416        assert!(res.is_ok());
417        assert_eq!(res.unwrap(), ("test", 2));
418        let _ = format!("{svc:?}");
419
420        assert_eq!(svc.ready().await, Ok(()));
421        svc.shutdown().await;
422        assert_eq!(cnt_sht.get(), 1);
423    }
424}