中间件

在《axum 的状态共享》中，我们已经用到了中间件：用于添加共享数据的 AddExtension 中间件——它应用于AddExtensionLayer和Extension。

axum 直接使用了 tower 中间件的生态。

使用中间件

除了之前章节介绍的共享状态的中间件，我们再演示一个使用 tower 的 tower_http::trace::TraceLayer 中间件的示例。

定义处理函数

async fn foo() -> &'static str {
    "Welcome to axum.rs"
}

async fn bar() -> &'static str {
    "Powered by axum.rs"
}

定义路由并加上中间件

let app = Router::new()
        .route("/foo", get(foo))
        .route("/bar", get(bar))
        .layer(TraceLayer::new_for_http());

提示，为了能让该中间件正确地打印出日志，需要增加 tracing和tracing-subscriber 依赖，并在代码中对其进行初始化：

if std::env::var_os("RUST_LOG").is_none() {
    std::env::set_var("RUST_LOG", "tower_http=debug,middleware=debug");
}
tracing_subscriber::fmt::init();

自定义中间件

你可以通过实现 tower::Service 来自定义中间件，但它非常复杂，有兴趣可以看看官方示例。

作为示例，我们将实现一个只能使用 Firefox 浏览器访问的中间件。如果用户使用 Firefox 浏览器访问，将展示正常信息，如果是用其它浏览器访问，将显示提示信息。

定义 extractor

pub struct UserAgentInfo;

#[async_trait]
impl<B> FromRequest<B> for UserAgentInfo
where
    B: Send,
{
    type Rejection = (StatusCode, String);
    async fn from_request(req: &mut RequestParts<B>) -> Result<Self, Self::Rejection> {
        let user_agent = req
            .headers()
            .and_then(|headers| headers.get(axum::http::header::USER_AGENT))
            .and_then(|value| value.to_str().ok())
            .unwrap_or("");
        tracing::debug!("该用户UserAgent是：{:?}", user_agent);
        if !user_agent.contains("Firefox") {
            tracing::error!("非Firefox浏览器，禁止访问");
            return Err((
                StatusCode::BAD_REQUEST,
                "You MUST use Firefox to visit this page.".to_string(),
            ));
        }
        Ok(UserAgentInfo {})
    }
}

作为中间件加到路由上

.layer(extractor_middleware::<user_agent::UserAgentInfo>());

本章使用了多个第三方库，如果你对这些库暂时无法理解，可以先跳过。重点是理解如何使用及自定义中间件。