From e6a32eab513d5f369ff96d77d242fc207f143fc0 Mon Sep 17 00:00:00 2001
From: brian <BrianLTJ@users.noreply.github.com>
Date: Sun, 8 Jun 2025 22:24:44 +0800
Subject: [PATCH] feat: add auth

---
 src/api/book.rs        | 34 ++++++++++++++++++----------------
 src/main.rs            |  3 +++
 src/middleware/auth.rs | 33 +++++++++++++++++++--------------
 3 files changed, 40 insertions(+), 30 deletions(-)

diff --git a/src/api/book.rs b/src/api/book.rs
index a48bc02..18f78fb 100644
--- a/src/api/book.rs
+++ b/src/api/book.rs
@@ -6,6 +6,8 @@ use axum::{
 };
 use axum_macros::debug_handler;
 
+use crate::middleware::auth;
+use crate::middleware::auth::Claims;
 use crate::model::db::book::ActiveModel as BookActiveModel;
 use crate::model::db::book::Column as BookColumn;
 use crate::model::db::book::Model as BookModel;
@@ -31,19 +33,19 @@ pub fn get_nest_handlers() -> Router<crate::AppState> {
 //
 #[debug_handler]
 async fn get_all_books_handler(
-    state: State<AppState>
-) -> Result<Json<Vec<book::BookItem>>,(StatusCode,String)> {
-    // let conn = state.conn.get_postgres_connection_pool();
-    let uid :i64 = 1;
-    let all_books  = Book::find()
+    state: State<AppState>,
+    claims: Claims,
+) -> Result<Json<Vec<book::BookItem>>, (StatusCode, String)> {
+    let uid: i64 = claims.uid.clone();
+    let all_books = Book::find()
         .filter(BookColumn::Uid.eq(uid))
         .all(&state.conn)
         .await
-        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR,e.to_string()))?;
+        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
 
     let mut books: Vec<BookItem> = Vec::new();
     for b in all_books {
-        let book_resp = BookItem{
+        let book_resp = BookItem {
             id: b.id,
             name: b.name,
         };
@@ -54,11 +56,11 @@ async fn get_all_books_handler(
 
 #[debug_handler]
 async fn get_book_by_id_handler(
-    Path(id):Path<i64>,
-    state:State<AppState>,
-) -> Result<Json<BookItem>,(StatusCode,String)> {
-    // let conn = state.conn.get_postgres_connection_pool();
-    let uid: i64 = 1;
+    Path(id): Path<i64>,
+    state: State<AppState>,
+    claims: Claims,
+) -> Result<Json<BookItem>, (StatusCode, String)> {
+    let uid: i64 = claims.uid.clone();
     let book_query = Book::find()
         .filter(BookColumn::Uid.eq(uid))
         .filter(BookColumn::Id.eq(id))
@@ -84,10 +86,10 @@ async fn get_book_by_id_handler(
 #[debug_handler]
 async fn create_book_handler(
     state: State<AppState>,
+    claims: Claims,
     Json(payload): Json<BookInfo>,
 ) -> Result<Json<SimpleResponse>, (StatusCode, String)> {
-    // let conn = state.conn.get_postgres_connection_pool();
-    let uid: i64 = 1;
+    let uid: i64 = claims.uid.clone();
 
     let book = BookActiveModel {
         name: Set(payload.name.clone().to_owned()),
@@ -121,10 +123,10 @@ async fn create_book_handler(
 async fn update_book_handler(
     Path(id): Path<i64>,
     state: State<AppState>,
+    claims: Claims,
     Json(payload): Json<BookInfo>,
 ) -> Result<Json<SimpleResponse>, (StatusCode, String)> {
-    // let conn = state.conn.get_postgres_connection_pool();
-    let uid: i64 = 1;
+    let uid: i64 = claims.uid.clone();
 
     let exist_book = Book::find()
         .filter(BookColumn::Uid.eq(uid))
diff --git a/src/main.rs b/src/main.rs
index 4b2a53d..dc13b8b 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,3 +1,4 @@
+use crate::middleware::auth;
 use axum::{http::Method, Router};
 use clap::Parser;
 use sea_orm::{Database, DatabaseConnection};
@@ -91,6 +92,8 @@ async fn start_server(config: &Config) {
         .await
         .expect("Database connection failed.");
 
+    auth::initialize_jwt_key(config.keys.jwt.clone());
+
     let state = AppState { conn };
     // Build router
     let cors_layer = CorsLayer::new()
diff --git a/src/middleware/auth.rs b/src/middleware/auth.rs
index 6921cd6..328bf73 100644
--- a/src/middleware/auth.rs
+++ b/src/middleware/auth.rs
@@ -1,22 +1,18 @@
 use axum::{
     extract::FromRequestParts,
-    http::{
-        request::Parts,
-        StatusCode,
-    },
-    Json, RequestPartsExt,
+    http::{request::Parts, StatusCode},
     response::{IntoResponse, Response},
+    Json, RequestPartsExt,
 };
 use axum_extra::{
     headers::{authorization::Bearer, Authorization},
     TypedHeader,
 };
+use jsonwebtoken::{decode, encode, DecodingKey, EncodingKey, Header, Validation};
 use serde::{Deserialize, Serialize};
 use serde_json::json;
-use jsonwebtoken::{decode,encode, DecodingKey, EncodingKey, Header, Validation};
 use std::fmt::Display;
-use once_cell::sync::Lazy;
-use crate::util;
+use std::sync::OnceLock;
 
 #[derive(Debug, Serialize, Deserialize)]
 pub struct Claims {
@@ -46,10 +42,15 @@ pub enum AuthError {
     InvalidToken,
 }
 
-static KEYS: Lazy<Keys> = Lazy::new(|| {
-    let secret = std::env::var("JWT_SECRET").expect("JWT_SECRET must be set");
-    Keys::new(secret.as_bytes())
-});
+static KEYS: OnceLock<Keys> = OnceLock::new();
+
+pub fn initialize_jwt_key(key_str: String) {
+    let res = KEYS.set(Keys::new(key_str.as_bytes()));
+    match res {
+        Ok(_) => {}
+        Err(_) => panic!("jwt key initialize failed"),
+    }
+}
 
 struct Keys {
     encoding: EncodingKey,
@@ -93,8 +94,12 @@ where
             .await
             .map_err(|_| AuthError::InvalidToken)?;
         // Decode the user data
-        let token_data = decode::<Claims>(bearer.token(), &KEYS.decoding, &Validation::default())
-            .map_err(|_| AuthError::InvalidToken)?;
+        let token_data = decode::<Claims>(
+            bearer.token(),
+            &KEYS.get().unwrap().decoding,
+            &Validation::default(),
+        )
+        .map_err(|_| AuthError::InvalidToken)?;
 
         Ok(token_data.claims)
     }