Skip to content
Clone
13aafa3
718 lines 25.5 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::{Code, MfaStatus as ProtoMfaStatus, MfaType as ProtoMfaType};
use crate::protos::auth::v1::mfa_service_server::{MfaService, MfaServiceServer};
use crate::protos::auth::v1::*;
use sea_orm::{ActiveModelTrait, ColumnTrait, EntityTrait, QueryFilter, Set};
use serde_json::json;
use tonic::{Request, Response, Status};
use crate::AuthService;
use crate::api::{context, convert, error, store};
use crate::errors::{AuthError, AuthResult};
#[derive(Clone)]
pub struct MfaApi {
svc: AuthService,
}
impl MfaApi {
pub fn new(svc: AuthService) -> Self {
Self { svc }
}
pub fn server(
self,
) -> tonic::service::interceptor::InterceptedService<MfaServiceServer<Self>, context::Interceptor>
{
let gate = context::Interceptor::new(self.svc.auth().rpc_key.clone());
MfaServiceServer::with_interceptor(self, gate)
}
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.svc.db().read())
.await?;
Ok(rows
.iter()
.filter(|row| row.expires_at.is_none_or(|expires| expires > now))
.count() as u32)
}
#[tracing::instrument(skip_all)]
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.svc.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(proto_type(&row.mfa_type)),
status: store::mfa_status_code(proto_status(&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.svc.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.svc.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 = recovery_code();
let normalized = normalize_recovery_code(&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.svc.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.svc.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()))]
async fn configure_totp_impl(
&self,
request: ConfigureTotpRequest,
request_id: Option<&str>,
) -> AuthResult<ConfigureTotpResponse> {
let db = self.svc.db();
let user = store::load_active_user(db, request.user_id).await?;
let key = self.svc.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()))]
async fn enable_totp_impl(
&self,
request: EnableTotpRequest,
request_id: Option<&str>,
) -> AuthResult<EnableTotpResponse> {
let db = self.svc.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.svc.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()))]
async fn disable_totp_impl(
&self,
request: DisableTotpRequest,
request_id: Option<&str>,
) -> AuthResult<DisableTotpResponse> {
let db = self.svc.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.svc.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()))]
async fn regenerate_recovery_codes_impl(
&self,
request: RegenerateRecoveryCodesRequest,
request_id: Option<&str>,
) -> AuthResult<RegenerateRecoveryCodesResponse> {
let db = self.svc.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,
})
}
#[tracing::instrument(skip_all, fields(request_id = %request_id.unwrap_or_default()))]
async fn verify_totp_impl(
&self,
request: VerifyTotpRequest,
request_id: Option<&str>,
) -> AuthResult<VerifyTotpResponse> {
let db = self.svc.db();
let user = match store::load_active_user(db, request.user_id).await {
Ok(user) => user,
Err(
AuthError::AccountUnavailable { .. }
| AuthError::LoginForbidden { .. }
| AuthError::UserNotFound,
) => {
return Ok(VerifyTotpResponse {
valid: false,
reason: Code::UserDisabled as i32,
matched_step: 0,
attempt_count: 0,
});
}
// 数据库/缓存故障不能伪装成「用户被禁用」,否则上游会把可用性故障当业务结果。
Err(error) => return Err(error),
};
let factor = user_mfa::Entity::find()
.filter(user_mfa::Column::UserId.eq(user.id))
.filter(user_mfa::Column::MfaType.eq(MfaType::Totp))
.filter(user_mfa::Column::MfaStatus.eq(MfaStatus::Enabled))
.filter(user_mfa::Column::RevokedAt.is_null())
.one(db.read())
.await?;
let Some(factor) = factor else {
return Ok(VerifyTotpResponse {
valid: false,
reason: Code::MfaNotConfigured as i32,
matched_step: 0,
attempt_count: 0,
});
};
match crate::share::mfa::consume_totp(
db,
&factor,
self.svc.auth().totp_encryption_key.as_deref(),
&request.code,
request_id,
)
.await
{
Ok(step) => {
if let Err(error) = store::audit(
db,
user.id,
Some(user.id),
UserEventType::Login,
Some(json!({ "action": "totp_verified" })),
request_id,
non_empty(&request.client_ip),
None,
)
.await
{
// TOTP 步长已推进,审计失败不能让客户端重试(重试只会得到 TotpStepReused)。
tracing::warn!(%error, user_id = user.id, "failed to record totp verification");
}
Ok(VerifyTotpResponse {
valid: true,
reason: Code::Unspecified as i32,
matched_step: step as i64,
attempt_count: 0,
})
}
Err(AuthError::TotpStepReused) => Ok(VerifyTotpResponse {
valid: false,
reason: Code::TotpStepReused as i32,
matched_step: 0,
attempt_count: 0,
}),
Err(AuthError::TotpCodeInvalid) => Ok(VerifyTotpResponse {
valid: false,
reason: Code::InvalidCredentials as i32,
matched_step: 0,
attempt_count: 0,
}),
// 其余错误(密钥缺失、解密失败、数据库故障)必须上抛,不能归为「凭据错误」。
Err(error) => Err(error),
}
}
#[tracing::instrument(skip_all, fields(request_id = %request_id.unwrap_or_default()))]
async fn redeem_recovery_code_impl(
&self,
request: RedeemRecoveryCodeRequest,
request_id: Option<&str>,
) -> AuthResult<RedeemRecoveryCodeResponse> {
let db = self.svc.db();
let user = match store::load_active_user(db, request.user_id).await {
Ok(user) => user,
Err(
AuthError::AccountUnavailable { .. }
| AuthError::LoginForbidden { .. }
| AuthError::UserNotFound,
) => {
return Ok(RedeemRecoveryCodeResponse {
valid: false,
reason: Code::UserDisabled as i32,
remaining_codes: 0,
});
}
Err(error) => return Err(error),
};
let remaining_before = self.remaining_codes(user.id).await?;
match crate::share::mfa::consume_recovery(db, user.id, &request.code, request_id).await {
Ok(()) => {}
Err(AuthError::RecoveryCodeInvalid) => {
return Ok(RedeemRecoveryCodeResponse {
valid: false,
reason: Code::InvalidCredentials as i32,
remaining_codes: remaining_before,
});
}
Err(error) => return Err(error),
}
store::audit(
db,
user.id,
Some(user.id),
UserEventType::Login,
Some(json!({ "action": "recovery_code_redeemed" })),
request_id,
non_empty(&request.client_ip),
None,
)
.await?;
Ok(RedeemRecoveryCodeResponse {
valid: true,
reason: Code::Unspecified as i32,
remaining_codes: remaining_before.saturating_sub(1),
})
}
}
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";
fn normalize_recovery_code(code: &str) -> String {
code.trim()
.to_uppercase()
.chars()
.filter(|value| *value != '-')
.collect()
}
fn recovery_code() -> String {
let left = crate::share::crypto::random_alnum(6);
let right = crate::share::crypto::random_alnum(6);
format!("{left}-{right}")
}
fn proto_type(value: &MfaType) -> ProtoMfaType {
match value {
MfaType::Totp => ProtoMfaType::Totp,
MfaType::Webauthn => ProtoMfaType::Webauthn,
MfaType::Sms => ProtoMfaType::Sms,
MfaType::Email => ProtoMfaType::Email,
}
}
fn proto_status(value: &MfaStatus) -> ProtoMfaStatus {
match value {
MfaStatus::Enabled => ProtoMfaStatus::Enabled,
MfaStatus::Disabled => ProtoMfaStatus::Disabled,
MfaStatus::Pending => ProtoMfaStatus::Pending,
MfaStatus::Revoked => ProtoMfaStatus::Revoked,
}
}
#[tonic::async_trait]
impl MfaService for MfaApi {
async fn get_mfa_status(
&self,
request: Request<GetMfaStatusRequest>,
) -> Result<Response<GetMfaStatusResponse>, Status> {
let meta = context::meta(&request).await?;
let request = request.into_inner();
let request_id = meta.request_id().map(str::to_owned);
let timer = crate::oltp::metrics::ActiveRpcGuard::new("auth.v1.MfaService");
let result = self
.status_impl(request)
.await
.map(Response::new)
.map_err(|failure| error::status(&failure, request_id.as_deref()));
timer.finish("GetMfaStatus", &result);
result
}
async fn verify_totp(
&self,
request: Request<VerifyTotpRequest>,
) -> Result<Response<VerifyTotpResponse>, Status> {
let meta = context::meta(&request).await?;
let request = request.into_inner();
let request_id = meta.request_id().map(str::to_owned);
let timer = crate::oltp::metrics::ActiveRpcGuard::new("auth.v1.MfaService");
let result = self
.verify_totp_impl(request, request_id.as_deref())
.await
.map(Response::new)
.map_err(|failure| error::status(&failure, request_id.as_deref()));
timer.finish("VerifyTotp", &result);
result
}
async fn redeem_recovery_code(
&self,
request: Request<RedeemRecoveryCodeRequest>,
) -> Result<Response<RedeemRecoveryCodeResponse>, Status> {
let meta = context::meta(&request).await?;
let request = request.into_inner();
let request_id = meta.request_id().map(str::to_owned);
let timer = crate::oltp::metrics::ActiveRpcGuard::new("auth.v1.MfaService");
let result = self
.redeem_recovery_code_impl(request, request_id.as_deref())
.await
.map(Response::new)
.map_err(|failure| error::status(&failure, request_id.as_deref()));
timer.finish("RedeemRecoveryCode", &result);
result
}
async fn configure_totp(
&self,
request: Request<ConfigureTotpRequest>,
) -> Result<Response<ConfigureTotpResponse>, Status> {
let meta = context::meta(&request).await?;
let request = request.into_inner();
let request_id = meta.request_id().map(str::to_owned);
let timer = crate::oltp::metrics::ActiveRpcGuard::new("auth.v1.MfaService");
let result = self
.configure_totp_impl(request, request_id.as_deref())
.await
.map(Response::new)
.map_err(|failure| error::status(&failure, request_id.as_deref()));
timer.finish("ConfigureTotp", &result);
result
}
async fn enable_totp(
&self,
request: Request<EnableTotpRequest>,
) -> Result<Response<EnableTotpResponse>, Status> {
let meta = context::meta(&request).await?;
let request = request.into_inner();
let request_id = meta.request_id().map(str::to_owned);
let timer = crate::oltp::metrics::ActiveRpcGuard::new("auth.v1.MfaService");
let result = self
.enable_totp_impl(request, request_id.as_deref())
.await
.map(Response::new)
.map_err(|failure| error::status(&failure, request_id.as_deref()));
timer.finish("EnableTotp", &result);
result
}
async fn disable_totp(
&self,
request: Request<DisableTotpRequest>,
) -> Result<Response<DisableTotpResponse>, Status> {
let meta = context::meta(&request).await?;
let request = request.into_inner();
let request_id = meta.request_id().map(str::to_owned);
let timer = crate::oltp::metrics::ActiveRpcGuard::new("auth.v1.MfaService");
let result = self
.disable_totp_impl(request, request_id.as_deref())
.await
.map(Response::new)
.map_err(|failure| error::status(&failure, request_id.as_deref()));
timer.finish("DisableTotp", &result);
result
}
async fn regenerate_recovery_codes(
&self,
request: Request<RegenerateRecoveryCodesRequest>,
) -> Result<Response<RegenerateRecoveryCodesResponse>, Status> {
let meta = context::meta(&request).await?;
let request = request.into_inner();
let request_id = meta.request_id().map(str::to_owned);
let timer = crate::oltp::metrics::ActiveRpcGuard::new("auth.v1.MfaService");
let result = self
.regenerate_recovery_codes_impl(request, request_id.as_deref())
.await
.map(Response::new)
.map_err(|failure| error::status(&failure, request_id.as_deref()));
timer.finish("RegenerateRecoveryCodes", &result);
result
}
}