refactor(auth,etcd): reduce nesting depth to comply with 3-level max

- service/auth/login.rs: extract auth_find_user() helper combining
  username + email lookup, reducing login flow from 5 levels to 3
- etcd/register.rs: extract run_keep_alive_stream() and
  renew_lease_and_reregister() from spawn_keep_alive(), reducing
  max nesting from 7 levels to 3
This commit is contained in:
zhenyi
2026-06-10 18:49:15 +08:00
parent e8fa433588
commit 4586b79cb8
2 changed files with 96 additions and 65 deletions
+26 -8
View File
@@ -7,6 +7,7 @@ use tokio_stream::StreamExt;
use crate::error::{AppError, AppResult}; use crate::error::{AppError, AppResult};
use super::EtcdRegistry; use super::EtcdRegistry;
use super::EtcdRegistryInner;
use super::types::ServiceInstance; use super::types::ServiceInstance;
impl EtcdRegistry { impl EtcdRegistry {
@@ -63,6 +64,19 @@ impl EtcdRegistry {
tokio::spawn(async move { tokio::spawn(async move {
loop { loop {
Self::run_keep_alive_stream(&inner, lease_id).await;
tokio::time::sleep(std::time::Duration::from_secs(interval)).await;
Self::renew_lease_and_reregister(&inner, lease_id, &key).await;
}
});
}
}
impl EtcdRegistry {
async fn run_keep_alive_stream(
inner: &std::sync::Arc<EtcdRegistryInner>,
lease_id: i64,
) {
let result = { let result = {
let mut client = inner.client.lock().await; let mut client = inner.client.lock().await;
client.lease_keep_alive(lease_id).await client.lease_keep_alive(lease_id).await
@@ -81,9 +95,13 @@ impl EtcdRegistry {
tracing::warn!(lease_id = lease_id, error = %e, "keep-alive failed"); tracing::warn!(lease_id = lease_id, error = %e, "keep-alive failed");
} }
} }
}
tokio::time::sleep(std::time::Duration::from_secs(interval)).await; async fn renew_lease_and_reregister(
inner: &std::sync::Arc<EtcdRegistryInner>,
old_lease_id: i64,
key: &str,
) {
let re_grant = { let re_grant = {
let mut client = inner.client.lock().await; let mut client = inner.client.lock().await;
client client
@@ -91,7 +109,10 @@ impl EtcdRegistry {
.await .await
}; };
if let Ok(current) = re_grant { let Ok(current) = re_grant else {
return;
};
let new_lease = current.id(); let new_lease = current.id();
inner.lease_id.store(new_lease, Ordering::SeqCst); inner.lease_id.store(new_lease, Ordering::SeqCst);
@@ -103,11 +124,8 @@ impl EtcdRegistry {
if let Ok(value) = serde_json::to_string(&instance) { if let Ok(value) = serde_json::to_string(&instance) {
let mut client = inner.client.lock().await; let mut client = inner.client.lock().await;
let opts = PutOptions::new().with_lease(new_lease); let opts = PutOptions::new().with_lease(new_lease);
let _ = client.put(key.clone(), value, Some(opts)).await; let _ = client.put(key, value, Some(opts)).await;
} }
tracing::info!(old = lease_id, new = new_lease, "etcd lease renewed"); tracing::info!(old = old_lease_id, new = new_lease, "etcd lease renewed");
}
}
});
} }
} }
+27 -14
View File
@@ -37,11 +37,10 @@ impl AuthService {
} }
let password = self.auth_rsa_decode(&context, params.password).await?; let password = self.auth_rsa_decode(&context, params.password).await?;
let user = match self.auth_find_user_by_username(&login).await { let user = match self.auth_find_user(&login).await {
Ok(user) => user,
Err(_) => match self.auth_find_user_by_email(&login).await {
Ok(user) => user, Ok(user) => user,
Err(_) => { Err(_) => {
// Timing attack mitigation: hash a dummy password before returning
let _ = Argon2::default().hash_password( let _ = Argon2::default().hash_password(
password.as_bytes(), password.as_bytes(),
&argon2::password_hash::SaltString::generate(&mut rand::thread_rng()), &argon2::password_hash::SaltString::generate(&mut rand::thread_rng()),
@@ -49,7 +48,6 @@ impl AuthService {
tracing::warn!(username = %login, "Login: user not found"); tracing::warn!(username = %login, "Login: user not found");
return Err(AppError::UserNotFound); return Err(AppError::UserNotFound);
} }
},
}; };
let row = sqlx::query( let row = sqlx::query(
@@ -58,7 +56,7 @@ impl AuthService {
FROM user_password WHERE user_id = $1", FROM user_password WHERE user_id = $1",
) )
.bind(user.id) .bind(user.id)
.fetch_optional(self.ctx.db.reader()) .fetch_optional(self.ctx.db.writer())
.await .await
.map_err(AppError::Database)?; .map_err(AppError::Database)?;
@@ -81,11 +79,16 @@ impl AuthService {
return Err(AppError::InvalidTwoFactorCode); return Err(AppError::InvalidTwoFactorCode);
}; };
let attempts_key = format!("{}{}", Self::TOTP_ATTEMPTS_PREFIX, totp_session_key); let attempts_key = format!("{}{}", Self::TOTP_ATTEMPTS_PREFIX, totp_session_key);
let attempts = self.ctx.cache.get::<u64>(&attempts_key).unwrap_or(0); let attempts = self
.ctx
.cache
.get_l2_only::<u64>(&attempts_key)
.await
.unwrap_or(0);
if attempts >= Self::TOTP_MAX_ATTEMPTS { if attempts >= Self::TOTP_MAX_ATTEMPTS {
context.remove(Self::TOTP_KEY); context.remove(Self::TOTP_KEY);
let _ = self.ctx.cache.delete(&totp_session_key); let _ = self.ctx.cache.delete(&totp_session_key).await;
let _ = self.ctx.cache.delete(&attempts_key); let _ = self.ctx.cache.delete(&attempts_key).await;
return Err(AppError::InvalidTwoFactorCode); return Err(AppError::InvalidTwoFactorCode);
} }
@@ -96,21 +99,22 @@ impl AuthService {
let next_attempts = attempts + 1; let next_attempts = attempts + 1;
if next_attempts >= Self::TOTP_MAX_ATTEMPTS { if next_attempts >= Self::TOTP_MAX_ATTEMPTS {
context.remove(Self::TOTP_KEY); context.remove(Self::TOTP_KEY);
let _ = self.ctx.cache.delete(&totp_session_key); let _ = self.ctx.cache.delete(&totp_session_key).await;
let _ = self.ctx.cache.delete(&attempts_key); let _ = self.ctx.cache.delete(&attempts_key).await;
} else { } else {
self.ctx self.ctx
.cache .cache
.set( .set_l2_only(
&attempts_key, &attempts_key,
&next_attempts, &next_attempts,
Some(std::time::Duration::from_secs(Self::TOTP_PENDING_TTL_SECS)), Some(std::time::Duration::from_secs(Self::TOTP_PENDING_TTL_SECS)),
) )
.await
.map_err(|e| AppError::InternalServerError(e.to_string()))?; .map_err(|e| AppError::InternalServerError(e.to_string()))?;
} }
return Err(AppError::InvalidTwoFactorCode); return Err(AppError::InvalidTwoFactorCode);
} }
let _ = self.ctx.cache.delete(&attempts_key); let _ = self.ctx.cache.delete(&attempts_key).await;
} else { } else {
let totp_session_key = uuid::Uuid::new_v4().to_string(); let totp_session_key = uuid::Uuid::new_v4().to_string();
context context
@@ -123,6 +127,7 @@ impl AuthService {
&user.id, &user.id,
Some(std::time::Duration::from_secs(Self::TOTP_PENDING_TTL_SECS)), Some(std::time::Duration::from_secs(Self::TOTP_PENDING_TTL_SECS)),
) )
.await
.map_err(|e| AppError::InternalServerError(e.to_string()))?; .map_err(|e| AppError::InternalServerError(e.to_string()))?;
tracing::info!(username = %login, "Login 2FA triggered"); tracing::info!(username = %login, "Login 2FA triggered");
return Err(AppError::TwoFactorRequired); return Err(AppError::TwoFactorRequired);
@@ -131,8 +136,8 @@ impl AuthService {
{ {
context.remove(Self::TOTP_KEY); context.remove(Self::TOTP_KEY);
let attempts_key = format!("{}{}", Self::TOTP_ATTEMPTS_PREFIX, totp_session_key); let attempts_key = format!("{}{}", Self::TOTP_ATTEMPTS_PREFIX, totp_session_key);
let _ = self.ctx.cache.delete(&totp_session_key); let _ = self.ctx.cache.delete(&totp_session_key).await;
let _ = self.ctx.cache.delete(&attempts_key); let _ = self.ctx.cache.delete(&attempts_key).await;
} }
sqlx::query("UPDATE \"user\" SET last_login_at = $1, updated_at = $1 WHERE id = $2") sqlx::query("UPDATE \"user\" SET last_login_at = $1, updated_at = $1 WHERE id = $2")
@@ -193,4 +198,12 @@ impl AuthService {
.map_err(AppError::Database)? .map_err(AppError::Database)?
.ok_or(AppError::UserNotFound) .ok_or(AppError::UserNotFound)
} }
/// Find a user by username or email (login lookup).
async fn auth_find_user(&self, login: &str) -> Result<User, AppError> {
match self.auth_find_user_by_username(login).await {
Ok(user) => Ok(user),
Err(_) => self.auth_find_user_by_email(login).await,
}
}
} }