Module axum::middleware

source ·
Expand description

Utilities for writing middleware

Table of contents

Intro

axum is unique in that it doesn’t have its own bespoke middleware system and instead integrates with tower. This means the ecosystem of tower and tower-http middleware all work with axum.

While its not necessary to fully understand tower to write or use middleware with axum, having at least a basic understanding of tower’s concepts is recommended. See tower’s guides for a general introduction. Reading the documentation for tower::ServiceBuilder is also recommended.

Applying middleware

axum allows you to add middleware just about anywhere

Applying multiple middleware

Its recommended to use tower::ServiceBuilder to apply multiple middleware at once, instead of calling layer (or route_layer) repeatedly:

use axum::{
    routing::get,
    Extension,
    Router,
};
use tower_http::{trace::TraceLayer};
use tower::ServiceBuilder;

async fn handler() {}

#[derive(Clone)]
struct State {}

let app = Router::new()
    .route("/", get(handler))
    .layer(
        ServiceBuilder::new()
            .layer(TraceLayer::new_for_http())
            .layer(Extension(State {}))
    );

Commonly used middleware

Some commonly used middleware are:

Ordering

When you add middleware with Router::layer (or similar) all previously added routes will be wrapped in the middleware. Generally speaking, this results in middleware being executed from bottom to top.

So if you do this:

use axum::{routing::get, Router};

async fn handler() {}

let app = Router::new()
    .route("/", get(handler))
    .layer(layer_one)
    .layer(layer_two)
    .layer(layer_three);

Think of the middleware as being layered like an onion where each new layer wraps all previous layers:

        requests
           |
           v
+----- layer_three -----+
| +---- layer_two ----+ |
| | +-- layer_one --+ | |
| | |               | | |
| | |    handler    | | |
| | |               | | |
| | +-- layer_one --+ | |
| +---- layer_two ----+ |
+----- layer_three -----+
           |
           v
        responses

That is:

  • First layer_three receives the request
  • It then does its thing and passes the request onto layer_two
  • Which passes the request onto layer_one
  • Which passes the request onto handler where a response is produced
  • That response is then passes to layer_one
  • Then to layer_two
  • And finally to layer_three where it’s returned out of your app

It’s a little more complicated in practice because any middleware is free to return early and not call the next layer, for example if a request cannot be authorized, but its a useful mental model to have.

As previously mentioned its recommended to add multiple middleware using tower::ServiceBuilder, however this impacts ordering:

use tower::ServiceBuilder;
use axum::{routing::get, Router};

async fn handler() {}

let app = Router::new()
    .route("/", get(handler))
    .layer(
        ServiceBuilder::new()
            .layer(layer_one)
            .layer(layer_two)
            .layer(layer_three),
    );

ServiceBuilder works by composing all layers into one such that they run top to bottom. So with the previous code layer_one would receive the request first, then layer_two, then layer_three, then handler, and then the response would bubble back up through layer_three, then layer_two, and finally layer_one.

Executing middleware top to bottom is generally easier to understand and follow mentally which is one of the reasons ServiceBuilder is recommended.

Writing middleware

axum offers many ways of writing middleware, at different levels of abstraction and with different pros and cons.

axum::middleware::from_fn

Use axum::middleware::from_fn to write your middleware when:

  • You’re not comfortable with implementing your own futures and would rather use the familiar async/await syntax.
  • You don’t intend to publish your middleware as a crate for others to use. Middleware written like this are only compatible with axum.

axum::middleware::from_extractor

Use axum::middleware::from_extractor to write your middleware when:

  • You have a type that you sometimes want to use as an extractor and sometimes as a middleware. If you only need your type as a middleware prefer middleware::from_fn.

tower’s combinators

tower has several utility combinators that can be used to perform simple modifications to requests or responses. The most commonly used ones are

You should use these when

  • You want to perform a small ad hoc operation, such as adding a header.
  • You don’t intend to publish your middleware as a crate for others to use.

tower::Service and Pin<Box<dyn Future>>

For maximum control (and a more low level API) you can write you own middleware by implementing tower::Service:

Use tower::Service with Pin<Box<dyn Future>> to write your middleware when:

  • Your middleware needs to be configurable for example via builder methods on your tower::Layer such as tower_http::trace::TraceLayer.
  • You do intend to publish your middleware as a crate for others to use.
  • You’re not comfortable with implementing your own futures.

A decent template for such a middleware could be:

use axum::{
    response::Response,
    body::Body,
    http::Request,
};
use futures::future::BoxFuture;
use tower::{Service, Layer};
use std::task::{Context, Poll};

struct MyLayer;

