Skip to main content

Axum的middleware模块

什么是middleware模块

用于编写中间件的实用程序

axum 在于其没有自己独特的专门的中间件系统,而是与 tower 集成。

这意味着 towertower-http 中间件的生态系统都与 axum 配合使用。

虽然不必完全了解 tower,才能编写或使用 axum 的中间件, 但建议至少对 tower 的概念有基本了解。

请查看 tower 的指南以获取一般介绍。 同时建议阅读 tower::ServiceBuilder 的文档。

应用中间件

axum允许您几乎可以在任何地方添加中间件

  • 使用 Router::layerRouter::route_layer 对整个路由器进行中间件处理。
  • 通过 MethodRouter::layerMethodRouter::route_layer 添加到方法路由器。
  • 对特定的处理程序使用 Handler::layer

应用多个中间件

建议一次使用 tower::ServiceBuilder 来应用多个中间件,而不是重复调用 layer(或 route_layer):

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 {}))
);

常用中间件

一些常用的中间件包括:

  • 用于高级跟踪/日志记录的TraceLayer
  • 用于处理跨源资源共享 (CORS) 的CorsLayer
  • 自动压缩响应的CompressionLayer
  • RequestIdLayerPropagateRequestIdLayer 设置并传播请求标识。
  • 对于超时的TimeoutLayer

顺序

当您使用 Router::layer(或类似方法)添加中间件时,所有先前添加的路由都将被包裹在中间件中。

一般来说,这导致中间件从底部到顶部执行

因此,如果你这样做:

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

async fn handler() {}

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

将中间件想象成像洋葱一样分层,每一层都包裹着前面所有的层:

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

也就是说:

  • 首先,layer_three 接收请求
  • 然后执行其操作,并将请求传递到 layer_two
  • 然后传递请求到 layer_one
  • 然后传递请求到处理程序生成响应,该响应然后传递给 layer_one
  • 然后到 layer_two
  • 最后传递给 layer_three,最终从您的应用程序返回

实际上更复杂一点,因为任何中间件都可以自由地提前返回并不调用下一层, 例如,如果请求无法授权,但这是一个有用的心智模型。

如前所述,建议使用tower::ServiceBuilder 添加多个中间件,但这会影响顺序

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工作原理将所有层组合在一起,使它们从上到下依次运行

因此,使用上述代码,layer_one 首先会接收请求,然后是 layer_two, 然后是 layer_three,然后是 handler, 最后响应会通过 layer_three,然后 layer_two,最后是 layer_one

从上到下执行中间件通常更容易理解和跟踪,这也是为什么建议使用ServiceBuilder的原因。

编写中间件

axum提供了许多编写中间件的方式,不同的抽象级别和不同的优缺点

axum::middleware::from_fn

在需要时,使用 axum::middleware::from_fn 来编写中间件

  • 您不熟悉实现自己的futures,并更倾向于使用熟悉的async/await语法。
  • 您不打算将自己的中间件作为crate发布供他人使用。像这样编写的中间件只与axum兼容。

axum::middleware::from_extractor

使用axum::middleware::from_extractor来编写中间件时:

  • 您有一种类型,有时希望将其用作提取器,有时希望将其用作中间件。 如果您只需要将您的类型作为中间件,请优先选择 middleware::from_fn

tower的组合子

tower 有几个实用的组合器,可用于对请求或响应进行简单的修改。

其中最常用的是:

  • ServiceBuilder::map_request
  • ServiceBuilder::map_response
  • ServiceBuilder::then
  • ServiceBuilder::and_then

在这种情况下应该使用这些:

  • 您想执行一个小的临时操作,比如添加一个标头(header)。
  • 您不打算将您的中间件作为一个包发布供他人使用。

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

为了获得最大控制权(以及更低级别的API), 您可以通过实现 tower::Service 来编写自己的中间件:

在需要使用 tower::ServicePin<Box<dyn Future>> 编写中间件时。

  • 您的中间件需要是可配置的,例如通过 tower::Layer 上的构建方法,比如 tower_http::trace::TraceLayer
  • 您确实打算将您的中间件发布为其他人使用的 crate。
  • 您不太熟悉实现自己的 futures

一种合适的模板可以是:

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)
})
}
}

请注意,将你的错误类型定义为 S::Error 意味着你的中间件通常不返回错误。 作为原则,请始终尝试返回一个响应,不要用自定义错误类型退出。

例如,如果你在新中间件中使用的第三方库返回其自己专门的错误类型, 尝试将其转换为某种合理的响应,并以该响应返回 Ok

如果选择实现自定义的错误类型, 比如 type Error = BoxError(一个封装的不透明错误), 或者任何不是 Infallible 的其他错误类型, 您必须使用 HandleErrorLayer,在这里是使用 ServiceBuilder 的示例:

warning
ServiceBuilder::new()
.layer(HandleErrorLayer::new(|_: BoxError| async {
// because Axum uses infallible errors, you must handle your custom error type from your middleware here
StatusCode::BAD_REQUEST
}))
.layer(
// <your actual layer which DOES return an error>
);

tower::Service和自定义futures

如果您习惯实现自己的 futures(或想要学习它)并且需要尽可能多的控制, 则使用 tower::Service 而不是futures 是最佳选择。

使用tower::Service和手动futures编写中间件时:

  • 您希望您的中间件拥有尽可能低的开销。
  • 您的中间件需要通过 tower::Layer 上的构建器方法(例如 tower_http::trace::TraceLayer)进行配置。
  • 您确实打算将您的中间件作为一个 crate 发布供他人使用,可能作为 tower-http 的一部分。
  • 您乐意实现自己的 futures,或者想了解异步 Rust 的底层工作原理。

