[网络协议]Rust axum—HTTP服务端

1. Hello, axum!

于 Cargo.toml 文件中添加依赖：

[dependencies]
axum = { version = "0.2" }
tokio = { version = "1", features = ["full"] }
tower = { version = "0.4" }
futures = { version = "0.3" }
serde = { version = "1.0", features = ["derive"] }
serde_json = { version = "1.0" }

use axum::{
    handler::get,
    Router,
    response::IntoResponse,
};


#[tokio::main]
async fn main() {
    let app = Router::new()
        .route("/", get(index));
        
    axum::Server::bind(&"127.0.0.1:8888".parse().unwrap())
        .serve(app.into_make_service())
        .await
        .unwrap();
}

async fn index() -> impl IntoResponse {
    "Hello, axum!"
}

2. 嵌套路由

通过调用 Router::nest() 来实现，可以实现路由分组的功能。

use axum::{
    handler::get,
    Router,
    response::IntoResponse,
    routing::BoxRoute
};


#[tokio::main]
async fn main() {
    let app = Router::new()
        .nest("/user", user_routes());

    axum::Server::bind(&"127.0.0.1:8888".parse().unwrap())
        .serve(app.into_make_service())
        .await
        .unwrap();
}

fn user_routes() -> Router<BoxRoute> {
    Router::new()
        .route("/list", get(list_user))
        .route("/get", get(get_user))
        .boxed()
}

async fn list_user() -> impl IntoResponse {
    r#"
    "Mike"
    "Jane"
    "#
}

async fn get_user() -> impl IntoResponse {
    "Mike"
}

对应的 URL：

• http://127.0.0.1:8888/user/list
• http://127.0.0.1:8888/user/get

3. 提取器

“提取器”可以从请求中提取出请求处理器所需的参数。

Json

use axum::{
    handler::post,
    extract::Json,
    Router,
    response::IntoResponse,
    http::StatusCode,
};
use serde::{Serialize, Deserialize};

#[derive(Serialize, Deserialize)]
struct User {
    email: String,
    password: String,
}

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

    axum::Server::bind(&"127.0.0.1:8888".parse().unwrap())
        .serve(app.into_make_service())
        .await
        .unwrap();
}

async fn create_user(Json(payload): Json<User>) -> impl IntoResponse {
    let user = User {
        email: payload.email,
        password: payload.password,
    };
    (StatusCode::CREATED, Json(user))
}

表单

use axum::{
    handler::post,
    extract::{Form, Json},
    Router,
    response::IntoResponse,
    http::StatusCode,
};
use serde::{Serialize, Deserialize};

#[derive(Serialize, Deserialize)]
struct User {
    email: String,
    password: String,
}

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

    axum::Server::bind(&"127.0.0.1:8888".parse().unwrap())
        .serve(app.into_make_service())
        .await
        .unwrap();
}

async fn create_user(Form(payload): Form<User>) -> impl IntoResponse {
    let user = User {
        email: payload.email,
        password: payload.password,
    };
    (StatusCode::CREATED, Json(user))
}

路径参数

use axum::{
    handler::get,
    Router,
    response::IntoResponse,
    extract::Path,
};
use std::collections::HashMap;

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

    axum::Server::bind(&"127.0.0.1:8888".parse().unwrap())
        .serve(app.into_make_service())
        .await
        .unwrap();
}

async fn get_book(Path(params): Path<HashMap<String, String>>) -> impl IntoResponse {
    let category = params.get("category").unwrap();
    let name = params.get("name").unwrap();

    format!("category: {}, name: {}", category, name)
}

如，请求 http://127.0.0.1:8888/book/programming/python

查询参数

use axum::{
    handler::get,
    Router,
    response::IntoResponse,
    extract::Query,
};
use std::collections::HashMap;

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

    axum::Server::bind(&"127.0.0.1:8888".parse().unwrap())
        .serve(app.into_make_service())
        .await
        .unwrap();
}

async fn get_book(Query(query): Query<HashMap<String, String>>) -> impl IntoResponse {
    let category = query.get("category").unwrap();
    let name = query.get("name").unwrap();

    format!("category: {}, name: {}", category, name)
}

如，请求 http://127.0.0.1:8888/book?category=programming&name=python

请求头

use axum::{
    handler::get,
    Router,
    response::IntoResponse,
    http::header::HeaderMap,
};

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

    axum::Server::bind(&"127.0.0.1:8888".parse().unwrap())
        .serve(app.into_make_service())
        .await
        .unwrap();
}

async fn get_user_agent(headers: HeaderMap) -> impl IntoResponse {
    let agent = headers.get("User-Agent").unwrap().to_str().unwrap();
    format!("agent: {}", agent)
}

4. 响应

请求处理器的返回值只需实现 IntoResponse trait 即可。

此处演示如何设置响应状态码、响应头、响应体。