impl<S> Layer<S> for MyLayer {
    type Service = MyMiddleware<S>;

    fn layer(&self, inner: S) -> Self::Service {
        MyMiddleware { inner }
    }
}

#[derive(Clone)]
struct MyMiddleware<S> {
    inner: S,
}

impl<S> Service<Request<Body>> for MyMiddleware<S>
where
    S: Service<Request<Body>, Response = Response> + Send + 'static,
    S::Future: Send + 'static,
{
    type Response = S::Response;
    type Error = S::Error;
    // `BoxFuture` is a type alias for `Pin<Box<dyn Future + Send + 'a>>`
    type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.inner.poll_ready(cx)
    }

    fn call(&mut self, mut request: Request<Body>) -> Self::Future {
        let future = self.inner.call(request);
        Box::pin(async move {
            let response: Response = future.await?;
            Ok(response)
        })
    }
}

tower::Service and custom futures

If you’re comfortable implementing your own futures (or want to learn it) and need as much control as possible then using tower::Service without boxed futures is the way to go.

Use tower::Service with manual futures to write your middleware when:

  • You want your middleware to have the lowest possible overhead.
  • Your middleware needs to be configurable for example via builder methods on your tower::Layer such as tower_http::trace::TraceLayer.
  • You do intend to publish your middleware as a crate for others to use, perhaps as part of tower-http.
  • You’re comfortable with implementing your own futures, or want to learn how the lower levels of async Rust works.

tower’s “Building a middleware from scratch” guide is a good place to learn how to do this.

Error handling for middleware

axum’s error handling model requires handlers to always return a response. However middleware is one possible way to introduce errors into an application. If hyper receives an error the connection will be closed without sending a response. Thus axum requires those errors to be handled gracefully:

use axum::{
    routing::get,
    error_handling::HandleErrorLayer,
    http::StatusCode,
    BoxError,
    Router,
};
use tower::{ServiceBuilder, timeout::TimeoutLayer};
use std::time::Duration;

async fn handler() {}

let app = Router::new()
    .route("/", get(handler))
    .layer(
        ServiceBuilder::new()
            // this middleware goes above `TimeoutLayer` because it will receive
            // errors returned by `TimeoutLayer`
            .layer(HandleErrorLayer::new(|_: BoxError| async {
                StatusCode::REQUEST_TIMEOUT
            }))
            .layer(TimeoutLayer::new(Duration::from_secs(10)))
    );

See error_handling for more details on axum’s error handling model.

Routing to services/middleware and backpressure

Generally routing to one of multiple services and backpressure doesn’t mix well. Ideally you would want ensure a service is ready to receive a request before calling it. However, in order to know which service to call, you need the request…

One approach is to not consider the router service itself ready until all destination services are ready. That is the approach used by tower::steer::Steer.

Another approach is to always consider all services ready (always return Poll::Ready(Ok(()))) from Service::poll_ready and then actually drive readiness inside the response future returned by Service::call. This works well when your services don’t care about backpressure and are always ready anyway.

axum expects that all services used in your app wont care about backpressure and so it uses the latter strategy. However that means you should avoid routing to a service (or using a middleware) that does care about backpressure. At the very least you should [load shed] so requests are dropped quickly and don’t keep piling up.

It also means that if poll_ready returns an error then that error will be returned in the response future from call and not from poll_ready. In that case, the underlying service will not be discarded and will continue to be used for future requests. Services that expect to be discarded if poll_ready fails should not be used with axum.

One possible approach is to only apply backpressure sensitive middleware around your entire app. This is possible because axum applications are themselves services:

use axum::{
    routing::get,
    Router,
};
use tower::ServiceBuilder;

async fn handler() { /* ... */ }

let app = Router::new().route("/", get(handler));

let app = ServiceBuilder::new()
    .layer(some_backpressure_sensitive_middleware)
    .service(app);

However when applying middleware around your whole application in this way you have to take care that errors are still being handled with appropriately.

Also note that handlers created from async functions don’t care about backpressure and are always ready. So if you’re not using any Tower middleware you don’t have to worry about any of this.

Accessing state in middleware

How to make state available to middleware depends on how the middleware is written.

Accessing state in axum::middleware::from_fn

Use axum::middleware::from_fn_with_state.

Accessing state in custom tower::Layers

use axum::{
    Router,
    routing::get,
    middleware::{self, Next},
    response::Response,
    extract::State,
    http::Request,
};
use tower::{Layer, Service};
use std::task::{Context, Poll};

#[derive(Clone)]
struct AppState {}

#[derive(Clone)]
struct MyLayer {
    state: AppState,
}

