How to Implement Middleware in the Axum Framework

Axum is a Rust web framework that uses tower as its internal abstraction. Therefore, it does not have its own custom middleware system. Instead, it integrates with tower. This means that the tower ecosystem and tower-http middleware work seamlessly with axum.

While a deep understanding of tower is not necessary to write or use axum middleware, it is recommended to have at least a basic understanding of tower’s concepts. For a general introduction, refer to the tower guide. It is also recommended to read the documentation for tower::ServiceBuilder.

How to Implement Axum Middleware

In fact, axum provides multiple ways to write middleware with different levels of abstraction, each with its own advantages and disadvantages. In simple terms, the lower the abstraction level, the more complex the implementation.

There are some limitations:

Middleware written this way is only compatible with axum.

`axum::middleware::from_fn` and `axum::middleware::from_extractor_with_state`

axum::middleware::from_fn is the simplest method. It accepts a function that takes a Request and a Next, and returns a Response. Its signature looks like this:

async fn custom_middleware(
    request: Request,
    next: Next,
) -> Response {
    let response = next.run(request).await;
    response
}

In fact, it can use extractors (Extractors) for dependency injection, and the return type can be Result, for example:

async fn my_middleware(
    headers: HeaderMap,
    request: Request,
    next: Next,
) -> Result<Response, StatusCode>  {
    let response = next.run(request).await;
    Ok(response)
}

axum::middleware::from_fn_with_state functions similarly to axum::middleware::from_fn, but it includes an additional State parameter, which is convenient for injecting State to implement more functional middleware (such as authentication, logging, etc.).

The application of middleware differs as follows:

// axum::middleware::from_fn
let app = Router::new()
    .route("/", get(index))
    .layer(middleware::from_fn(custom_middleware));

// axum::middleware::from_fn_with_state
let app = Router::new()
    .route("/", get(index))
    .route_layer(middleware::from_fn_with_state(state.clone(), custom_middleware))
    .with_state(state);

Data Passing

There is a scenario where we need to retrieve some data in the middleware and use it in subsequent processes (especially in the response handling phase). For example, retrieving the user identity from the request header, we first define a struct:

#[derive(Clone)]
pub struct CurrentUser {
    pub user_id: String,
    pub email: String,
}

Then define a middleware:

pub(crate) async fn auth(
    State(ref state): State<AppState>,
    mut req: Request,
    next: Next,
) -> Result<Response, StatusCode> {
    // Retrieve the api_key from the authorization header
    let api_key = req
        .headers()
        .get(http::header::AUTHORIZATION)
        .and_then(|h| h.to_str().ok())
        .ok_or(StatusCode::UNAUTHORIZED)?
        .strip_prefix("Bearer ")
        .ok_or(StatusCode::UNAUTHORIZED)?.trim();

    let (user_id, email) = get_user_info(api_key).await?;

    let current_user = CurrentUser {
        user_id,
        email,
    };

    req.extensions_mut().insert(current_user); // current_user = 1

    Ok(next.run(req).await)
}

Using it in the handler:

pub(crate) async fn index(
    Extension(current_user): Extension<CurrentUserFromApiKey>,
    State(ref db): State<DatabaseConnection>,
    State(ref config): State<Settings>,
) -> Result<impl IntoResponse> {
    // current_user = 1

    not_implemented!()
}

`axum::middleware::from_extractor` and `axum::middleware::from_extractor_with_state`

The from_extractor series is similar to from_fn, but from_extractor is different in that you have a type that you sometimes want to use as an extractor and other times as middleware.

struct RequireAuth;

#[async_trait]
impl<S> FromRequestParts<S> for RequireAuth
where
    S: Send + Sync,
{
    type Rejection = StatusCode;

    async fn from_request_parts(parts: &mut Parts, state: &S) -> Result<Self, Self::Rejection> {
        let auth_header = parts
            .headers
            .get(header::AUTHORIZATION)
            .and_then(|value| value.to_str().ok());

        match auth_header {
            Some(auth_header) if token_is_valid(auth_header) => {
                Ok(Self)
            }
            _ => Err(StatusCode::UNAUTHORIZED),
        }
    }
}

Usage:

async fn index() {
    // todo
}
let app = Router::new()
    .route("/", get(index))
    .route_layer(from_extractor::<RequireAuth>());

axum::middleware::from_extractor_with_state functions similarly to axum::middleware::from_extractor, but it includes an additional State parameter, which is convenient for injecting State to implement more functional middleware (such as authentication, logging, etc.).

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:

However, these combinators do not meet all requirements and only support simple modifications.

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

This is a lower-level API that allows you to have maximum control over request handling. Typically, tower::Service and Pin<Box<dyn Future>> are suitable for the following scenarios:

Your middleware needs to be configurable.
You want to publish your middleware as crates for others to use.
You are uncomfortable implementing your own futures.

Here is an example using the official tower::Service template:

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

#[derive(Clone)]
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> for MyMiddleware<S>
where
    S: Service<Request, 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, request: Request) -> Self::Future {
        let future = self.inner.call(request);
        Box::pin(async move {
            let response: Response = future.await?;
            Ok(response)
        })
    }
}

Applying the middleware:

let app = Router::new()
    .route("/", get(handler))
    .layer(
        ServiceBuilder::new()
        .layer(HandleErrorLayer::new(|_: BoxError| async {
            StatusCode::BAD_REQUEST
        }))
    )

Summary

For simple requirements, it is best to use axum::middleware::from_fn and axum::middleware::from_extractor_with_state.

If you need more complex requirements, such as configurable middleware or handling futures, then using tower::Service is a better choice.


希望这能帮到您！如果有任何其他问题或需要进一步的帮助，请随时告诉我。