feat: add auth

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,10 +33,10 @@ pub fn get_nest_handlers() -> Router<crate::AppState> {
 //
 #[debug_handler]
 async fn get_all_books_handler(
-    state: State<AppState>
+    state: State<AppState>,
+    claims: Claims,
 ) -> Result<Json<Vec<book::BookItem>>, (StatusCode, String)> {
-    // let conn = state.conn.get_postgres_connection_pool();
-    let uid :i64 = 1;
+    let uid: i64 = claims.uid.clone();
     let all_books = Book::find()
         .filter(BookColumn::Uid.eq(uid))
         .all(&state.conn)
@@ -56,9 +58,9 @@ async fn get_all_books_handler(
 async fn get_book_by_id_handler(
     Path(id): Path<i64>,
     state: State<AppState>,
+    claims: Claims,
 ) -> Result<Json<BookItem>, (StatusCode, String)> {
-    // let conn = state.conn.get_postgres_connection_pool();
-    let uid: i64 = 1;
+    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))

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

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,7 +94,11 @@ where
             .await
             .map_err(|_| AuthError::InvalidToken)?;
         // Decode the user data
-        let token_data = decode::<Claims>(bearer.token(), &KEYS.decoding, &Validation::default())
+        let token_data = decode::<Claims>(
+            bearer.token(),
+            &KEYS.get().unwrap().decoding,
+            &Validation::default(),
+        )
         .map_err(|_| AuthError::InvalidToken)?;
 
         Ok(token_data.claims)