如何在Rust Axum嵌套路由中访问父路由传递的pool状态？

我在Axum项目中，父路由通过with_state传递了数据库连接池pool给子路由，但在create_auth_router的/profile路由中间件中，不知道如何访问这个pool状态。

主路由代码：

let app = Router::new()
    .nest("/products", create_products_router())
    .nest("/auth", create_auth_router())
    .with_state(pool)
    .layer(
        CorsLayer::new()
            .allow_origin("http://localhost:4200".parse::<HeaderValue>().unwrap())
            .allow_methods([Method::GET, Method::POST, Method::DELETE, Method::PUT])
            .allow_headers(vec!["content-type".parse().unwrap()]),
    );

子路由create_auth_router代码：

pub fn create_auth_router() -> Router<PgPool> {
    Router::new()
        .route("/login", post(login))
        .route("/register", post(register))
        .route(
            "/profile",
            get(profile).layer(middleware::from_fn_with_state(
                How to access state "pool" here!!!,
                authorization_middleware,
            )),
        )
}

方案1：利用Axum的状态提取器（推荐） #

Axum中父路由设置的状态会自动传递给子路由，无需手动在from_fn_with_state中传递状态，直接使用middleware::from_fn，然后在中间件函数中通过State提取器获取PgPool：

1. 修改子路由的中间件配置：

pub fn create_auth_router() -> Router<PgPool> {
    Router::new()
        .route("/login", post(login))
        .route("/register", post(register))
        .route(
            "/profile",
            get(profile).layer(middleware::from_fn(authorization_middleware)),
        )
}

2. 调整中间件函数定义，添加State<PgPool>作为参数：

use axum::{extract::State, http::{Request, Response, StatusCode}, middleware::Next, body::Body};
use sqlx::PgPool;

async fn authorization_middleware(
    req: Request<Body>,
    next: Next<Body>,
    State(pool): State<PgPool>, // 提取父路由传递的Pool状态
) -> Result<Response<Body>, StatusCode> {
    // 在这里使用pool执行数据库操作，比如验证用户权限
    // 示例：let user = sqlx::query!("SELECT * FROM users WHERE id = $1", user_id).fetch_one(&pool).await?;
    
    Ok(next.run(req).await)
}

方案2：手动传递Pool给子路由 #

如果倾向于使用from_fn_with_state，可以修改create_auth_router的签名，让它接收pool作为参数再传递给中间件：

1. 修改子路由函数：

pub fn create_auth_router(pool: PgPool) -> Router<PgPool> {
    Router::new()
        .route("/login", post(login))
        .route("/register", post(register))
        .route(
            "/profile",
            get(profile).layer(middleware::from_fn_with_state(
                pool.clone(), // 传递Pool的克隆
                authorization_middleware,
            )),
        )
        .with_state(pool) // 确保子路由其他handler也能获取状态
}

2. 在主路由中调用时传递Pool：

let app = Router::new()
    .nest("/products", create_products_router())
    .nest("/auth", create_auth_router(pool.clone())) // 传递克隆后的Pool
    .with_state(pool)
    .layer(/* CORS配置 */);

3. 中间件函数接收传递的状态：

async fn authorization_middleware(
    State(pool): State<PgPool>,
    req: Request<Body>,
    next: Next<Body>,
) -> Result<Response<Body>, StatusCode> {
    // 使用pool进行操作
    Ok(next.run(req).await)
}

内容的提问来源于stack exchange，提问作者Kadiem Alqazzaz