Refactor; Fix CORS when login via passkeys

register.html

@@ -110,10 +110,16 @@
                         },
                     }),
                 })
-                .then((response) => {
+                .then(async (response) => {
                     if (response.ok){
                         console.log("Logged In!");
+                        var redir = await response.json();
+                        if(redir && redir.status === "ok"){
+                            console.log("Redirecting");
+                            window.location.replace(redir.redirect);
+                        } else {
                             window.location.replace("/");
+                        }
                     } else {
                         console.log("Error");
                     }

src/controllers/mod.rs

@@ -1,2 +1,3 @@
 pub mod otp_controllers;
 pub mod passkey_login_controllers;
+pub mod passkey_register_controllers;

src/controllers/passkey_login_controllers.rs

@@ -130,7 +130,7 @@ pub async fn finish_authentication(
             cookies.add(new_cookie);
             // 从Cookie中恢复重定向信息
             match original_uri {
-                Some(redirect) => return Ok(Redirect::to(redirect.value())),
+                Some(redirect) => return Ok(format!("{{\"status\": \"ok\", \"redirect\":\"{}\"}}",redirect.value().to_string())),
                 _ => (),
             };
             // 处理完成重定向后，清除Cookie

src/controllers/passkey_register_controllers.rs

@@ -0,0 +1,169 @@
+
+use axum::http::StatusCode;
+use uuid::Uuid;
+use webauthn_rs::prelude::{CreationChallengeResponse, CredentialID, PasskeyAuthentication, PasskeyRegistration, RegisterPublicKeyCredential, WebauthnError};
+use crate::dto::credential_mapper::add_credential_by_id;
+
+use crate::dto::user_mapper::{create_user_if_non_existent, get_user_count_by_uid};
+
+use crate::SESSION_ACTIVE_TIME;
+
+use axum::response::IntoResponse;
+
+use crate::dto::{credential_mapper::get_credential_from_uid, user_mapper::get_uid_by_name};
+
+
+use crate::config::COOKIE_NAME;
+
+use tracing::{info,debug,error};
+
+use axum::extract::{Json, Path};
+
+
+use tower_sessions::Session;
+
+use tower_cookies::Cookies;
+
+use crate::entities::ServerState;
+
+use std::collections::HashMap;
+use std::sync::Arc;
+use std::time::{Duration, Instant};
+
+use axum::extract::State;
+
+
+// TODO - Improve error handling and messages
+fn auth_user(cookie_content: String, session_table: &mut HashMap<Uuid, Instant>) -> bool {
+    let Ok(uuid) = Uuid::parse_str(&cookie_content) else {
+        info!("此用户Session不在表中");
+        return false;
+    };
+    let Some(expire) = session_table.get(&uuid) else {
+        info!("此用户Session已过期");
+        return false;
+    };
+    if *expire <= Instant::now() {
+        info!("此用户Session已过期");
+        return false;
+    }
+    session_table.insert(
+        uuid,
+        Instant::now() + Duration::from_secs(*SESSION_ACTIVE_TIME),
+    );
+    tracing::info!("valid cookie {}", uuid);
+    true
+}
+
+pub async fn start_register(
+    State(state): State<Arc<ServerState>>,
+    cookies: Cookies,
+    session: Session,
+    Path(username): Path<String>,
+) -> Result<Json<CreationChallengeResponse>, String> {
+    tracing::info!("开始注册");
+    // todo!("Auth User");
+    let Some(cookie_content) = cookies.get(&COOKIE_NAME) else {
+        tracing::info!("用户没有Cookie");
+        return Err(WebauthnError::AuthenticationFailure.to_string());
+    };
+    let mut session_table = state.session.lock().await;
+    if !auth_user(cookie_content.value().to_string(), &mut session_table) {
+        tracing::info!("用户Cookie无效，不在Session表中");
+        return Err(WebauthnError::AuthenticationFailure.to_string());
+    }
+    // Remove any previous registrations that may have occurred from the session.
+    let _ = session
+        .remove::<(String, Uuid, PasskeyRegistration)>("reg_state")
+        .await
+        .map_err(|_| WebauthnError::AuthenticationFailure.to_string())?;
+    let pool = &state.db;
+    let username_tmp = username.clone();
+    let (user_id, exclude_credentials): (Uuid, Option<Vec<CredentialID>>) =
+        match get_uid_by_name(username_tmp, pool).await {
+            Ok(uid) => {
+                let exclude_passkeys = get_credential_from_uid(uid, pool).await?;
+                (
+                    uid,
+                    Some(
+                        exclude_passkeys
+                            .into_iter()
+                            .map(|k| k.cred_id().clone())
+                            .collect(),
+                    ),
+                )
+            }
+            Err(_) => (Uuid::new_v4(), None),
+        };
+
+    let res = match state.webauthn.start_passkey_registration(
+        user_id,
+        &username,
+        &username,
+        exclude_credentials,
+    ) {
+        Ok((ccr, reg_state)) => {
+            // Note that due to the session store in use being a server side memory store, this is
+            // safe to store the reg_state into the session since it is not client controlled and
+            // not open to replay attacks. If this was a cookie store, this would be UNSAFE.
+            session
+                .insert("reg_state", (username, user_id, reg_state))
+                .await
+                .map_err(|_| WebauthnError::ChallengePersistenceError.to_string())?;
+            Json(ccr)
+        }
+        Err(e) => {
+            debug!("challenge_register -> {:?}", e);
+            return Err(WebauthnError::ChallengePersistenceError.to_string());
+        }
+    };
+    Ok(res)
+}
+
+pub async fn finish_register(
+    State(state): State<Arc<ServerState>>,
+    session: Session,
+    Json(reg): Json<RegisterPublicKeyCredential>,
+) -> Result<impl IntoResponse, String> {
+    info!("完成注册...");
+    let pool = &state.db;
+    let Ok(Some((user_name, user_id, reg_state))) = session
+        .get::<(String, Uuid, PasskeyRegistration)>("reg_state")
+        .await
+    else {
+        return Err(WebauthnError::AuthenticationFailure.to_string()); //Corrupt Session
+    };
+
+    let _ = session
+        .remove::<(Uuid, PasskeyAuthentication)>("reg_state")
+        .await;
+
+    let res = match state.webauthn.finish_passkey_registration(&reg, &reg_state) {
+        Ok(key) => {
+            info!("Passkey 正常");
+            info!("检查用户是否存在");
+            // Check if the user_id already exists
+            let user_count = get_user_count_by_uid(user_id, pool).await?;
+            // If the user doesn't exist, insert them into the users table
+            if user_count == 0
+                && create_user_if_non_existent(user_id, user_name.to_string(), pool).await? != 1
+            {
+                return Err(WebauthnError::AuthenticationFailure.to_string());
+            }
+
+            // Insert the key into the auth table
+            if add_credential_by_id(user_id, &key, pool).await? != 1 {
+                error!("将用户凭据持久化时失败,rows_affected!=1");
+                return Err(WebauthnError::AuthenticationFailure.to_string());
+            }
+
+            StatusCode::OK
+        }
+        Err(e) => {
+            error!("challenge_register -> {:?}", e);
+            StatusCode::BAD_REQUEST
+        }
+    };
+
+    Ok(res)
+}

src/main.rs

@@ -6,7 +6,6 @@ use entities::*;
 
 use controllers::passkey_login_controllers::{start_authentication, finish_authentication};
 use services::{
-    auth::start_register,
     gc_services::gc_task,
     login, login_page, register_page,
 };
@@ -22,7 +21,8 @@ pub mod controllers;
 pub mod dto;
 pub mod entities;
 pub mod services;
-use crate::{config::*, services::auth::finish_register};
+use config::*;
+use controllers::passkey_register_controllers::{finish_register,start_register};
 #[tokio::main]
 async fn main() {
     // 初始化日志记录器

src/services/auth.rs

@@ -1,231 +0,0 @@
-use std::collections::HashMap;
-use std::sync::Arc;
-use std::time::Duration;
-
-use axum::extract::State;
-use axum::{
-    extract::{Json, Path},
-    http::StatusCode,
-    response::IntoResponse,
-};
-use std::time::Instant;
-use tower_cookies::Cookies;
-use tower_sessions::Session;
-use tracing::*;
-
-/*
- * Webauthn RS auth handlers.
- * These files use webauthn to process the data received from each route, and are closely tied to axum
- */
-
-// 1. Import the prelude - this contains everything needed for the server to function.
-use webauthn_rs::prelude::*;
-
-use crate::config::{COOKIE_NAME, SESSION_ACTIVE_TIME};
-use crate::dto::credential_mapper::{add_credential_by_id, get_credential_from_uid};
-use crate::dto::user_mapper::{
-    create_user_if_non_existent, get_uid_by_name, get_user_count_by_uid,
-};
-use crate::ServerState;
-
-// 2. The first step a client (user) will carry out is requesting a credential to be
-// registered. We need to provide a challenge for this. The work flow will be:
-//
-//          ┌───────────────┐     ┌───────────────┐      ┌───────────────┐
-//          │ Authenticator │     │    Browser    │      │     Site      │
-//          └───────────────┘     └───────────────┘      └───────────────┘
-//                  │                     │                      │
-//                  │                     │     1. Start Reg     │
-//                  │                     │─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─▶│
-//                  │                     │                      │
-//                  │                     │     2. Challenge     │
-//                  │                     │◀ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┤
-//                  │                     │                      │
-//                  │  3. Select Token    │                      │
-//             ─ ─ ─│◀ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─│                      │
-//  4. Verify │     │                     │                      │
-//                  │  4. Yield PubKey    │                      │
-//            └ ─ ─▶│─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─▶                      │
-//                  │                     │                      │
-//                  │                     │  5. Send Reg Opts    │
-//                  │                     │─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─▶│─ ─ ─
-//                  │                     │                      │     │ 5. Verify
-//                  │                     │                      │         PubKey
-//                  │                     │                      │◀─ ─ ┘
-//                  │                     │                      │─ ─ ─
-//                  │                     │                      │     │ 6. Persist
-//                  │                     │                      │       Credential
-//                  │                     │                      │◀─ ─ ┘
-//                  │                     │                      │
-//                  │                     │                      │
-//
-// In this step, we are responding to the start reg(istration) request, and providing
-// the challenge to the browser.
-
-// TODO - Improve error handling and messages
-fn auth_user(cookie_content: String, session_table: &mut HashMap<Uuid, Instant>) -> bool {
-    let Ok(uuid) = Uuid::parse_str(&cookie_content) else {
-        info!("此用户Session不在表中");
-        return false;
-    };
-    let Some(expire) = session_table.get(&uuid) else {
-        info!("此用户Session已过期");
-        return false;
-    };
-    if *expire <= Instant::now() {
-        info!("此用户Session已过期");
-        return false;
-    }
-    session_table.insert(
-        uuid,
-        Instant::now() + Duration::from_secs(*SESSION_ACTIVE_TIME),
-    );
-    tracing::info!("valid cookie {}", uuid);
-    true
-}
-
-pub async fn start_register(
-    State(state): State<Arc<ServerState>>,
-    cookies: Cookies,
-    session: Session,
-    Path(username): Path<String>,
-) -> Result<Json<CreationChallengeResponse>, String> {
-    tracing::info!("开始注册");
-    // todo!("Auth User");
-    let Some(cookie_content) = cookies.get(&COOKIE_NAME) else {
-        tracing::info!("用户没有Cookie");
-        return Err(WebauthnError::AuthenticationFailure.to_string());
-    };
-    let mut session_table = state.session.lock().await;
-    if !auth_user(cookie_content.value().to_string(), &mut session_table) {
-        tracing::info!("用户Cookie无效，不在Session表中");
-        return Err(WebauthnError::AuthenticationFailure.to_string());
-    }
-    // Remove any previous registrations that may have occurred from the session.
-    let _ = session
-        .remove::<(String, Uuid, PasskeyRegistration)>("reg_state")
-        .await
-        .map_err(|_| WebauthnError::AuthenticationFailure.to_string())?;
-    let pool = &state.db;
-    let username_tmp = username.clone();
-    let (user_id, exclude_credentials): (Uuid, Option<Vec<CredentialID>>) =
-        match get_uid_by_name(username_tmp, pool).await {
-            Ok(uid) => {
-                let exclude_passkeys = get_credential_from_uid(uid, pool).await?;
-                (
-                    uid,
-                    Some(
-                        exclude_passkeys
-                            .into_iter()
-                            .map(|k| k.cred_id().clone())
-                            .collect(),
-                    ),
-                )
-            }
-            Err(_) => (Uuid::new_v4(), None),
-        };
-
-    let res = match state.webauthn.start_passkey_registration(
-        user_id,
-        &username,
-        &username,
-        exclude_credentials,
-    ) {
-        Ok((ccr, reg_state)) => {
-            // Note that due to the session store in use being a server side memory store, this is
-            // safe to store the reg_state into the session since it is not client controlled and
-            // not open to replay attacks. If this was a cookie store, this would be UNSAFE.
-            session
-                .insert("reg_state", (username, user_id, reg_state))
-                .await
-                .map_err(|_| WebauthnError::ChallengePersistenceError.to_string())?;
-            Json(ccr)
-        }
-        Err(e) => {
-            debug!("challenge_register -> {:?}", e);
-            return Err(WebauthnError::ChallengePersistenceError.to_string());
-        }
-    };
-    Ok(res)
-}
-
-// 3. The browser has completed it's steps and the user has created a public key
-// on their device. Now we have the registration options sent to us, and we need
-// to verify these and persist them.
-
-pub async fn finish_register(
-    State(state): State<Arc<ServerState>>,
-    session: Session,
-    Json(reg): Json<RegisterPublicKeyCredential>,
-) -> Result<impl IntoResponse, String> {
-    info!("完成注册...");
-    let pool = &state.db;
-    let Ok(Some((user_name, user_id, reg_state))) = session
-        .get::<(String, Uuid, PasskeyRegistration)>("reg_state")
-        .await
-    else {
-        return Err(WebauthnError::AuthenticationFailure.to_string()); //Corrupt Session
-    };
-
-    let _ = session
-        .remove::<(Uuid, PasskeyAuthentication)>("reg_state")
-        .await;
-
-    let res = match state.webauthn.finish_passkey_registration(&reg, &reg_state) {
-        Ok(key) => {
-            info!("Passkey 正常");
-            info!("检查用户是否存在");
-            // Check if the user_id already exists
-            let user_count = get_user_count_by_uid(user_id, pool).await?;
-            // If the user doesn't exist, insert them into the users table
-            if user_count == 0
-                && create_user_if_non_existent(user_id, user_name.to_string(), pool).await? != 1
-            {
-                return Err(WebauthnError::AuthenticationFailure.to_string());
-            }
-
-            // Insert the key into the auth table
-            if add_credential_by_id(user_id, &key, pool).await? != 1 {
-                error!("将用户凭据持久化时失败,rows_affected!=1");
-                return Err(WebauthnError::AuthenticationFailure.to_string());
-            }
-
-            StatusCode::OK
-        }
-        Err(e) => {
-            error!("challenge_register -> {:?}", e);
-            StatusCode::BAD_REQUEST
-        }
-    };
-
-    Ok(res)
-}
-
-// 4. Now that our public key has been registered, we can authenticate a user and verify
-// that they are the holder of that security token. The work flow is similar to registration.
-//
-//          ┌───────────────┐     ┌───────────────┐      ┌───────────────┐
-//          │ Authenticator │     │    Browser    │      │     Site      │
-//          └───────────────┘     └───────────────┘      └───────────────┘
-//                  │                     │                      │
-//                  │                     │     1. Start Auth    │
-//                  │                     │─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─▶│
-//                  │                     │                      │
-//                  │                     │     2. Challenge     │
-//                  │                     │◀ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┤
-//                  │                     │                      │
-//                  │  3. Select Token    │                      │
-//             ─ ─ ─│◀ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─│                      │
-//  4. Verify │     │                     │                      │
-//                  │    4. Yield Sig     │                      │
-//            └ ─ ─▶│─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─▶                      │
-//                  │                     │    5. Send Auth      │
-//                  │                     │        Opts          │
-//                  │                     │─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─▶│─ ─ ─
-//                  │                     │                      │     │ 5. Verify
-//                  │                     │                      │          Sig
-//                  │                     │                      │◀─ ─ ┘
-//                  │                     │                      │
-//                  │                     │                      │
-//
-// The user indicates the wish to start authentication and we need to provide a challenge.

src/services/mod.rs

@@ -133,6 +133,5 @@ pub fn check_otp(key_from_db: String, user_input_otp: String) -> bool {
     false
 }
 
-pub mod auth;
 pub mod cookie_services;
 pub mod gc_services;