impl<S> Layer<S> for MyLayer {
    type Service = MyService<S>;

    fn layer(&self, inner: S) -> Self::Service {
        MyService {
            inner,
            state: self.state.clone(),
        }
    }
}

#[derive(Clone)]
struct MyService<S> {
    inner: S,
    state: AppState,
}

impl<S, B> Service<Request<B>> for MyService<S>
where
    S: Service<Request<B>>,
{
    type Response = S::Response;
    type Error = S::Error;
    type Future = S::Future;

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.inner.poll_ready(cx)
    }

    fn call(&mut self, req: Request<B>) -> Self::Future {
        // Do something with `self.state`.
        //
        // See `axum::RequestExt` for how to run extractors directly from
        // a `Request`.

        self.inner.call(req)
    }
}

async fn handler(_: State<AppState>) {}

let state = AppState {};

let app = Router::new()
    .route("/", get(handler))
    .layer(MyLayer { state: state.clone() })
    .with_state(state);

Passing state from middleware to handlers

State can be passed from middleware to handlers using request extensions:

use axum::{
    Router,
    http::{Request, StatusCode},
    routing::get,
    response::{IntoResponse, Response},
    middleware::{self, Next},
    extract::Extension,
};

#[derive(Clone)]
struct CurrentUser { /* ... */ }

async fn auth<B>(mut req: Request<B>, next: Next<B>) -> Result<Response, StatusCode> {
    let auth_header = req.headers()
        .get(http::header::AUTHORIZATION)
        .and_then(|header| header.to_str().ok());

    let auth_header = if let Some(auth_header) = auth_header {
        auth_header
    } else {
        return Err(StatusCode::UNAUTHORIZED);
    };

    if let Some(current_user) = authorize_current_user(auth_header).await {
        // insert the current user into a request extension so the handler can
        // extract it
        req.extensions_mut().insert(current_user);
        Ok(next.run(req).await)
    } else {
        Err(StatusCode::UNAUTHORIZED)
    }
}

async fn authorize_current_user(auth_token: &str) -> Option<CurrentUser> {
    // ...
}

async fn handler(
    // extract the current user, set by the middleware
    Extension(current_user): Extension<CurrentUser>,
) {
    // ...
}

let app = Router::new()
    .route("/", get(handler))
    .route_layer(middleware::from_fn(auth));

Response extensions can also be used but note that request extensions are not automatically moved to response extensions. You need to manually do that for the extensions you need.

Rewriting request URI in middleware

Middleware added with Router::layer will run after routing. That means it cannot be used to run middleware that rewrites the request URI. By the time the middleware runs the routing is already done.

The workaround is to wrap the middleware around the entire Router (this works because Router implements Service):

use tower::Layer;
use axum::{
    Router,
    ServiceExt, // for `into_make_service`
    response::Response,
    middleware::Next,
    http::Request,
};

async fn rewrite_request_uri<B>(req: Request<B>, next: Next<B>) -> Response {
    // ...
}

// this can be any `tower::Layer`
let middleware = axum::middleware::from_fn(rewrite_request_uri);

let app = Router::new();

// apply the layer around the whole `Router`
// this way the middleware will run before `Router` receives the request
let app_with_middleware = middleware.layer(app);

axum::Server::bind(&"0.0.0.0:3000".parse().unwrap())
    .serve(app_with_middleware.into_make_service())
    .await
    .unwrap();

Modules

future
Future types.

Structs

AddExtension
Middleware for adding some shareable value to request extensions.
FromExtractor
Middleware that runs an extractor and discards the value.
FromExtractorLayer
Layer that applies FromExtractor that runs an extractor and discards the value.
FromFn
A middleware created from an async function.
FromFnLayer
A tower::Layer from an async function.
MapRequest
A middleware created from an async function that transforms a request.
MapRequestLayer
A tower::Layer from an async function that transforms a request.
MapResponse
A middleware created from an async function that transforms a response.
MapResponseLayer
A tower::Layer from an async function that transforms a response.
Next
The remainder of a middleware stack, including the handler.

Traits

IntoMapRequestResult
Trait implemented by types that can be returned from map_request, map_request_with_state.

Functions

from_extractor
Create a middleware from an extractor.
from_extractor_with_state
Create a middleware from an extractor with the given state.
from_fn
Create a middleware from an async function.
from_fn_with_state
Create a middleware from an async function with the given state.
map_request
Create a middleware from an async function that transforms a request.
map_request_with_state
Create a middleware from an async function that transforms a request, with the given state.
map_response
Create a middleware from an async function that transforms a response.
map_response_with_state
Create a middleware from an async function that transforms a response, with the given state.