tower从头开始构建中间件 指南是学习如何做到这一点的好地方。

中间件的错误处理

axum的错误处理模型要求处理程序始终返回响应。 但是,中间件是将错误引入应用程序的一种可能方式。

如果 hyper 收到错误,则连接将在不发送响应的情况下关闭。 因此,axum 要求对这些错误进行优雅处理:

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()
// 这个中间件位于 `TimeoutLayer` 之上,因为它将接收 `TimeoutLayer` 返回的错误。
.layer(HandleErrorLayer::new(|_: BoxError| async {
StatusCode::REQUEST_TIMEOUT
}))
.layer(TimeoutLayer::new(Duration::from_secs(10)))
);

有关axum错误处理模型的更多详细信息,请参阅error_handling

路由到服务/中间件和背压

通常,将路由到多个服务和背压不太相容。理想情况下,您希望在调用服务之前确保服务已准备好接收请求。 然而,为了知道要调用哪个服务,您需要该请求…

其中一种方法是在所有目标服务准备就绪之前,不认为路由器服务本身已准备就绪。 这是 tower::steer::Steer 使用的方法。

另一种方法是始终考虑所有服务都已准备就绪(始终从 Service::poll_ready 返回 Poll::Ready(Ok(())) ), 然后实际上在 Service::call 返回的响应未来中驱动准备就绪。

当您的服务不关心背压并且始终保持准备就绪时,这种方法非常有效。

axum 希望在您的应用程序中使用的所有服务都不关心后压力,因此它使用后一种策略。 但这意味着您应该避免路由到关心后压力的服务(或使用中间件)。

至少,您应该卸载,以便快速丢弃请求,不要继续堆积。

这也意味着,如果 poll_ready 返回错误,则该错误将在 call 的响应未来中返回, 而不是在 poll_ready 中返回。

在这种情况下,底层服务将不会被丢弃,并将继续用于未来的请求。 如果服务期望在 poll_ready 失败时被丢弃,则不应与 axum 一起使用。

一个可能的方法是只在整个应用程序周围应用对背压敏感的中间件。

这是可能的,因为axum应用程序本身就是服务:

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);

然而,当以这种方式应用中间件到整个应用程序时,您必须小心确保错误仍然被适当处理。

还要注意,从异步函数创建的处理程序不关心背压,并且始终准备就绪。

因此,如果您不使用任何 Tower 中间件,您不必担心这些问题

在中间件中访问状态

如何使状态在中间件中可用取决于中间件的编写方式。

axum::middleware::from_fn 中访问状态

使用axum::middleware::from_fn_with_state

在自定义 tower::Layer 中访问状态

use axum::{
Router,
routing::get,
middleware::{self, Next},
response::Response,
extract::{State, 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);

从中间件传递状态到处理程序

中间件可以使用请求扩展 将状态从中间件传递给处理程序。

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

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

async fn auth(mut req: Request, next: Next) -> 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));

还可以使用响应扩展, 但请注意,请求扩展不会自动转移到响应扩展。 您需要为所需的扩展手动执行此操作。

在中间件中重写请求URI

使用 Router::layer 添加的中间件将在路由之后运行

这意味着它不能用于运行重写请求 URI 的中间件。在中间件运行时,路由已经完成

解决方法将中间件包装在整个路由器周围(这是有效的因为路由器实现了 Service ):

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

fn rewrite_request_uri<B>(req: Request<B>) -> Request<B> {
// ...
}

// 这可以是任何 `tower::Layer`
let middleware = tower::util::MapRequestLayer::new(rewrite_request_uri);

let app = Router::new();

// 将层应用于整个 `Router`。
// 这样中间件将在 `Router` 接收请求之前运行。
let app_with_middleware = middleware.layer(app);

let listener = tokio::net::TcpListener::bind("0.0.0.0:3000").await.unwrap();
axum::serve(listener, app_with_middleware.into_make_service()).await.unwrap();

模块

  • future: Future types

结构体

  • AddExtension: 用于向请求扩展添加一些可共享值的中间件。
  • FromExtractor: 运行提取器并丢弃值的中间件。
  • FromExtractorLayer: 应用运行提取器并丢弃值的FromExtractor的层。
  • FromFn: 从异步函数创建的中间件。
  • FromFnLayer: 来自异步函数的tower::Layer
  • MapRequest: 从转换请求的异步函数创建的中间件。
  • MapRequestLayer: 来自转换请求的异步函数的tower::Layer
  • MapResponse: 从转换响应的异步函数创建的中间件。
  • MapResponseLayer: 来自转换响应的异步函数的tower::Layer
  • Next: 包括处理程序在内的中间件堆栈的剩余部分。

Traits

  • IntoMapRequestResult: 可从map_requestmap_request_with_state返回的类型实现的Trait。

函数

  • from_extractor: 从提取器创建一个中间件
  • from_extractor_with_state: 从具有给定状态的提取器创建一个中间件
  • from_fn: 从异步函数创建一个中间件
  • from_fn_with_state: 从具有给定状态的异步函数创建一个中间件
  • map_request: 从转换请求的异步函数创建一个中间件
  • map_request_with_state: 从转换请求与给定状态的异步函数创建一个中间件
  • map_response: 从转换响应的异步函数创建一个中间件
  • ap_response_with_state: 从转换响应与给定状态的异步函数创建一个中间件