如何在Rust Axum自定义ConnectInfo实现中获取User Agent？

问题背景 #

我用Axum的into_make_service_with_connect_info方法获取远程IP，自定义了MyConnectionInfo结构体，能通过target.remote_addr()拿到IP，但User Agent只能硬编码。想知道能不能在Connected实现里获取User Agent等请求头？还是我的实现方式错了？

自定义实现代码：

#[derive(Clone, Debug)]
struct MyConnectionInfo {
    ip: String,
    user_agent: String,
}

impl Connected<IncomingStream<'_, TcpListener>> for MyConnectionInfo {
    fn connect_info(target: IncomingStream<'_, TcpListener>) -> Self {
        MyConnectionInfo {
            ip: target.remote_addr().to_string(), user_agent: "Firefox".to_string()
        }
    }
}

核心结论 #

不能在Connected trait的实现中获取请求头，原因很明确：

• Connected是在TCP连接建立阶段被调用的，此时Axum还未解析HTTP请求，仅能拿到底层TCP流的连接信息（比如远程IP）。
• 请求头属于HTTP请求层面的内容，只有当Axum开始处理HTTP请求时，才能被解析和获取。

正确实现方式 #

下面提供几种常见的可行方案：

方案1：在Handler中直接提取多个信息 #

直接在handler参数中同时声明ConnectInfo和TypedHeader<UserAgent>，Axum会自动帮你提取这两个信息：

use axum::{
    extract::{ConnectInfo, TypedHeader},
    routing::get,
    Router,
};
use std::net::SocketAddr;
use axum::headers::UserAgent;

async fn handler(
    ConnectInfo(addr): ConnectInfo<SocketAddr>,
    TypedHeader(user_agent): TypedHeader<UserAgent>,
) -> String {
    format!("IP: {}, User Agent: {}", addr, user_agent.as_str())
}

#[tokio::main]
async fn main() {
    let app = Router::new().route("/", get(handler));

    axum::Server::bind(&"0.0.0.0:3000".parse().unwrap())
        .serve(app.into_make_service_with_connect_info::<SocketAddr>())
        .await
        .unwrap();
}

方案2：自定义组合提取器 #

如果需要在多个handler中复用这些信息，可以自定义一个提取器，把IP和User Agent组合成一个结构体：

use axum::{
    extract::{ConnectInfo, FromRequest, TypedHeader},
    headers::UserAgent,
    http::Request,
};
use std::net::SocketAddr;

#[derive(Debug, Clone)]
struct RequestContext {
    ip: String,
    user_agent: String,
}

#[axum::async_trait]
impl<S> FromRequest<S> for RequestContext
where
    S: Send + Sync,
    ConnectInfo<SocketAddr>: FromRequest<S>,
    TypedHeader<UserAgent>: FromRequest<S>,
{
    type Rejection = axum::extract::rejection::TypedHeaderRejection;

    async fn from_request(req: Request<axum::body::Body>, state: &S) -> Result<Self, Self::Rejection> {
        let ConnectInfo(addr) = ConnectInfo::<SocketAddr>::from_request(req.clone(), state).await.unwrap();
        let TypedHeader(user_agent) = TypedHeader::<UserAgent>::from_request(req, state).await?;

        Ok(Self {
            ip: addr.to_string(),
            user_agent: user_agent.to_string(),
        })
    }
}

// 使用示例
async fn handler(context: RequestContext) -> String {
    format!("IP: {}, User Agent: {}", context.ip, context.user_agent)
}

方案3：通过中间件存入请求扩展 #

如果需要在中间件或多个handler中共享这些信息，可以用中间件提取后存入请求的Extensions：

use axum::{
    extract::{ConnectInfo, TypedHeader},
    headers::UserAgent,
    http::Request,
    middleware::Next,
    response::Response,
};
use std::net::SocketAddr;

#[derive(Debug, Clone)]
struct RequestContext {
    ip: String,
    user_agent: String,
}

async fn request_context_middleware<B>(
    mut req: Request<B>,
    next: Next<B>,
) -> Result<Response, axum::extract::rejection::TypedHeaderRejection> {
    // 从扩展中获取已由into_make_service_with_connect_info注入的ConnectInfo
    let addr = req.extensions().get::<ConnectInfo<SocketAddr>>().unwrap().0;
    // 提取User Agent
    let TypedHeader(user_agent) = TypedHeader::<UserAgent>::from_request(req.clone(), &()).await?;

    // 将组合后的信息存入扩展
    req.extensions_mut().insert(RequestContext {
        ip: addr.to_string(),
        user_agent: user_agent.to_string(),
    });

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

// 使用示例
async fn handler(context: axum::extract::Extension<RequestContext>) -> String {
    format!("IP: {}, User Agent: {}", context.ip, context.user_agent)
}

#[tokio::main]
async fn main() {
    let app = Router::new()
        .route("/", get(handler))
        .layer(axum::middleware::from_fn(request_context_middleware));

    axum::Server::bind(&"0.0.0.0:3000".parse().unwrap())
        .serve(app.into_make_service_with_connect_info::<SocketAddr>())
        .await
        .unwrap();
}

总结 #

Connected trait仅负责获取连接层元数据，无法触及HTTP请求内容。要获取User Agent等请求头，必须在HTTP请求处理阶段通过提取器或中间件实现。

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