Skip to content
Clone
1480255
438 lines 15.7 KB Rust
use crate::models::users::user_events::UserEventType;
use crate::models::users::user_mfa::{self, MfaStatus, MfaType};
use crate::models::users::user_mfa_backup_codes;
use crate::protos::auth::common::v1::{MfaStatus as ProtoMfaStatus, MfaType as ProtoMfaType};
use crate::protos::auth::v1::*;
use sea_orm::{ActiveModelTrait, ColumnTrait, EntityTrait, QueryFilter, Set};
use serde_json::json;
use crate::AuthService;
use crate::errors::{AuthError, AuthResult};
use crate::share::{auth_convert as convert, auth_store as store};
impl AuthService {
async fn remaining_codes(&self, user_id: i64) -> AuthResult<u32> {
let now = chrono::Utc::now();
let rows = user_mfa_backup_codes::Entity::find()
.filter(user_mfa_backup_codes::Column::UserId.eq(user_id))
.filter(user_mfa_backup_codes::Column::ConsumedAt.is_null())
.filter(user_mfa_backup_codes::Column::RevokedAt.is_null())
.all(self.db().read())
.await?;
Ok(rows
.iter()
.filter(|row| row.expires_at.is_none_or(|expires| expires > now))
.count() as u32)
}
#[tracing::instrument(skip_all)]
pub(crate) async fn status_impl(
&self,
request: GetMfaStatusRequest,
) -> AuthResult<GetMfaStatusResponse> {
let rows = user_mfa::Entity::find()
.filter(user_mfa::Column::UserId.eq(request.user_id))
.all(self.db().read())
.await?;
let mut login_requires_mfa = false;
let mut factors = Vec::with_capacity(rows.len());
for row in rows {
if row.mfa_status == MfaStatus::Enabled && row.revoked_at.is_none() {
login_requires_mfa = true;
}
factors.push(get_mfa_status_response::Factor {
factor_id: row.id.to_string(),
mfa_type: store::mfa_type_code({
let __inline = |value: &MfaType| -> ProtoMfaType {
match value {
MfaType::Totp => ProtoMfaType::Totp,
MfaType::Webauthn => ProtoMfaType::Webauthn,
MfaType::Sms => ProtoMfaType::Sms,
MfaType::Email => ProtoMfaType::Email,
}
};
__inline(&row.mfa_type)
}),
status: store::mfa_status_code({
let __inline = |value: &MfaStatus| -> ProtoMfaStatus {
match value {
MfaStatus::Enabled => ProtoMfaStatus::Enabled,
MfaStatus::Disabled => ProtoMfaStatus::Disabled,
MfaStatus::Pending => ProtoMfaStatus::Pending,
MfaStatus::Revoked => ProtoMfaStatus::Revoked,
}
};
__inline(&row.mfa_status)
}),
name: row.name,
has_secret: row.secret_ciphertext.is_some(),
last_used_step: row.last_totp_step.unwrap_or(0),
verified_at: convert::datetime_opt(row.verified_at),
last_used_at: convert::datetime_opt(row.mfa_last_used_at),
revoked_at: convert::datetime_opt(row.revoked_at),
});
}
let (recovery_codes_total, recovery_codes_used) =
self.recovery_code_stats(request.user_id).await?;
Ok(GetMfaStatusResponse {
login_requires_mfa,
factors,
recovery_codes_remaining: self.remaining_codes(request.user_id).await?,
recovery_codes_total,
recovery_codes_used,
})
}
async fn recovery_code_stats(&self, user_id: i64) -> AuthResult<(u32, u32)> {
let rows = user_mfa_backup_codes::Entity::find()
.filter(user_mfa_backup_codes::Column::UserId.eq(user_id))
.filter(user_mfa_backup_codes::Column::RevokedAt.is_null())
.all(self.db().read())
.await?;
let used = rows.iter().filter(|row| row.consumed_at.is_some()).count();
Ok((rows.len() as u32, used as u32))
}
async fn revoke_recovery_codes(&self, user_id: i64) -> AuthResult<u64> {
let now = store::now();
let result = user_mfa_backup_codes::Entity::update_many()
.col_expr(
user_mfa_backup_codes::Column::RevokedAt,
sea_orm::sea_query::Expr::value(now),
)
.filter(user_mfa_backup_codes::Column::UserId.eq(user_id))
.filter(user_mfa_backup_codes::Column::RevokedAt.is_null())
.exec(self.db().write())
.await?;
Ok(result.rows_affected)
}
async fn issue_recovery_codes(&self, user_id: i64, count: usize) -> AuthResult<Vec<String>> {
self.revoke_recovery_codes(user_id).await?;
let batch_id = uuid::Uuid::new_v4();
let expires_at = store::now() + chrono::Duration::days(RECOVERY_CODE_TTL_DAYS);
let mut codes = Vec::with_capacity(count);
for _ in 0..count {
let code = {
let __inline = || -> String {
let left = crate::share::crypto::random_alnum(6);
let right = crate::share::crypto::random_alnum(6);
format!("{left}-{right}")
};
__inline()
};
let normalized = {
let __inline = |code: &str| -> String {
code.trim()
.to_uppercase()
.chars()
.filter(|value| *value != '-')
.collect()
};
__inline(&code)
};
let code_hash = crate::share::crypto::sha256_hex(normalized.as_bytes());
user_mfa_backup_codes::ActiveModel {
id: Set(uuid::Uuid::new_v4()),
user_id: Set(user_id),
batch_id: Set(batch_id),
code_hash: Set(code_hash),
hash_key_id: Set(RECOVERY_CODE_HASH_KEY_ID.to_owned()),
created_at: Set(store::now()),
expires_at: Set(Some(expires_at)),
consumed_at: Set(None),
revoked_at: Set(None),
version: Set(1),
}
.insert(self.db().write())
.await?;
codes.push(code);
}
Ok(codes)
}
async fn verify_password_hash(&self, user_id: i64, client_digest: &str) -> AuthResult<()> {
let user = store::load_active_user(self.db(), user_id).await?;
let stored = user.password_hash.as_str();
if stored.is_empty() {
return Err(AuthError::InvalidPassword);
}
let check =
crate::share::password::verify(client_digest.to_owned(), stored.to_owned(), None)
.await?;
if check == crate::share::password::PasswordCheck::Mismatch {
return Err(AuthError::InvalidPassword);
}
Ok(())
}
fn enabled_factor<'a>(&self, rows: &'a [user_mfa::Model]) -> Option<&'a user_mfa::Model> {
rows.iter().find(|row| {
row.mfa_type == MfaType::Totp
&& row.mfa_status == MfaStatus::Enabled
&& row.revoked_at.is_none()
})
}
#[tracing::instrument(skip_all, fields(request_id = %request_id.unwrap_or_default()))]
pub(crate) async fn configure_totp_impl(
&self,
request: ConfigureTotpRequest,
request_id: Option<&str>,
) -> AuthResult<ConfigureTotpResponse> {
let db = self.db();
let user = store::load_active_user(db, request.user_id).await?;
let key = self.auth().totp_encryption_key.as_deref();
let rows = user_mfa::Entity::find()
.filter(user_mfa::Column::UserId.eq(user.id))
.filter(user_mfa::Column::RevokedAt.is_null())
.all(db.read())
.await?;
if self.enabled_factor(&rows).is_some() {
return Err(AuthError::TotpAlreadyEnabled);
}
let now = store::now();
for row in rows.iter() {
user_mfa::ActiveModel {
id: Set(row.id),
revoked_at: Set(Some(now)),
updated_at: Set(now),
version: Set(row.version + 1),
..Default::default()
}
.update(db.write())
.await?;
}
let secret = crate::share::totp::generate_secret();
let ciphertext = crate::share::totp::encrypt_secret(key, &secret)?;
let uri = crate::share::totp::otpauth_uri(&secret, &user.username);
let name = request.name.trim();
let name = if name.is_empty() { "TOTP" } else { name };
let inserted = user_mfa::ActiveModel {
id: Set(uuid::Uuid::new_v4()),
user_id: Set(user.id),
mfa_type: Set(MfaType::Totp),
mfa_status: Set(MfaStatus::Pending),
name: Set(name.to_owned()),
secret_ciphertext: Set(Some(ciphertext)),
secret_key_id: Set(None),
last_totp_step: Set(None),
mfa_last_used_at: Set(None),
created_at: Set(now),
updated_at: Set(now),
version: Set(1),
verified_at: Set(None),
revoked_at: Set(None),
}
.insert(db.write())
.await?;
store::audit(
db,
user.id,
Some(user.id),
UserEventType::TwoFactorEnabled,
Some(json!({ "action": "totp_configured", "factor_id": inserted.id })),
request_id,
non_empty(&request.client_ip),
None,
)
.await?;
Ok(ConfigureTotpResponse {
factor_id: inserted.id.to_string(),
secret,
otpauth_uri: uri,
})
}
#[tracing::instrument(skip_all, fields(request_id = %request_id.unwrap_or_default()))]
pub(crate) async fn enable_totp_impl(
&self,
request: EnableTotpRequest,
request_id: Option<&str>,
) -> AuthResult<EnableTotpResponse> {
let db = self.db();
let user = store::load_active_user(db, request.user_id).await?;
self.verify_password_hash(user.id, &request.password_hash)
.await?;
let rows = user_mfa::Entity::find()
.filter(user_mfa::Column::UserId.eq(user.id))
.filter(user_mfa::Column::RevokedAt.is_null())
.all(db.read())
.await?;
if self.enabled_factor(&rows).is_some() {
return Err(AuthError::TotpAlreadyEnabled);
}
let pending = rows
.iter()
.find(|row| row.mfa_type == MfaType::Totp && row.mfa_status == MfaStatus::Pending)
.ok_or(AuthError::TotpNotConfigured)?;
let key = self.auth().totp_encryption_key.as_deref();
let ciphertext = pending
.secret_ciphertext
.as_deref()
.ok_or(AuthError::TotpNotConfigured)?;
let secret = crate::share::totp::decrypt_secret(key, ciphertext)?;
let (step, value) = crate::share::totp::verify_totp_code(&secret, &request.code, None)?;
let _ = step;
let now = store::now();
user_mfa::ActiveModel {
id: Set(pending.id),
mfa_status: Set(MfaStatus::Enabled),
verified_at: Set(Some(now)),
last_totp_step: Set(Some(value)),
mfa_last_used_at: Set(Some(now)),
updated_at: Set(now),
version: Set(pending.version + 1),
..Default::default()
}
.update(db.write())
.await?;
let codes = self
.issue_recovery_codes(user.id, RECOVERY_CODE_COUNT)
.await?;
store::audit(
db,
user.id,
Some(user.id),
UserEventType::TwoFactorEnabled,
Some(json!({ "action": "totp_enabled", "recovery_codes": codes.len() })),
request_id,
non_empty(&request.client_ip),
None,
)
.await?;
Ok(EnableTotpResponse {
recovery_codes: codes,
enabled_at: convert::datetime_opt(Some(now)),
})
}
#[tracing::instrument(skip_all, fields(request_id = %request_id.unwrap_or_default()))]
pub(crate) async fn disable_totp_impl(
&self,
request: DisableTotpRequest,
request_id: Option<&str>,
) -> AuthResult<DisableTotpResponse> {
let db = self.db();
let user = store::load_active_user(db, request.user_id).await?;
let rows = user_mfa::Entity::find()
.filter(user_mfa::Column::UserId.eq(user.id))
.filter(user_mfa::Column::RevokedAt.is_null())
.all(db.read())
.await?;
let Some(factor) = self.enabled_factor(&rows).cloned() else {
return Ok(DisableTotpResponse { disabled: false });
};
let code = request.code.trim();
let key = self.auth().totp_encryption_key.as_deref();
let consumed = crate::share::mfa::consume_totp(db, &factor, key, code, None)
.await
.is_ok();
let consumed = match consumed {
true => true,
false => crate::share::mfa::consume_recovery(db, user.id, code, None)
.await
.is_ok(),
};
if !consumed {
return Err(AuthError::TotpCodeInvalid);
}
let now = store::now();
user_mfa::ActiveModel {
id: Set(factor.id),
revoked_at: Set(Some(now)),
updated_at: Set(now),
version: Set(factor.version + 1),
..Default::default()
}
.update(db.write())
.await?;
self.revoke_recovery_codes(user.id).await?;
store::audit(
db,
user.id,
Some(user.id),
UserEventType::TwoFactorDisabled,
Some(json!({ "action": "totp_disabled", "factor_id": factor.id })),
request_id,
non_empty(&request.client_ip),
None,
)
.await?;
Ok(DisableTotpResponse { disabled: true })
}
#[tracing::instrument(skip_all, fields(request_id = %request_id.unwrap_or_default()))]
pub(crate) async fn regenerate_recovery_codes_impl(
&self,
request: RegenerateRecoveryCodesRequest,
request_id: Option<&str>,
) -> AuthResult<RegenerateRecoveryCodesResponse> {
let db = self.db();
let user = store::load_active_user(db, request.user_id).await?;
self.verify_password_hash(user.id, &request.password_hash)
.await?;
let rows = user_mfa::Entity::find()
.filter(user_mfa::Column::UserId.eq(user.id))
.filter(user_mfa::Column::RevokedAt.is_null())
.all(db.read())
.await?;
if self.enabled_factor(&rows).is_none() {
return Err(AuthError::TotpNotConfigured);
}
let codes = self
.issue_recovery_codes(user.id, RECOVERY_CODE_COUNT)
.await?;
store::audit(
db,
user.id,
Some(user.id),
UserEventType::TwoFactorEnabled,
Some(json!({ "action": "recovery_codes_regenerated", "count": codes.len() })),
request_id,
non_empty(&request.client_ip),
None,
)
.await?;
Ok(RegenerateRecoveryCodesResponse {
recovery_codes: codes,
})
}
}
fn non_empty(value: &str) -> Option<String> {
let trimmed = value.trim();
if trimmed.is_empty() {
None
} else {
Some(trimmed.to_owned())
}
}
const RECOVERY_CODE_COUNT: usize = 10;
const RECOVERY_CODE_TTL_DAYS: i64 = 365;
const RECOVERY_CODE_HASH_KEY_ID: &str = "v1";