Axum中间件如何获取并修改Extensions中的ServiceConfig对象

在Rust Axum应用中，创建Router时注入了Extension<ServiceConfig>，但jwt_validator中间件始终无法获取该对象，会进入错误分支返回500状态码；而后续请求处理handler却能正常读取这个配置。需要实现中间件对ServiceConfig的获取、修改，并传递给后续handler。

相关代码 #

路由创建代码 #

pub fn create_routes(sc: &ServiceConfig) -> Router {
    let api_routes = Router::new().nest("/", create_workflow_router());

    let router = Router::new()
        .route("/", get(welcome))
        .nest("/api", api_routes)
        .layer(Extension(sc.clone()))
        .layer(from_fn(jwt_validator))
        .layer(ServiceBuilder::new().layer(TraceLayer::new_for_http()));

    router
}

中间件代码 #

pub async fn jwt_validator(
    req: Request<axum::body::Body>,
    next: Next,
) -> Result<Response, StatusCode> {
      // 尝试从请求扩展中获取ServiceConfig
    let sc = if let Some(Extension(sc)) = req.extensions().get::<Extension<ServiceConfig>>() {
        sc.clone()
    } else {
        trace!("ServiceConfig not found in extensions.");
        return Err(StatusCode::INTERNAL_SERVER_ERROR);
    };

    println!("SC: {:?}", sc);

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

请求处理handler代码 #

pub async fn get_workflows(
    Query(params): Query<WorkflowQueryParams>,
    Extension(sc): Extension<ServiceConfig>,
) -> impl IntoResponse {
    println!("SC: {:?}", sc);
    Json("Hello")
}

运行日志 #

错误日志（中间件触发500） #

[2024-10-01T06:02:35Z TRACE axum::serve] connection 127.0.0.1:34734 accepted
[2024-10-01T06:02:35Z DEBUG tower_http::trace::make_span] request; method=GET uri=/api/v1/workflows version=HTTP/1.1
[2024-10-01T06:02:35Z TRACE tracing::span::active] -> request;
[2024-10-01T06:02:35Z DEBUG tower_http::trace::on_request] started processing request
[2024-10-01T06:02:35Z TRACE tracing::span::active] <- request;
[2024-10-01T06:02:35Z TRACE tracing::span::active] -> request;
[2024-10-01T06:02:35Z TRACE workflow_svc::middlewares::jwt_validator] ServiceConfig not found in extensions.
[2024-10-01T06:02:35Z DEBUG tower_http::trace::on_response] finished processing request latency=0 ms status=500
[2024-10-01T06:02:35Z ERROR tower_http::trace::on_failure] response failed classification=Status code: 500 Internal Server Error latency=0 ms
[2024-10-01T06:02:35Z TRACE tracing::span::active] <- request;
[2024-10-01T06:02:35Z TRACE tracing::span] -- request;
[2024-10-01T06:02:35Z TRACE tracing::span] -- request;

正常日志（注释中间件错误分支后） #

[2024-10-01T06:09:04Z TRACE axum::serve] connection 127.0.0.1:38964 accepted
[2024-10-01T06:09:04Z DEBUG tower_http::trace::make_span] request; method=GET uri=/api/v1/workflows version=HTTP/1.1
[2024-10-01T06:09:04Z TRACE tracing::span::active] -> request;
[2024-10-01T06:09:04Z DEBUG tower_http::trace::on_request] started processing request
[2024-10-01T06:09:04Z TRACE tracing::span::active] <- request;
[2024-10-01T06:09:04Z TRACE tracing::span::active] -> request;
SC: ServiceConfig { auth_enabled: false, jwt_secret: None, host: "localhost", port: 4001, dbc: SqlxPostgresPoolConnection }
[2024-10-01T06:09:04Z DEBUG tower_http::trace::on_response] finished processing request latency=0 ms status=200
[2024-10-01T06:09:04Z TRACE tracing::span::active] <- request;
[2024-10-01T06:09:04Z TRACE tracing::span] -- request;
[2024-10-01T06:09:04Z TRACE tracing::span::active] -> request;
[2024-10-01T06:09:04Z TRACE tracing::span::active] <- request;
[2024-10-01T06:09:04Z TRACE tracing::span] -- request;
[2024-10-01T06:09:04Z TRACE axum::serve] connection 127.0.0.1:38964 closed

1. 修复Layer执行顺序问题 #

Axum的Layer是栈式执行，后添加的Layer会先处理请求。你当前的代码中，Extension Layer先于中间件添加，导致中间件执行时，配置还未被注入到请求扩展中。

调整Layer顺序，将Extension放在中间件之后添加，确保配置先被注入，再执行中间件：

pub fn create_routes(sc: &ServiceConfig) -> Router {
    let api_routes = Router::new().nest("/", create_workflow_router());

    let router = Router::new()
        .route("/", get(welcome))
        .nest("/api", api_routes)
        .layer(from_fn(jwt_validator)) // 先添加中间件
        .layer(Extension(sc.clone()))  // 后添加配置扩展，确保先注入配置
        .layer(ServiceBuilder::new().layer(TraceLayer::new_for_http()));

    router
}

2. 修正中间件中配置的获取方式 #

请求扩展中存储的是ServiceConfig实例，而非Extension<ServiceConfig>包装器，直接获取对应类型即可：

pub async fn jwt_validator(
    mut req: Request<axum::body::Body>,
    next: Next,
) -> Result<Response, StatusCode> {
    // 正确获取ServiceConfig
    let mut sc = req.extensions()
        .get::<ServiceConfig>()
        .cloned()
        .ok_or_else(|| {
            trace!("ServiceConfig not found in extensions.");
            StatusCode::INTERNAL_SERVER_ERROR
        })?;

    println!("Original SC: {:?}", sc);
    Ok(next.run(req).await)
}

3. 修改配置并传递给后续handler #

如果需要修改配置，可移除原有实例并插入修改后的版本，后续handler会读取到更新后的配置：

pub async fn jwt_validator(
    mut req: Request<axum::body::Body>,
    next: Next,
) -> Result<Response, StatusCode> {
    // 获取并克隆原配置
    let mut sc = req.extensions()
        .get::<ServiceConfig>()
        .cloned()
        .ok_or(StatusCode::INTERNAL_SERVER_ERROR)?;

    // 修改配置示例
    sc.auth_enabled = true;
    sc.jwt_secret = Some("updated_secret".to_string());

    // 替换请求扩展中的配置
    req.extensions().remove::<ServiceConfig>();
    req.extensions().insert(sc);

    println!("Modified SC: {:?}", req.extensions().get::<ServiceConfig>().unwrap());
    Ok(next.run(req).await)
}

内容的提问来源于stack exchange，提问作者Mopparthy Ravindranath