use std::sync::Arc;
use axum::Router;
pub const HTTP_INTERCEPTOR_ORDER: [&str; 5] =
["tracing", "logging", "auth", "validation", "metrics"];
pub const HTTP_BASELINE_LAYER_ORDER: [&str; 5] = [
"request_id",
"cors",
"security_headers",
"body_limit",
"timeout",
];
pub type RouterTransform = Arc<dyn Fn(Router) -> Router + Send + Sync + 'static>;
#[derive(Clone, Default)]
pub struct HttpMiddlewareStack {
tracing: Vec<RouterTransform>,
logging: Vec<RouterTransform>,
auth: Vec<RouterTransform>,
validation: Vec<RouterTransform>,
metrics: Vec<RouterTransform>,
}
impl HttpMiddlewareStack {
#[must_use]
pub fn new() -> Self {
Self::default()
}
#[must_use]
pub fn with_tracing_transform<F>(mut self, transform: F) -> Self
where
F: Fn(Router) -> Router + Send + Sync + 'static,
{
self.tracing.push(Arc::new(transform));
self
}
#[must_use]
pub fn with_logging_transform<F>(mut self, transform: F) -> Self
where
F: Fn(Router) -> Router + Send + Sync + 'static,
{
self.logging.push(Arc::new(transform));
self
}
#[must_use]
pub fn with_auth_transform<F>(mut self, transform: F) -> Self
where
F: Fn(Router) -> Router + Send + Sync + 'static,
{
self.auth.push(Arc::new(transform));
self
}
#[must_use]
pub fn with_validation_transform<F>(mut self, transform: F) -> Self
where
F: Fn(Router) -> Router + Send + Sync + 'static,
{
self.validation.push(Arc::new(transform));
self
}
#[must_use]
pub fn with_metrics_transform<F>(mut self, transform: F) -> Self
where
F: Fn(Router) -> Router + Send + Sync + 'static,
{
self.metrics.push(Arc::new(transform));
self
}
pub fn apply(&self, router: Router) -> Router {
let router = apply_phase(router, &self.metrics);
let router = apply_phase(router, &self.validation);
let router = apply_phase(router, &self.auth);
let router = apply_phase(router, &self.logging);
apply_phase(router, &self.tracing)
}
}
fn apply_phase(router: Router, transforms: &[RouterTransform]) -> Router {
transforms
.iter()
.fold(router, |router, transform| transform(router))
}
#[cfg(test)]
mod tests {
use std::convert::Infallible;
use std::sync::Arc;
use axum::{Router, body::Body, http::Request, response::Response, routing::get};
use parking_lot::Mutex;
use tower::{Layer, Service, ServiceExt};
use super::HttpMiddlewareStack;
#[derive(Clone)]
struct RecordLayer {
name: &'static str,
events: Arc<Mutex<Vec<&'static str>>>,
}
impl RecordLayer {
fn new(name: &'static str, events: Arc<Mutex<Vec<&'static str>>>) -> Self {
Self { name, events }
}
}
impl<S> Layer<S> for RecordLayer {
type Service = RecordService<S>;
fn layer(&self, inner: S) -> Self::Service {
RecordService {
inner,
name: self.name,
events: Arc::clone(&self.events),
}
}
}
#[derive(Clone)]
struct RecordService<S> {
inner: S,
name: &'static str,
events: Arc<Mutex<Vec<&'static str>>>,
}
impl<S> Service<Request<Body>> for RecordService<S>
where
S: Service<Request<Body>, Response = Response, Error = Infallible> + Clone + Send + 'static,
S::Future: Send + 'static,
{
type Response = Response;
type Error = Infallible;
type Future = futures::future::BoxFuture<'static, Result<Response, Infallible>>;
fn poll_ready(
&mut self,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Result<(), Self::Error>> {
self.inner.poll_ready(cx)
}
fn call(&mut self, req: Request<Body>) -> Self::Future {
let mut inner = self.inner.clone();
let events = Arc::clone(&self.events);
let name = self.name;
Box::pin(async move {
events.lock().push(name);
inner.call(req).await
})
}
}
#[tokio::test]
async fn stack_applies_request_phases_in_locked_order() {
let events = Arc::new(Mutex::new(Vec::new()));
let app = Router::new().route("/", get(|| async { "ok" }));
let app = HttpMiddlewareStack::new()
.with_tracing_transform({
let events = Arc::clone(&events);
move |router| router.layer(RecordLayer::new("tracing", Arc::clone(&events)))
})
.with_logging_transform({
let events = Arc::clone(&events);
move |router| router.layer(RecordLayer::new("logging", Arc::clone(&events)))
})
.with_auth_transform({
let events = Arc::clone(&events);
move |router| router.layer(RecordLayer::new("auth", Arc::clone(&events)))
})
.with_validation_transform({
let events = Arc::clone(&events);
move |router| router.layer(RecordLayer::new("validation", Arc::clone(&events)))
})
.with_metrics_transform({
let events = Arc::clone(&events);
move |router| router.layer(RecordLayer::new("metrics", Arc::clone(&events)))
})
.apply(app);
let response = app
.oneshot(Request::builder().uri("/").body(Body::empty()).unwrap())
.await
.unwrap();
assert_eq!(response.status(), axum::http::StatusCode::OK);
assert_eq!(
*events.lock(),
vec!["tracing", "logging", "auth", "validation", "metrics"]
);
}
}