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

spire_core/worker/middleware/
mod.rs
