spire_core/worker/middleware/
mod.rs

1//! Middleware, which runs around a [`Worker`] for every request.
2//!
3//! A [`Middleware`] receives the [`Context`] before any handler, and the
4//! [`Next`] worker to pass it on to. Before passing it on, it can change the
5//! request, e.g. add headers, or answer without fetching. Around, it can
6//! check the response's head: fetched once, it is shared with the handlers.
7//! After, it can change the [`Outcome`].
8//!
9//! [`WorkerExt::layer`] wraps a worker, such as a router, in middleware. The
10//! middleware added last runs first.
11//!
12//! # Examples
13//!
14//! ```
15//! use std::time::Duration;
16//!
17//! use spire_core::backend::Noop;
18//! use spire_core::http::http::HeaderValue;
19//! use spire_core::http::http::header::USER_AGENT;
20//! use spire_core::worker::middleware::{self, Next, SetHeader};
21//! use spire_core::worker::{Outcome, WorkerExt, from_fn};
22//! use spire_core::crawler::Crawler;
23//! use spire_core::{Context, Result};
24//!
25//! /// Retries pages that answer with a captcha instead of content.
26//! async fn captcha(cx: Context<Noop>, next: Next<Noop>) -> Result<Outcome> {
27//!     if cx.response().await?.headers().contains_key("x-captcha") {
28//!         return Ok(Outcome::Retry { after: Some(Duration::from_secs(60)) });
29//!     }
30//!     Ok(next.run(cx).await)
31//! }
32//!
33//! let worker = from_fn(|_: Context<Noop>| async {})
34//!     .layer(middleware::from_fn(captcha))
35//!     .layer(SetHeader::new(USER_AGENT, HeaderValue::from_static("MyBot/1.0")));
36//! let crawler = Crawler::new(Noop, worker);
37//! ```
38
39mod header;
40mod revalidate;
41
42use std::fmt;
43use std::future::Future;
44use std::sync::Arc;
45
46use futures::future::BoxFuture;
47pub use header::SetHeader;
48pub use revalidate::{InMemValidators, Revalidate, Validator, Validators};
49
50use crate::context::Context;
51use crate::worker::{IntoOutcome, Outcome, Worker};
52
53/// Runs around a [`Worker`] for every request, see the [module](self).
54pub trait Middleware<C>: Send + Sync + 'static {
55    /// Handles the request in `cx`, passing it on to `next` or not.
56    fn handle(&self, cx: Context<C>, next: Next<C>) -> impl Future<Output = Outcome> + Send;
57}
58
59/// The rest of the chain a [`Middleware`] passes a request on to: further
60/// middleware, then the worker.
61pub struct Next<C> {
62    worker: Arc<dyn ErasedWorker<C>>,
63}
64
65impl<C> Next<C> {
66    /// Runs the rest of the chain on `cx`.
67    pub async fn run(self, cx: Context<C>) -> Outcome
68    where
69        C: Send + 'static,
70    {
71        self.worker.handle(cx).await
72    }
73}
74
75impl<C> fmt::Debug for Next<C> {
76    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
77        f.debug_struct("Next").finish_non_exhaustive()
78    }
79}
80
81/// A [`Worker`] behind a pointer, so [`Next`] does not name its type.
82trait ErasedWorker<C>: Send + Sync + 'static {
83    fn handle(&self, cx: Context<C>) -> BoxFuture<'_, Outcome>;
84}
85
86impl<C: Send + 'static, W: Worker<C>> ErasedWorker<C> for W {
87    fn handle(&self, cx: Context<C>) -> BoxFuture<'_, Outcome> {
88        Box::pin(Worker::handle(self, cx))
89    }
90}
91
92/// Wraps workers in [`Middleware`].
93pub trait WorkerExt: Sized {
94    /// Wraps this worker in `middleware`, which runs before it, and before
95    /// any middleware added earlier.
96    fn layer<C, M>(self, middleware: M) -> Layered<M, C>
97    where
98        C: Send + 'static,
99        Self: Worker<C>,
100        M: Middleware<C>,
101    {
102        Layered {
103            middleware,
104            worker: Arc::new(self),
105        }
106    }
107}
108
109impl<W> WorkerExt for W {}
110
111/// A [`Worker`] wrapped in [`Middleware`], created by [`WorkerExt::layer`].
112pub struct Layered<M, C> {
113    middleware: M,
114    worker: Arc<dyn ErasedWorker<C>>,
115}
116
117impl<M, C> Worker<C> for Layered<M, C>
118where
119    M: Middleware<C>,
120    C: Send + 'static,
121{
122    fn handle(&self, cx: Context<C>) -> impl Future<Output = Outcome> + Send {
123        let next = Next {
124            worker: Arc::clone(&self.worker),
125        };
126        self.middleware.handle(cx, next)
127    }
128}
129
130impl<M: Clone, C> Clone for Layered<M, C> {
131    fn clone(&self) -> Self {
132        Self {
133            middleware: self.middleware.clone(),
134            worker: Arc::clone(&self.worker),
135        }
136    }
137}
138
139impl<M: fmt::Debug, C> fmt::Debug for Layered<M, C> {
140    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
141        f.debug_struct("Layered")
142            .field("middleware", &self.middleware)
143            .finish_non_exhaustive()
144    }
145}
146
147/// Creates a [`Middleware`] from an async function of the [`Context`] and
148/// the [`Next`] worker, returning anything [`IntoOutcome`], e.g.
149/// `Result<Outcome>`.
150pub const fn from_fn<F>(f: F) -> FromFn<F> {
151    FromFn(f)
152}
153
154/// A [`Middleware`] backed by an async function, created by [`from_fn`].
155#[derive(Clone, Copy)]
156pub struct FromFn<F>(F);
157
158impl<C, F, Fut> Middleware<C> for FromFn<F>
159where
160    C: Send + 'static,
161    F: Fn(Context<C>, Next<C>) -> Fut + Send + Sync + 'static,
162    Fut: Future<Output: IntoOutcome> + Send,
163{
164    fn handle(&self, cx: Context<C>, next: Next<C>) -> impl Future<Output = Outcome> + Send {
165        let future = (self.0)(cx, next);
166        async move { future.await.into_outcome() }
167    }
168}
169
170impl<F> fmt::Debug for FromFn<F> {
171    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
172        f.debug_struct("FromFn").finish_non_exhaustive()
173    }
174}
175
176#[cfg(test)]
177mod tests;
spire_core/worker/middleware/mod.rs

spire_core/worker/middleware/
mod.rs
