feat: use extractor to load user

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"

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<S> FromRequestParts<S> for AuthUser
+    where S: Send + Sync {
+    type Rejection = ErrorMessage;
+
+    async fn from_request_parts(parts: &mut Parts, _state: &S) -> Result<Self, Self::Rejection> {
+        let headers = &parts.headers;
+        let user = login_by_token(headers).await
+            .ok_or(ErrorMessage::InvalidToken)?;
+
+        Ok(AuthUser(user))
+    }
+}

src/extractor/mod.rs

@@ -0,0 +1 @@
+pub mod auth;

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();

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<CreateFolderRequest>) -> 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<CreateFolderRequest>) -> Result<(), ErrorMessage> {
     let parent = Folder::find_by_id(query.parent).one(&*DATABASE).await.unwrap()
         .ok_or(ErrorMessage::InvalidParams("parent".to_string()))?;
 

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<HashMap<String, String>>) -> Result<String, ErrorMessage> {
-    let user = login_by_token(headers).await
-        .ok_or(ErrorMessage::InvalidToken)?;
+pub async fn get_folder_info(AuthUser(user): AuthUser, Query(query): Query<HashMap<String, String>>) -> Result<String, ErrorMessage> {
     let query_id: i64 = query.get("id")
         .ok_or(ErrorMessage::InvalidParams("id".to_string()))?
         .parse().map_err(|_| ErrorMessage::InvalidParams("folder_id".to_string()))?;

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<RenameFolderRequest>) -> Result<String, ErrorMessage> {
-    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<RenameFolderRequest>) -> Result<String, ErrorMessage> {
     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);

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;

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<DeletePictureRequest>, headers: HeaderMap) -> Result<String, ErrorMessage> {
+pub async fn delete_picture(Query(picture_id): Query<DeletePictureRequest>, AuthUser(user): AuthUser) -> Result<String, ErrorMessage> {
     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

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<ListPictureRequest>) -> impl IntoResponse {
+pub async fn list_picture(AuthUser(user): AuthUser, Query(query): Query<ListPictureRequest>) -> 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<ListPicture
     Ok(Json(response))
 }
 
-pub async fn get_picture_preview(header_map: HeaderMap, Query(query): Query<PictureGetPreviewRequest>)
+pub async fn get_picture_preview(AuthUser(user): AuthUser, Query(query): Query<PictureGetPreviewRequest>)
                                  -> Result<(HeaderMap, Vec<u8>), 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<i64> = serde_json::from_str(&e).unwrap();

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<RenamePictureRequest>) -> Result<String, ErrorMessage> {
-    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<RenamePictureRequest>) -> Result<String, ErrorMessage> {
     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);

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<String, ErrorMessage> {
+pub async fn post_picture(AuthUser(user): AuthUser, mut multipart: Multipart) -> Result<String, ErrorMessage> {
     let mut file: Option<Bytes> = None;
     let mut file_name: Option<String> = None;
-    let user = login_by_token(headers).await
-        .ok_or(ErrorMessage::InvalidToken)?;
     let mut resource_type = None;
     let mut dir = user.root;
 

src/service/share/create.rs

@@ -1,28 +1,24 @@
-use axum::http::HeaderMap;
 use axum::Json;
 use chrono::{Days, Utc};
 use sea_orm::{ActiveModelTrait, NotSet};
 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<CreateShareRequest>) -> Result<String, ErrorMessage> {
-    let user_id = login_by_token(headers).await
-        .ok_or(ErrorMessage::InvalidToken)?;
-
+pub async fn create_share(AuthUser(user): AuthUser, Json(body): Json<CreateShareRequest>) -> Result<String, ErrorMessage> {
     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<i64> = 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<CreateShareReques
         id: NotSet,
         content: Set(body.content),
         password: Set(password),
-        user_id: Set(user_id.id),
+        user_id: Set(user.id),
         mode: Set(share_level as i16),
         create_time: NotSet,
         valid_time: Set(Utc::now().checked_add_days(Days::new(7)).unwrap().naive_local()),

src/service/share/get.rs

@@ -8,6 +8,7 @@ use sea_orm::prelude::Uuid;
 use serde::{Deserialize, Serialize};
 
 use crate::DATABASE;
+use crate::extractor::auth::AuthUser;
 use crate::model::prelude::{Folder, Share, User, UserImage};
 use crate::model::share::ShareInfo;
 use crate::service::error::ErrorMessage;
@@ -119,10 +120,7 @@ pub async fn get_share_folder(Query(query): Query<GetShareFolderRequest>) -> Res
     Ok(serde_json::to_string(&response).unwrap())
 }
 
-pub async fn list_all_share(headers: HeaderMap, Query(query): Query<ListShareRequest>) -> Result<String, ErrorMessage> {
-    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<ListShareRequest>) -> Result<String, ErrorMessage> {
     if query.size > 100 {
         return Err(ErrorMessage::SizeTooLarge);
     }

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<CreatePayRequest>) -> Result<String, ErrorMessage> {
-    let user = login_by_token(headers).await.ok_or(ErrorMessage::InvalidToken)?;
-
+pub async fn creat_wechat_pay(AuthUser(user): AuthUser, Json(request): Json<CreatePayRequest>) -> Result<String, ErrorMessage> {
     let (result, code) = start_wechat(user.id, request.level, request.period, request.start_date).await;
     if !result {
         return Err(ErrorMessage::Other(code));

src/service/user/login.rs

@@ -32,7 +32,7 @@ pub async fn login_user(headers: HeaderMap, Query(request): Query<LoginRequest>)
 
     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<LoginRequest>)
     Ok((headers, serde_json::to_string(&user).unwrap()))
 }
 
-pub async fn login_by_token(header: HeaderMap) -> Option<crate::model::user::Model> {
+pub async fn login_by_token(header: &HeaderMap) -> Option<crate::model::user::Model> {
     if !header.contains_key("token") { return None; }
     let token = header.get("token").unwrap().to_str().unwrap();
     let uid = TOKEN_CACHE.get(token).await;