auth/src/services/auth.rs
2024-03-14 00:24:04 +08:00

414 lines
17 KiB
Rust

#![allow(non_snake_case)]
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use axum::debug_handler;
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::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 {
return false;
};
let Some(expire) = session_table.get(&uuid) else {
return false;
};
if *expire <= Instant::now() {
return false;
}
session_table.insert(uuid, Instant::now()+Duration::from_secs(*SESSION_ACTIVE_TIME));
tracing::info!("valid cookie {}",uuid);
return true;
}
#[debug_handler]
pub async fn start_register(
State(state): State<Arc<ServerState>>,
cookies: Cookies,
session: Session,
Path(username): Path<String>,
) -> Result<Json<CreationChallengeResponse>, String> {
tracing::info!("Start register");
// todo!("Auth User");
let Some(cookie_content) = cookies.get(&COOKIE_NAME) else {
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){
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 = username.clone();
let (user_id, exclude_credentials): (Uuid, Option<Vec<CredentialID>>) =
match sqlx::query!("SELECT key FROM users WHERE name = $1;", username)
.fetch_optional(pool)
.await
.map_err(|_| WebauthnError::AuthenticationFailure.to_string())?
{
Some(record) => {
let uid = record.KEY.clone();
let records = sqlx::query!(
"SELECT credential FROM CREDENTIALS WHERE user_id = $1;",
uid
)
.fetch_all(pool)
.await
.map_err(|_| WebauthnError::AuthenticationFailure.to_string())?;
(
Uuid::parse_str(&record.KEY)
.map_err(|_| WebauthnError::AuthenticationFailure.to_string())?,
Some(
records
.iter()
.map(|record| serde_json::from_str::<Passkey>(&record.CREDENTIAL))
.collect::<Result<Vec<Passkey>, _>>()
.map_err(|_| WebauthnError::CredentialPersistenceError.to_string())?
.iter()
.map(|passkey| passkey.cred_id().clone())
.collect(),
),
)
}
None => (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!("Confirming registration....");
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 is okay");
let uid = &user_id.to_string();
let username = &user_name.to_string();
// Check if the user_id already exists
let record = sqlx::query!(
"SELECT COUNT(KEY) AS count FROM users WHERE KEY = $1;",
uid
)
.fetch_one(pool)
.await
.map_err(|_| WebauthnError::AuthenticationFailure.to_string())?;
// If the user doesn't exist, insert them into the users table
if record.count == 0
&& sqlx::query!(
"INSERT INTO users(KEY, NAME) VALUES($1, $2);",
uid,
username
)
.execute(pool)
.await
.map_err(|_| WebauthnError::UserNotPresent.to_string())?
.rows_affected()
!= 1
{
return Err(WebauthnError::AuthenticationFailure.to_string());
}
// Serialise the key
let serialised_key = serde_json::ser::to_string(&key)
.map_err(|_| WebauthnError::CredentialPersistenceError.to_string())?;
// Insert the key into the auth table
if sqlx::query!(
"INSERT INTO CREDENTIALS(user_id, credential) VALUES($1, $2);",
uid,
serialised_key
)
.execute(pool)
.await
.map_err(|_| WebauthnError::UserNotPresent.to_string())?
.rows_affected()
!= 1
{
return Err(WebauthnError::AuthenticationFailure.to_string());
}
StatusCode::OK
}
Err(e) => {
debug!("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.
pub async fn start_authentication(
State(state): State<Arc<ServerState>>,
session: Session,
Path(user_name): Path<String>,
) -> Result<impl IntoResponse, String> {
info!("Start Authentication");
let pool = &state.db;
// Remove any previous authentication that may have occurred from the session.
let _ = session.remove::<(Uuid, PasskeyAuthentication)>("auth_state");
let user_id = sqlx::query!("SELECT KEY FROM users WHERE NAME = $1;", user_name)
.fetch_one(pool)
.await
.map_err(|_| WebauthnError::UserNotPresent.to_string())?
.KEY;
let records = sqlx::query!(
"SELECT credential FROM CREDENTIALS WHERE user_id = $1;",
user_id
)
.fetch_all(pool)
.await
.map_err(|_| WebauthnError::AuthenticationFailure.to_string())?;
if records.is_empty() {
return Err(WebauthnError::CredentialNotFound.to_string());
}
let allow_credentials: Vec<Passkey> = records
.iter()
.map(|record| serde_json::de::from_str::<Passkey>(&record.CREDENTIAL))
.collect::<Result<Vec<Passkey>, _>>()
.map_err(|_| WebauthnError::CredentialPersistenceError.to_string())?;
let res = match state
.webauthn
.start_passkey_authentication(&allow_credentials)
{
Ok((rcr, auth_state)) => {
// Note that due to the session store in use being a server side memory store, this is
// safe to store the auth_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("auth_state", (&user_id, auth_state))
.await
.map_err(|_| WebauthnError::AuthenticationFailure.to_string())?;
Json(rcr)
}
Err(e) => {
debug!("challenge_authenticate -> {:?}", e);
return Err(WebauthnError::MismatchedChallenge.to_string());
}
};
Ok(res)
}
// 5. The browser and user have completed their part of the processing. Only in the
// case that the webauthn authenticate call returns Ok, is authentication considered
// a success. If the browser does not complete this call, or *any* error occurs,
// this is an authentication failure.
pub async fn finish_authentication(
State(state): State<Arc<ServerState>>,
session: Session,
Json(auth): Json<PublicKeyCredential>,
) -> Result<impl IntoResponse, String> {
let pool = &state.db;
let (user_id, auth_state): (Uuid, PasskeyAuthentication) = session
.get("auth_state")
.await
.unwrap()
.ok_or(WebauthnError::AuthenticationFailure.to_string())?;
let _ = session.remove::<(Uuid, PasskeyAuthentication)>("auth_state");
let res = match state
.webauthn
.finish_passkey_authentication(&auth, &auth_state)
{
Ok(auth_result) => {
let uid = user_id.clone().to_string();
let records = sqlx::query!(
"SELECT CREDENTIAL FROM CREDENTIALS WHERE USER_ID = $1;",
uid
)
.fetch_all(pool)
.await
.map_err(|_| WebauthnError::AuthenticatorDataMissingExtension.to_string())?;
if records.is_empty() {
return Err(WebauthnError::UserNotPresent.to_string());
}
for record in records {
let mut credential = serde_json::from_str::<Passkey>(&record.CREDENTIAL)
.map_err(|_| WebauthnError::CredentialExistCheckError.to_string())?;
if credential.cred_id() == auth_result.cred_id() {
credential.update_credential(&auth_result);
let credential = serde_json::to_string(&credential)
.map_err(|_| WebauthnError::CredentialPersistenceError.to_string())?;
if sqlx::query!(
"UPDATE CREDENTIALS SET credential = $1 WHERE user_id = $2 AND credential = $3;",
credential,
uid,
record.CREDENTIAL
)
.execute(pool)
.await.is_ok_and(|e|e.rows_affected() !=1 )
{
return Err(WebauthnError::AuthenticationFailure.to_string());
}
break;
}
}
let user_name = sqlx::query!(
"SELECT NAME FROM users WHERE KEY = $1;",
uid
)
.fetch_one(pool)
.await
.map_err(|_| WebauthnError::AuthenticationFailure.to_string())?;
// Add our own values to the session
session
.insert("user_id", user_id)
.await
.map_err(|_| WebauthnError::InvalidUserUniqueId.to_string())?;
session
.insert("user_name", user_name.NAME)
.await
.map_err(|_| WebauthnError::InvalidUsername.to_string())?;
StatusCode::OK
}
Err(e) => {
debug!("challenge_register -> {:?}", e);
StatusCode::BAD_REQUEST
}
};
info!("Authentication Successful!");
Ok(res)
}