【Rust基础】使用Rocket从Token中提取用户信息

概要

Rocket没有像Spring那样提供了ContextHolder来直接获取当前请求对象，我们目前也不能直接通过一个静态工具类来获取请求对象，因为它是基于协程而不是线程的，所以我们只能通过接口中从request中提取需要的数据，作为函数参数传递下去。

Rocket提供了FromRequest这个trait，通过它可以拿到原始的request对象，从而返回我们需要的数据。

从Token中提取用户信息

用户登录后，将返回一个Token给客户端，客户端请求时将Token放到Header中，服务端从Token解析用户信息，这个Token可以是简单UUID，也可以是JWT等。我们可以从Request中提取这些Token。

rust 复制代码

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct UserPrincipal {
    /// 用户ID
    pub id: i64,
    /// 昵称
    pub nickname: String,
}

#[rocket::async_trait]
impl<'r> FromRequest<'r> for UserPrincipal {
    type Error = &'r str;

    async fn from_request(req: &'r Request<'_>) -> Outcome<Self, Self::Error> {
        // 从header提取Token
        let user_info = req.headers().get("Token").next().and_then(|token| {
            // 从redis获取token
            let user = cache::get::<UserPrincipal>(&format!(
                "{}{}",
                CachePrefix::AccessToken.to_string(),
                token
            ));
            if user.is_none() {
                return None;
            }
            let user = user.unwrap();
            Some(UserPrincipal {
                id: user.id.clone(),
                nickname: user.nickname.clone(),
            })
        });

        let user = match user_info {
            Some(user) => user,
            None => return Outcome::Error((Status::Unauthorized, "未登录")),
        };

        // 使用user验证其他场景

        Outcome::Success(user)
    }
}

以上，我们定义了一个UserPrincipal作为在系统中关联的用户认证主体，通过FromRequest从Header拿到Token后，再去redis中查询用户信息，如果未拿到则说明token已失效需要重新登录，返回401。

请求上下文

之前提到，我们无法直接拿到请求上下文，那我们可以自定义一个上下文，通过FromRequest这个trait来组装需要的数据。

例如：

rust 复制代码

pub struct ReqContext {
    /// 是否移动端
    pub is_mobile: bool,
    // 其他需要提取的数据
}

#[rocket::async_trait]
impl<'r> FromRequest<'r> for ReqContext {
    type Error = ();

    async fn from_request(req: &'r Request<'_>) -> Outcome<Self, Self::Error> {
        let re = Regex::new(MOBILE_USER_AGENT_REGEX).expect("Invalid regex");

        let user_agent = req.headers().get_one("User-Agent").unwrap_or("");

        Outcome::Success(ReqContext {
            is_mobile: re.is_match(user_agent),
        })
    }
}

这样以来，我们就可以在API中使用他们了。

rust 复制代码

#[post("/api", data = "<req>")]
pub async fn buy_member_card(
    req: Json<SomeData>,
    user: UserPrincipal,
    req_context: ReqContext,
) -> Res<OrderCreateRes> {
    match service::some_fn(req, user, req_context).await {
        Ok(res) => Res::success(res),
        Err(e) => Res::from_error(e),
    }
}

Rocket关闭时释放资源

对于一些连接池，我们需要在系统即将关闭的时候释放某些连接，或者等待未处理完成的数据处理完成后再关闭，Rocket提供了一个Fairing trait，可以attach自定义的Fairing来监听shutdown事件，从而实现释放资源。

rust 复制代码

pub struct ShutdownFairing;
#[rocket::async_trait]
impl Fairing for ShutdownFairing {
    fn info(&self) -> Info {
        Info {
            name: "Shutdown Handler",
            kind: Kind::Shutdown,
        }
    }

    async fn on_shutdown(&self, _rocket: &Rocket<Orbit>) {
        //释放资源
    }
}

需要注意的是：

释放资源时，不要存在阻塞操作，否则会导致进程阻塞而无法关闭程序。
对于kill -9可能无法正确释放资源，尽量不要强制kill。