use axum::{
    handler::get,
    Router,
    response::IntoResponse,
    http::{StatusCode, header::{HeaderMap, HeaderName, HeaderValue}},
};


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

    axum::Server::bind(&"127.0.0.1:8888".parse().unwrap())
        .serve(app.into_make_service())
        .await
        .unwrap();
}

async fn hello() -> impl IntoResponse {
    let mut headers = HeaderMap::new();
    headers.insert(
        HeaderName::from_static("token"),
        HeaderValue::from_static("123456")
    );

    (StatusCode::OK, headers, "hello, axum!")
}

5. 共享状态

use axum::{
    handler::get,
    Router,
    response::IntoResponse,
    AddExtensionLayer,
    extract,
};
use std::sync::{Arc, atomic::{AtomicU32, Ordering}};

struct State {
    count: AtomicU32,
}

#[tokio::main]
async fn main() {
    let shared_state = Arc::new(State{count: AtomicU32::new(0)});
    
    let app = Router::new()
        .route("/", get(hello))
        .layer(AddExtensionLayer::new(shared_state));

    axum::Server::bind(&"127.0.0.1:8888".parse().unwrap())
        .serve(app.into_make_service())
        .await
        .unwrap();
}

async fn hello(state: extract::Extension<Arc<State>>) -> impl IntoResponse {
    let state: Arc<State> = state.0;
    let response = format!("count: {}", state.count.load(Ordering::Relaxed));
    state.count.fetch_add(1, Ordering::Relaxed);
    response
}

6. 中间件

判断请求头中是否存在 token。

use axum::{
    handler::get,
    Router,
    response::IntoResponse,
    http::{Request, Response},
};
use futures::future::BoxFuture;
use tower::{Service, layer::layer_fn};
use std::task::{Context, Poll};
use std::fmt;
use std::error::Error;


#[derive(Clone)]
struct TokenMiddleware<S, T> {
    inner: S,
    fallback: T,
}

#[derive(Debug)]
struct TokenError;

impl fmt::Display for TokenError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "token error")
    }
}

impl Error for TokenError {}

impl<S, T, ReqBody, ResBody> Service<Request<ReqBody>> for TokenMiddleware<S, T>
    where
        S: Service<Request<ReqBody>, Response = Response<ResBody>> + Clone + Send + 'static,
        S::Future: Send + 'static,
        S::Error: Into<Box<dyn Error + Send + Sync>> + 'static,
        T: Service<Request<ReqBody>, Response = Response<ResBody>> + Clone + Send + 'static,
        T::Future: Send + 'static,
        T::Error: Into<Box<dyn Error + Send + Sync>> + 'static,
        ReqBody: Send + 'static,
        ResBody: Send + 'static,
{
    type Response = S::Response;
    type Error = Box<dyn Error + Send + Sync>;
    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).map_err(Into::into)
    }

    fn call(&mut self, mut req: Request<ReqBody>) -> Self::Future {
        let clone = self.inner.clone();
        let mut inner = std::mem::replace(&mut self.inner, clone);

        let clone = self.fallback.clone();
        let mut fallback = std::mem::replace(&mut self.fallback, clone);

        let future = async move {
            let token = req.headers().get("token");

            if let Some(value) = token {
                if let Ok(s) = value.to_str() {
                    if s == "123456" {
                        let res = inner.call(req).await;
                        if let Ok(res) = res {
                            println!("`MyMiddleware` received the response");
                            return Ok(res);
                        }

                        return Err(Box::new(TokenError) as Box<dyn Error + Send + Sync>);
                    }
                }
            }

            let res = fallback.call(req).await;
            if let Ok(res) = res {
                return Ok(res);
            }

            return Err(Box::new(TokenError) as Box<dyn Error + Send + Sync>);
        };

        Box::pin(future)
    }
}

#[tokio::main]
async fn main() {
    let app = Router::new()
        .route("/", get(hello))
        .layer(layer_fn(|inner| TokenMiddleware { inner, fallback: get(fallback) }));

    axum::Server::bind(&"127.0.0.1:8888".parse().unwrap())
        .serve(app.into_make_service())
        .await
        .unwrap();
}

async fn hello() -> impl IntoResponse {
    "hello, axum!"
}

async fn fallback() -> impl IntoResponse {
    "error!"
}

• 为单独的请求处理器添加中间件：

  let app = Router::new()
      .route(
          "/",
          get(handler.layer(ConcurrencyLimitLayer::new(100))),
      );
  

• 为一组请求处理器添加中间件：

  let app = Router::new()
      .route("/", get(handler))
      .route("/foo", post(handler))
      .layer(ConcurrencyLimitLayer::new(100));
  

  或，

  let foo = Router::new()
      .route("/", get(handler))
      .route("/foo", post(handler))
      .layer(ConcurrencyLimitLayer::new(100));
  
  let bar = Router::new()
      .route("/requires-auth", get(handler))
      .layer(MyAuthLayer::new());
  
  let app = foo.or(bar);