< 返回版块

c5soft 发表于 2022-03-02 09:28

Tags:axum

onion 中间件

axum中间件采用典型的洋葱模型,很像nodejs平台的koa,上图引自介绍koa插图,阅读本文时请忽略洋葱图每层中间件的定义。看代码 main.rs:

use axum::{
    body::{Bytes, HttpBody},
    headers::HeaderMap,
    http::{HeaderValue, Request},
    middleware::{self, Next},
    response::IntoResponse,
    routing::get,
    Router,
};
use tower::ServiceBuilder;

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

    println!("http://localhost:3000");
    axum::Server::bind(&"0.0.0.0:3000".parse().unwrap())
        .serve(app.into_make_service())
        .await
        .unwrap();
}

async fn handler(headers: HeaderMap) -> impl IntoResponse {
    let x_mid = get_header_value_or_empty(&headers, "x-mid");
    println!("Here handler called, x_mid={}", x_mid);
    "Hello World"
}

async fn mid_handler_a<B>(req: Request<B>, next: Next<B>) -> impl IntoResponse
where
    B: std::fmt::Debug,
{
    println!("");
    mid_handler(req, next, "A").await
}

async fn mid_handler_b<B>(req: Request<B>, next: Next<B>) -> impl IntoResponse
where
    B: std::fmt::Debug,
{
    mid_handler(req, next, "B").await
}

async fn mid_handler<B>(mut req: Request<B>, next: Next<B>, handler_name: &str) -> impl IntoResponse
where
    B: std::fmt::Debug,
{
    let uri = format!("{:?}", req.uri());
    let headers = req.headers_mut();
    let old_value = get_header_value_or_empty(headers, "x-mid");
    let new_value = old_value.clone() + &handler_name;
    if !old_value.is_empty() {
        headers.remove("x-mid");
    };
    headers.append("x-mid", HeaderValue::from_str(&new_value).unwrap());

    println!(
        "{} before{}, x-mid: {}=>{}",
        uri, handler_name, old_value, new_value
    );

    let response = next.run(req).await;

    let body_text = body_into_text(response.into_body()).await;
    let new_body = handler_name.to_owned() + "-" + &body_text + "-" + handler_name;
    println!(
        "{} after{}, response body:{}=>{}",
        uri, handler_name, body_text, new_body
    );
    new_body
}

async fn body_into_text<B>(body: B) -> String
where
    B: HttpBody<Data = Bytes> + Unpin,
    B::Error: std::fmt::Debug,
{
    let bytes = hyper::body::to_bytes(body).await.unwrap();
    String::from_utf8(bytes.to_vec()).unwrap()
}

fn get_header_value_or_empty(headers: &HeaderMap, key: &str) -> String {
    headers
        .get(key)
        .map(|x| x.to_owned())
        .unwrap_or(HeaderValue::from_static(""))
        .to_str()
        .unwrap()
        .to_owned()
}

程序处理流程:

      request
         |
         v
 +---- mid_a ----+ 
 | +-- mid_b --+ | 
 | |           | | 
 | |  handler  | | 
 | |           | | 
 | +-- mid_b --+ | 
 +---- mid_a ----+ 
         |
         v
      response

库依赖 Cargo.toml:

[dependencies]
axum = { version = "0.4.7", features = ["http2", "headers"] }
tokio = { version = "1.17.0", features = ["macros", "rt", "rt-multi-thread"] }
tower = { version = "0.4.12" }
hyper = { version = "0.14.17" }

评论区

写评论

还没有评论

1 共 0 条评论, 1 页