rust:axum+sqlx实战学习笔记3

linken_coding
2,094 阅读4分钟

携手创作，共同成长！这是我参与「掘金日新计划 · 8 月更文挑战」的第3天，点击查看活动详情

通过前面的两章的学习，我们掌握了axum的基本使用方法，如何使用handler函数获取路由中的参数以及对响应数据进行处理。

接下来，我们将继续深入探究axum的其他部分，逐步搭建起axum+sqlx的web后端服务。

那么，下面我们开始进入axum中间件领域。

1、首先，axum是在Tokio全家桶的一部分，所以axum的中间件其实与tower和tower-http一起工作的。

2、所以可以用tower或者用axum自带的函数编写中间件

3、axum 提供了许多编写中间件的方法，在不同的抽象级别和不同的优缺点。

方法一：使用axum::middleware::from_fn

特点：可以使用async/await语法 下面来看官方的文档例子：

use axum::{
    Router,
    http::{Request, StatusCode},
    routing::get,
    response::{IntoResponse, Response},
    middleware::{self, Next},
};

async fn auth<B>(req: Request<B>, next: Next<B>) -> Result<Response, StatusCode> {
    let auth_header = req.headers()
        .get(http::header::AUTHORIZATION)
        .and_then(|header| header.to_str().ok());

    match auth_header {
        Some(auth_header) if token_is_valid(auth_header) => {
            Ok(next.run(req).await)
        }
        _ => Err(StatusCode::UNAUTHORIZED),
    }
}

fn token_is_valid(token: &str) -> bool {
    // ...
}

let app = Router::new()
    .route("/", get(|| async { /* ... */ }))
    .route_layer(middleware::from_fn(auth));

如果想要传递一些共享的状态，可以这样：

use axum::{
    Router,
    http::{Request, StatusCode},
    routing::get,
    response::{IntoResponse, Response},
    middleware::{self, Next}
};

#[derive(Clone)]
struct State { /* ... */ }

async fn my_middleware<B>(
    req: Request<B>,
    next: Next<B>,
    state: State,
) -> Response {
    // ...
}

let state = State { /* ... */ };

let app = Router::new()
    .route("/", get(|| async { /* ... */ }))
    .route_layer(middleware::from_fn(move |req, next| {
        my_middleware(req, next, state.clone())
    }));

经过这两个例子，我们大概可以看出axum::middleware::from_fn的使用方法。其中真正要注意的是如何获取参数和返回数据。

  1. 参数req: Request<B>, next: Next<B>,分别表示HTTP请求和下一步的函数。所以中间件的参数基本是从req: Request<B>中获取，处理完成后用next: Next<B>跳转到下一步。
  2. 返回的响应数据要实现IntoResponse特质。

方法二：使用 axum::middleware::from_extractor

特点：从提取器创建中间件，下面看官方例子：

use axum::{
    extract::{FromRequest, RequestParts},
    middleware::from_extractor,
    routing::{get, post},
    Router,
};
use http::{header, StatusCode};
use async_trait::async_trait;

// An extractor that performs authorization.
struct RequireAuth;

#[async_trait]
impl<B> FromRequest<B> for RequireAuth
where
    B: Send,
{
    type Rejection = StatusCode;

    async fn from_request(req: &mut RequestParts<B>) -> Result<Self, Self::Rejection> {
        let auth_header = req
            .headers()
            .get(header::AUTHORIZATION)
            .and_then(|value| value.to_str().ok());

        match auth_header {
            Some(auth_header) if token_is_valid(auth_header) => {
                Ok(Self)
            }
            _ => Err(StatusCode::UNAUTHORIZED),
        }
    }
}

fn token_is_valid(token: &str) -> bool {
    // ...
}

async fn handler() {
    // If we get here the request has been authorized
}

async fn other_handler() {
    // If we get here the request has been authorized
}

