From b8b36d90d65cf7e522ff95723c81fe9cc3037174 Mon Sep 17 00:00:00 2001 From: zrll_ Date: Wed, 12 Jun 2024 17:34:20 +0800 Subject: [PATCH] feat: use extractor to load user --- Cargo.lock | 5 +++-- Cargo.toml | 1 + src/extractor/auth.rs | 22 ++++++++++++++++++++++ src/extractor/mod.rs | 1 + src/main.rs | 1 + src/service/folder/create.rs | 5 ++--- src/service/folder/get.rs | 5 ++--- src/service/folder/rename.rs | 6 ++---- src/service/mod.rs | 2 +- src/service/picture/delete.rs | 6 ++---- src/service/picture/get.rs | 10 +++------- src/service/picture/rename.rs | 6 ++---- src/service/picture/upload.rs | 7 ++----- src/service/share/create.rs | 12 ++++-------- src/service/share/get.rs | 6 ++---- src/service/trade/wechat/mod.rs | 6 ++---- src/service/user/login.rs | 4 ++-- 17 files changed, 54 insertions(+), 51 deletions(-) create mode 100644 src/extractor/auth.rs create mode 100644 src/extractor/mod.rs diff --git a/Cargo.lock b/Cargo.lock index 3803ade..69c38ec 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -250,9 +250,9 @@ dependencies = [ [[package]] name = "async-trait" -version = "0.1.77" +version = "0.1.80" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c980ee35e870bd1a4d2c8294d4c04d0499e67bca1e4b5cefcc693c2fa00caea9" +checksum = "c6fa2087f2753a7da8cc1c0dbfcf89579dd57458e36769de5ac750b4671737ca" dependencies = [ "proc-macro2", "quote", @@ -3811,6 +3811,7 @@ dependencies = [ "shadow-rs 0.26.1", "tokio", "toml", + "tower", "tower-http", "tracing", "tracing-appender", diff --git a/Cargo.toml b/Cargo.toml index 0809c53..b439ad4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -53,6 +53,7 @@ migration = { path = "migration" } lsys-lib-sms = "0.1.0" wechat-pay-rust-sdk = { version = "0.2.14", features = ["debug-print"] } rsa = "0.9.6" +tower = "0.4.13" [build-dependencies] shadow-rs = "0.24.1" diff --git a/src/extractor/auth.rs b/src/extractor/auth.rs new file mode 100644 index 0000000..77c0124 --- /dev/null +++ b/src/extractor/auth.rs @@ -0,0 +1,22 @@ +use axum::async_trait; +use axum::extract::{FromRequest, FromRequestParts}; +use axum::http::request::Parts; + +use crate::service::error::ErrorMessage; +use crate::service::user::login::login_by_token; + +pub struct AuthUser(pub crate::model::user::Model); + +#[async_trait] +impl FromRequestParts for AuthUser + where S: Send + Sync { + type Rejection = ErrorMessage; + + async fn from_request_parts(parts: &mut Parts, _state: &S) -> Result { + let headers = &parts.headers; + let user = login_by_token(headers).await + .ok_or(ErrorMessage::InvalidToken)?; + + Ok(AuthUser(user)) + } +} \ No newline at end of file diff --git a/src/extractor/mod.rs b/src/extractor/mod.rs new file mode 100644 index 0000000..5696e21 --- /dev/null +++ b/src/extractor/mod.rs @@ -0,0 +1 @@ +pub mod auth; \ No newline at end of file diff --git a/src/main.rs b/src/main.rs index e03db9b..681826e 100644 --- a/src/main.rs +++ b/src/main.rs @@ -44,6 +44,7 @@ use crate::service::user::register::register_user; mod config; mod model; mod service; +mod extractor; lazy_static! { static ref CONFIG: Config = Config::new(); diff --git a/src/service/folder/create.rs b/src/service/folder/create.rs index 137d611..b69ad8c 100644 --- a/src/service/folder/create.rs +++ b/src/service/folder/create.rs @@ -5,13 +5,12 @@ use sea_orm::ActiveValue::Set; use serde::Deserialize; use crate::DATABASE; +use crate::extractor::auth::AuthUser; use crate::model::prelude::Folder; use crate::service::error::ErrorMessage; use crate::service::user::login::login_by_token; -pub async fn create_folder(header_map: HeaderMap, Json(query): Json) -> Result<(), ErrorMessage> { - let user = login_by_token(header_map).await.ok_or(ErrorMessage::InvalidToken)?; - +pub async fn create_folder(AuthUser(user): AuthUser, Json(query): Json) -> Result<(), ErrorMessage> { let parent = Folder::find_by_id(query.parent).one(&*DATABASE).await.unwrap() .ok_or(ErrorMessage::InvalidParams("parent".to_string()))?; diff --git a/src/service/folder/get.rs b/src/service/folder/get.rs index a6aa935..dd498db 100644 --- a/src/service/folder/get.rs +++ b/src/service/folder/get.rs @@ -5,13 +5,12 @@ use axum::http::HeaderMap; use sea_orm::EntityTrait; use crate::DATABASE; +use crate::extractor::auth::AuthUser; use crate::model::prelude::Folder; use crate::service::error::ErrorMessage; use crate::service::user::login::login_by_token; -pub async fn get_folder_info(headers: HeaderMap, Query(query): Query>) -> Result { - let user = login_by_token(headers).await - .ok_or(ErrorMessage::InvalidToken)?; +pub async fn get_folder_info(AuthUser(user): AuthUser, Query(query): Query>) -> Result { let query_id: i64 = query.get("id") .ok_or(ErrorMessage::InvalidParams("id".to_string()))? .parse().map_err(|_| ErrorMessage::InvalidParams("folder_id".to_string()))?; diff --git a/src/service/folder/rename.rs b/src/service/folder/rename.rs index feb8077..0ef9ea3 100644 --- a/src/service/folder/rename.rs +++ b/src/service/folder/rename.rs @@ -5,13 +5,11 @@ use sea_orm::ActiveValue::Set; use serde::Deserialize; use crate::DATABASE; +use crate::extractor::auth::AuthUser; use crate::model::prelude::Folder; use crate::service::error::ErrorMessage; -pub async fn rename_folder(header_map: HeaderMap, Query(query): Query) -> Result { - let user = crate::service::user::login::login_by_token(header_map).await - .ok_or(ErrorMessage::InvalidToken)?; - +pub async fn rename_folder(AuthUser(user): AuthUser, Query(query): Query) -> Result { let folder = Folder::find_by_id(query.id).one(&*DATABASE).await.unwrap().ok_or(ErrorMessage::NotFound)?; if folder.user_id != user.id || folder.id == user.root { return Err(ErrorMessage::PermissionDenied); diff --git a/src/service/mod.rs b/src/service/mod.rs index 466e70d..ab121ec 100644 --- a/src/service/mod.rs +++ b/src/service/mod.rs @@ -3,4 +3,4 @@ pub mod user; pub mod folder; pub mod share; pub mod trade; -mod error; +pub mod error; diff --git a/src/service/picture/delete.rs b/src/service/picture/delete.rs index 859a5b4..4671d32 100644 --- a/src/service/picture/delete.rs +++ b/src/service/picture/delete.rs @@ -6,16 +6,14 @@ use serde::Deserialize; use tracing::{error, info}; use crate::DATABASE; +use crate::extractor::auth::AuthUser; use crate::model::prelude::{Folder, Image, UserImage}; use crate::service::error::ErrorMessage; use crate::service::picture::file::remove_file; -pub async fn delete_picture(Query(picture_id): Query, headers: HeaderMap) -> Result { +pub async fn delete_picture(Query(picture_id): Query, AuthUser(user): AuthUser) -> Result { info!("delete picture: {}", picture_id.image_id); - let user = crate::service::user::login::login_by_token(headers).await - .ok_or(ErrorMessage::InvalidToken)?; - let user_picture = UserImage::find_by_id(picture_id.image_id) .one(&*crate::DATABASE) .await diff --git a/src/service/picture/get.rs b/src/service/picture/get.rs index a287cdb..00feb66 100644 --- a/src/service/picture/get.rs +++ b/src/service/picture/get.rs @@ -6,17 +6,16 @@ use sea_orm::{ColumnTrait, EntityTrait, Order, PaginatorTrait, QueryFilter, Quer use serde::{Deserialize, Serialize}; use crate::DATABASE; +use crate::extractor::auth::AuthUser; use crate::model::prelude::UserImage; use crate::service::error::ErrorMessage; use crate::service::picture::compress::ImageFile; use crate::service::user::login::login_by_token; -pub async fn list_picture(header_map: HeaderMap, Query(query): Query) -> impl IntoResponse { +pub async fn list_picture(AuthUser(user): AuthUser, Query(query): Query) -> impl IntoResponse { if query.size > 100 { return Err(ErrorMessage::SizeTooLarge); } - let user = login_by_token(header_map).await - .ok_or(ErrorMessage::InvalidToken)?; let pictures = UserImage::find() .filter(crate::model::user_image::Column::UserId.eq(user.id)) @@ -33,11 +32,8 @@ pub async fn list_picture(header_map: HeaderMap, Query(query): Query) +pub async fn get_picture_preview(AuthUser(user): AuthUser, Query(query): Query) -> Result<(HeaderMap, Vec), ErrorMessage> { - let user = login_by_token(header_map).await - .ok_or(ErrorMessage::InvalidToken)?; - // let level: LevelInfo = user.level.into_iter() // .map(|e| { // let raw: Vec = serde_json::from_str(&e).unwrap(); diff --git a/src/service/picture/rename.rs b/src/service/picture/rename.rs index 0206e64..1c7d8bd 100644 --- a/src/service/picture/rename.rs +++ b/src/service/picture/rename.rs @@ -5,13 +5,11 @@ use sea_orm::ActiveValue::Set; use serde::Deserialize; use crate::DATABASE; +use crate::extractor::auth::AuthUser; use crate::model::prelude::UserImage; use crate::service::error::ErrorMessage; -pub async fn rename_picture(header_map: HeaderMap, Query(query): Query) -> Result { - let user = crate::service::user::login::login_by_token(header_map).await - .ok_or(ErrorMessage::InvalidToken)?; - +pub async fn rename_picture(AuthUser(user): AuthUser, Query(query): Query) -> Result { let picture = UserImage::find_by_id(query.id).one(&*DATABASE).await.unwrap().ok_or(ErrorMessage::NotFound)?; if picture.user_id != user.id { return Err(ErrorMessage::PermissionDenied); diff --git a/src/service/picture/upload.rs b/src/service/picture/upload.rs index c3c4bf9..0d69c60 100644 --- a/src/service/picture/upload.rs +++ b/src/service/picture/upload.rs @@ -1,21 +1,18 @@ use axum::body::Bytes; use axum::extract::Multipart; -use axum::http::HeaderMap; use sea_orm::{ActiveModelTrait, EntityTrait, IntoActiveModel, NotSet}; use sea_orm::ActiveValue::Set; use tracing::{debug, error}; use crate::DATABASE; +use crate::extractor::auth::AuthUser; use crate::model::prelude::Folder; use crate::service::error::ErrorMessage; use crate::service::picture::file::save_file; -use crate::service::user::login::login_by_token; -pub async fn post_picture(headers: HeaderMap, mut multipart: Multipart) -> Result { +pub async fn post_picture(AuthUser(user): AuthUser, mut multipart: Multipart) -> Result { let mut file: Option = None; let mut file_name: Option = None; - let user = login_by_token(headers).await - .ok_or(ErrorMessage::InvalidToken)?; let mut resource_type = None; let mut dir = user.root; diff --git a/src/service/share/create.rs b/src/service/share/create.rs index d928cd0..34b4ea9 100644 --- a/src/service/share/create.rs +++ b/src/service/share/create.rs @@ -1,4 +1,3 @@ -use axum::http::HeaderMap; use axum::Json; use chrono::{Days, Utc}; use sea_orm::{ActiveModelTrait, NotSet}; @@ -6,23 +5,20 @@ use sea_orm::ActiveValue::Set; use serde::Deserialize; use crate::DATABASE; +use crate::extractor::auth::AuthUser; use crate::service::error::ErrorMessage; use crate::service::share::content::ContentType; use crate::service::user::level::LevelInfo; -use crate::service::user::login::login_by_token; use crate::service::user::password::generate_password_hash; -pub async fn create_share(headers: HeaderMap, Json(body): Json) -> Result { - let user_id = login_by_token(headers).await - .ok_or(ErrorMessage::InvalidToken)?; - +pub async fn create_share(AuthUser(user): AuthUser, Json(body): Json) -> Result { for content in body.content.iter() { if !ContentType::verify(content) { return Err(ErrorMessage::InvalidParams(format!("content {}", content))); } } - let level: LevelInfo = user_id.level.into_iter() + let level: LevelInfo = user.level.into_iter() .map(|e| { let raw: Vec = serde_json::from_str(&e).unwrap(); LevelInfo::try_from(raw).unwrap_or_else(|_| LevelInfo::get_free_level()) @@ -42,7 +38,7 @@ pub async fn create_share(headers: HeaderMap, Json(body): Json) -> Res Ok(serde_json::to_string(&response).unwrap()) } -pub async fn list_all_share(headers: HeaderMap, Query(query): Query) -> Result { - let user = crate::service::user::login::login_by_token(headers).await - .ok_or(ErrorMessage::InvalidToken)?; - +pub async fn list_all_share(AuthUser(user): AuthUser, Query(query): Query) -> Result { if query.size > 100 { return Err(ErrorMessage::SizeTooLarge); } diff --git a/src/service/trade/wechat/mod.rs b/src/service/trade/wechat/mod.rs index 5b2352e..c80b609 100644 --- a/src/service/trade/wechat/mod.rs +++ b/src/service/trade/wechat/mod.rs @@ -2,18 +2,16 @@ use axum::http::HeaderMap; use axum::Json; use chrono::NaiveDate; use serde::Deserialize; +use crate::extractor::auth::AuthUser; use crate::service::error::ErrorMessage; use crate::service::trade::wechat::start::start_wechat; use crate::service::user::level::Level; -use crate::service::user::login::login_by_token; mod start; pub mod recall; -pub async fn creat_wechat_pay(headers: HeaderMap, Json(request): Json) -> Result { - let user = login_by_token(headers).await.ok_or(ErrorMessage::InvalidToken)?; - +pub async fn creat_wechat_pay(AuthUser(user): AuthUser, Json(request): Json) -> Result { let (result, code) = start_wechat(user.id, request.level, request.period, request.start_date).await; if !result { return Err(ErrorMessage::Other(code)); diff --git a/src/service/user/login.rs b/src/service/user/login.rs index 47f6978..1262bd6 100644 --- a/src/service/user/login.rs +++ b/src/service/user/login.rs @@ -32,7 +32,7 @@ pub async fn login_user(headers: HeaderMap, Query(request): Query) let user = if headers.contains_key("token") { //login with token debug!("Token: {}", headers.get("token").unwrap().to_str().unwrap()); - login_by_token(headers).await + login_by_token(&headers).await .ok_or(ErrorMessage::InvalidToken)? } else { //login with username and password need_token = true; @@ -97,7 +97,7 @@ pub async fn login_user(headers: HeaderMap, Query(request): Query) Ok((headers, serde_json::to_string(&user).unwrap())) } -pub async fn login_by_token(header: HeaderMap) -> Option { +pub async fn login_by_token(header: &HeaderMap) -> Option { if !header.contains_key("token") { return None; } let token = header.get("token").unwrap().to_str().unwrap(); let uid = TOKEN_CACHE.get(token).await;