29cc961
use models::users::user;
use models::users::user_recovery_codes;
use sea_orm::{
ActiveModelTrait, ColumnTrait, EntityTrait, IntoActiveModel, QueryFilter, Set, TransactionTrait,
};
use serde::{Deserialize, Serialize};
use time::OffsetDateTime;
use uuid::Uuid;
use super::AuthService;
use crate::share::crypto::{aes_decrypt, aes_encrypt, sha256_b64url};
use crate::errors::AuthError;
use crate::share::totp::{generate_secret, otpauth_uri, verify_totp_code};
const RECOVERY_CODE_COUNT: usize = 10;
#[derive(Debug, Clone, Serialize, Deserialize, utoipa::ToSchema)]
pub struct TotpStatusResp {
pub enabled: bool,
}
#[derive(Debug, Clone, Serialize, Deserialize, utoipa::ToSchema)]
pub struct TotpSetupResp {
pub secret: String,
pub otpauth_uri: String,
}
#[derive(Debug, Clone, Serialize, Deserialize, utoipa::ToSchema)]
pub struct TotpCodeCmd {
pub code: String,
}
#[derive(Debug, Clone, Serialize, Deserialize, utoipa::ToSchema)]
pub struct RecoveryCodesResp {
pub total: usize,
pub used: usize,
pub codes: Vec<String>,
}
fn recovery_code() -> String {
let segment = || {
let bytes = crate::share::crypto::random_bytes(3);
crate::share::totp::base32_encode(&bytes)
};
format!("{}-{}", segment(), segment())
}
impl AuthService {
pub async fn totp_status(&self, uid: Uuid) -> Result<TotpStatusResp, AuthError> {
let Some(user) = user::Entity::find_by_id(uid).one(self.db().read()).await? else {
return Err(AuthError::NotFound);
};
Ok(TotpStatusResp {
enabled: user.totp_enabled,
})
}
pub async fn totp_setup(&self, uid: Uuid) -> Result<TotpSetupResp, AuthError> {
let db = self.db();
let Some(user) = user::Entity::find_by_id(uid).one(db.write()).await? else {
return Err(AuthError::NotFound);
};
if user.totp_enabled {
return Err(AuthError::BadRequest(
"two-factor authentication is already enabled; disable it first".to_string(),
));
}
let key = self
.config()
.auth
.totp_encryption_key
.as_deref()
.ok_or_else(|| {
AuthError::Internal("totp encryption key not configured".to_string())
})?;
let secret = generate_secret();
let encrypted = aes_encrypt(key, &secret)?;
let otpauth = otpauth_uri(&secret, &user.username);
let now = OffsetDateTime::now_utc();
let mut active = user.into_active_model();
active.totp_encrypted_secret = Set(Some(encrypted));
active.totp_enabled = Set(false);
active.totp_last_counter = Set(None);
active.updated_at = Set(now);
active.update(db.write()).await?;
Ok(TotpSetupResp {
otpauth_uri: otpauth,
secret,
})
}
pub async fn totp_enable(&self, uid: Uuid, cmd: TotpCodeCmd) -> Result<(), AuthError> {
let db = self.db();
let Some(user) = user::Entity::find_by_id(uid).one(db.write()).await? else {
return Err(AuthError::NotFound);
};
let key = self
.config()
.auth
.totp_encryption_key
.as_deref()
.ok_or_else(|| {
AuthError::Internal("totp encryption key not configured".to_string())
})?;
let stored = user.totp_encrypted_secret.as_deref().ok_or_else(|| {
AuthError::BadRequest("run totp setup first".to_string())
})?;
let secret = aes_decrypt(key, stored)?;
let (_, counter) = verify_totp_code(&secret, &cmd.code, None)?;
let now = OffsetDateTime::now_utc();
let mut active = user.into_active_model();
active.totp_enabled = Set(true);
active.totp_last_used_at = Set(Some(now));
active.totp_last_counter = Set(Some(counter));
active.updated_at = Set(now);
active.update(db.write()).await?;
Ok(())
}
pub async fn totp_disable(&self, uid: Uuid, cmd: TotpCodeCmd) -> Result<(), AuthError> {
let db = self.db();
let Some(user) = user::Entity::find_by_id(uid).one(db.write()).await? else {
return Err(AuthError::NotFound);
};
let key = self
.config()
.auth
.totp_encryption_key
.as_deref()
.ok_or_else(|| {
AuthError::Internal("totp encryption key not configured".to_string())
})?;
let stored = user.totp_encrypted_secret.as_deref().ok_or_else(|| {
AuthError::BadRequest("two-factor authentication is not enabled".to_string())
})?;
let secret = aes_decrypt(key, stored)?;
verify_totp_code(&secret, &cmd.code, user.totp_last_counter)?;
let now = OffsetDateTime::now_utc();
TransactionTrait::transaction(db.primary(), |txn| {
Box::pin(async move {
let mut active = user.into_active_model();
active.totp_enabled = Set(false);
active.totp_encrypted_secret = Set(None);
active.totp_last_used_at = Set(None);
active.totp_last_counter = Set(None);
active.updated_at = Set(now);
active.update(txn).await?;
user_recovery_codes::Entity::delete_many()
.filter(user_recovery_codes::Column::UserId.eq(uid))
.exec(txn)
.await?;
Ok::<(), sea_orm::DbErr>(())
})
})
.await
.map_err(AuthError::from)
}
pub async fn recovery_codes(&self, uid: Uuid) -> Result<RecoveryCodesResp, AuthError> {
let db = self.db();
let rows = user_recovery_codes::Entity::find()
.filter(user_recovery_codes::Column::UserId.eq(uid))
.all(db.read())
.await?;
let used = rows.iter().filter(|row| row.used_at.is_some()).count();
Ok(RecoveryCodesResp {
total: rows.len(),
used,
codes: Vec::new(),
})
}
pub async fn regenerate_recovery_codes(&self, uid: Uuid) -> Result<RecoveryCodesResp, AuthError> {
let db = self.db();
let now = OffsetDateTime::now_utc();
let mut generated = Vec::with_capacity(RECOVERY_CODE_COUNT);
for _ in 0..RECOVERY_CODE_COUNT {
generated.push(recovery_code());
}
let codes = generated.clone();
TransactionTrait::transaction(db.primary(), |txn| {
Box::pin(async move {
user_recovery_codes::Entity::delete_many()
.filter(user_recovery_codes::Column::UserId.eq(uid))
.exec(txn)
.await?;
for code in &generated {
let row = user_recovery_codes::ActiveModel {
id: Set(Uuid::now_v7()),
user_id: Set(uid),
code_hash: Set(sha256_b64url(code.as_bytes())),
used_at: Set(None),
expires_at: Set(None),
created_at: Set(now),
};
row.insert(txn).await?;
}
Ok::<(), sea_orm::DbErr>(())
})
})
.await?;
Ok(RecoveryCodesResp {
total: codes.len(),
used: 0,
codes,
})
}
}