let app = Router::new()
    .route("/", get(handler))
    .route("/foo", post(other_handler))
    // The extractor will run before all routes
    .route_layer(from_extractor::<RequireAuth>());

从例子中，我们可以看出，这种中间件其实是为struct实现了FromRequest的特质。而参数RequestParts<B>和上面的Request<B>可以说是大同小异，都是HTTP请求的数据。同理，返回的数据也需要实现IntoResponse特质。

那么接下来，我们从这两种例子，仿写一个属于自己的中间件。因为我们搭建的web后端服务需要鉴权的模块，这个作为中间件自然能够满足我们需求。 首先，我们用到jwt的鉴权库，创建一个jwt加密密钥，并为密钥实现加密和解密。

use jsonwebtoken::{DecodingKey, EncodingKey};
use once_cell::sync::Lazy;
use std::env;

///环境变量密钥，
pub static KEYS: Lazy<Keys> = Lazy::new(|| {
    let secret = env::var("JWT_SECRET").expect("JWT_SECRET must be set");
    Keys::new(secret.as_bytes())
});

///认证错误类型
#[derive(Debug)]
pub enum AuthError {
    WrongCredentials,   //错误的凭据
    MissingCredentials, //丢失凭据
    TokenCreation,      //令牌创建
    InvalidToken,       //无效令牌
}
pub struct Keys {
    pub encoding: EncodingKey,
    pub decoding: DecodingKey,
}
impl Keys {
    fn new(secret: &[u8]) -> Self {
        Self {
            encoding: EncodingKey::from_secret(secret),
            decoding: DecodingKey::from_secret(secret),
        }
    }
}

然后编写属于自己的jwt中间件：

use crate::common::response::RespVO;
use crate::common::types::{Role, Sex, UserState};
use crate::jwt::KEYS;
use axum::{
    async_trait,
    extract::TypedHeader,
    extract::{FromRequest, RequestParts},
    headers::{authorization::Bearer, Authorization},
    http::StatusCode,
    response::Response,
    Json,
};
use jsonwebtoken::{decode, Validation};
use serde::{Deserialize, Serialize};
use sqlx::{Decode, Encode, Type};

#[derive(Debug, Clone, Serialize, Deserialize, Decode, Encode, Type)]
pub struct Claims {
    pub username: String,
    pub email: String,
    pub sex: Sex,
    pub avatar: Option<String>,
    pub role: Option<Role>,
    pub status: Option<UserState>,
    pub exp: Option<i32>,
}

#[async_trait]
impl<B> FromRequest<B> for Claims
where
    B: Send,
{
    type Rejection = Json<RespVO<String>>;
    async fn from_request(req: &mut RequestParts<B>) -> Result<Self, Self::Rejection> {
        // Extract the token from the authorization header
        let TypedHeader(Authorization(bearer)) =
            TypedHeader::<Authorization<Bearer>>::from_request(req)
                .await
                .map_err(|_| {
                    Json(RespVO::<String>::from_error_info(
                        StatusCode::UNAUTHORIZED,
                        "未认证",
                    ))
                })?;

        // Decode the user data
        let token_data = decode::<Claims>(bearer.token(), &KEYS.decoding, &Validation::default())
            .map_err(|_| {
            Json(RespVO::<String>::from_error_info(
                StatusCode::UNAUTHORIZED,
                "token无效",
            ))
        })?;
        Ok(token_data.claims)
    }
}

这里先忽视jwt如何进行加密，主要看中间件内部如何对jwt进行处理。同时这里也用到了上一章里面，我们写好的公共响应数据模型RespVO封装作为报错信息。

最后，照例到了总结环节：

  1. axum的中间件可以复用tower和tower-http
  2. axum自带编写中间件的函数有： axum::middleware::from_fn和 axum::middleware::from_extractor 两者大同小异，都是获取HTTP请求参数进行处理后，返回相应的数据或者实现了IntoResponse特质的报错信息。
avatar
文章
阅读
粉丝