appks/service/auth/login.rs

use argon2::{
    Argon2, PasswordHash,
    password_hash::{PasswordHasher, PasswordVerifier},
};
use serde::{Deserialize, Serialize};
use sqlx::Row;

use crate::error::AppError;
use crate::models::users::User;
use crate::service::AuthService;
use crate::session::Session;

#[derive(Deserialize, Serialize, Clone, Debug, utoipa::ToSchema)]
pub struct LoginParams {
    pub username: String,
    pub password: String,
    pub captcha: String,
    pub totp_code: Option<String>,
}

impl AuthService {
    pub const TOTP_KEY: &'static str = "totp_key";
    const TOTP_ATTEMPTS_PREFIX: &'static str = "auth:totp_attempts:";
    const TOTP_MAX_ATTEMPTS: u64 = 5;
    const TOTP_PENDING_TTL_SECS: u64 = 600;

    #[tracing::instrument(skip(self, params, context), fields(username = %params.username))]
    pub async fn auth_login(&self, params: LoginParams, context: Session) -> Result<(), AppError> {
        let login = params.username.trim().to_string();
        let totp_pending = context
            .get::<String>(Self::TOTP_KEY)
            .ok()
            .flatten()
            .is_some();
        if !totp_pending {
            self.auth_check_captcha(&context, params.captcha).await?;
        }
        let password = self.auth_rsa_decode(&context, params.password).await?;

        let user = match self.auth_find_user_by_username(&login).await {
            Ok(user) => user,
            Err(_) => match self.auth_find_user_by_email(&login).await {
                Ok(user) => user,
                Err(_) => {
                    let _ = Argon2::default().hash_password(
                        password.as_bytes(),
                        &argon2::password_hash::SaltString::generate(&mut rand::thread_rng()),
                    );
                    tracing::warn!(username = %login, "Login: user not found");
                    return Err(AppError::UserNotFound);
                }
            },
        };

        let row = sqlx::query(
            "SELECT user_id, password_hash, password_algo, password_salt, \
             must_change_password, password_updated_at, created_at, updated_at \
             FROM user_password WHERE user_id = $1",
        )
        .bind(user.id)
        .fetch_optional(self.ctx.db.reader())
        .await
        .map_err(AppError::Database)?;

        let row = row.ok_or(AppError::UserNotFound)?;
        let hash: String = row.try_get("password_hash").map_err(AppError::Database)?;

        let password_hash = PasswordHash::new(&hash).map_err(|_| AppError::UserNotFound)?;
        if Argon2::default()
            .verify_password(password.as_bytes(), &password_hash)
            .is_err()
        {
            tracing::warn!(username = %login, "Login: invalid password");
            return Err(AppError::UserNotFound);
        }

        let two_factor_enabled = self.auth_2fa_status_by_uid(user.id).await?.is_enabled;
        if two_factor_enabled {
            if let Some(totp_session_key) = context.get::<String>(Self::TOTP_KEY).ok().flatten() {
                let Some(ref totp_code) = params.totp_code else {
                    return Err(AppError::InvalidTwoFactorCode);
                };
                let attempts_key = format!("{}{}", Self::TOTP_ATTEMPTS_PREFIX, totp_session_key);
                let attempts = self.ctx.cache.get::<u64>(&attempts_key).unwrap_or(0);
                if attempts >= Self::TOTP_MAX_ATTEMPTS {
                    context.remove(Self::TOTP_KEY);
                    let _ = self.ctx.cache.delete(&totp_session_key);
                    let _ = self.ctx.cache.delete(&attempts_key);
                    return Err(AppError::InvalidTwoFactorCode);
                }

                if !self
                    .auth_2fa_verify_login(&context, user.id, totp_code)
                    .await?
                {
                    let next_attempts = attempts + 1;
                    if next_attempts >= Self::TOTP_MAX_ATTEMPTS {
                        context.remove(Self::TOTP_KEY);
                        let _ = self.ctx.cache.delete(&totp_session_key);
                        let _ = self.ctx.cache.delete(&attempts_key);
                    } else {
                        self.ctx
                            .cache
                            .set(
                                &attempts_key,
                                &next_attempts,
                                Some(std::time::Duration::from_secs(Self::TOTP_PENDING_TTL_SECS)),
                            )
                            .map_err(|e| AppError::InternalServerError(e.to_string()))?;
                    }
                    return Err(AppError::InvalidTwoFactorCode);
                }
                let _ = self.ctx.cache.delete(&attempts_key);
            } else {
                let totp_session_key = uuid::Uuid::new_v4().to_string();
                context
                    .insert(Self::TOTP_KEY, totp_session_key.clone())
                    .map_err(|_| AppError::InternalServerError("session insert failed".into()))?;
                self.ctx
                    .cache
                    .set(
                        &totp_session_key,
                        &user.id,
                        Some(std::time::Duration::from_secs(Self::TOTP_PENDING_TTL_SECS)),
                    )
                    .map_err(|e| AppError::InternalServerError(e.to_string()))?;
                tracing::info!(username = %login, "Login 2FA triggered");
                return Err(AppError::TwoFactorRequired);
            }
        } else if let Some(totp_session_key) = context.get::<String>(Self::TOTP_KEY).ok().flatten()
        {
            context.remove(Self::TOTP_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(&attempts_key);
        }

        sqlx::query("UPDATE \"user\" SET last_login_at = $1, updated_at = $1 WHERE id = $2")
            .bind(chrono::Utc::now())
            .bind(user.id)
            .execute(self.ctx.db.writer())
            .await
            .map_err(AppError::Database)?;

        context.renew();
        context.set_user(user.id);
        context.remove(Self::RSA_PRIVATE_KEY);
        context.remove(Self::RSA_PUBLIC_KEY);
        tracing::info!(user_uid = %user.id, username = %user.username, "User logged in");
        Ok(())
    }

    pub(crate) async fn auth_find_user_by_username(
        &self,
        username: &str,
    ) -> Result<User, AppError> {
        sqlx::query_as::<_, User>(
            "SELECT id, username, display_name, avatar_url, bio, status, role, visibility, \
             is_active, is_bot, last_login_at, created_at, updated_at, deleted_at \
             FROM \"user\" WHERE lower(username) = lower($1) AND is_active = true AND status = 'active' AND deleted_at IS NULL",
        )
        .bind(username)
        .fetch_optional(self.ctx.db.reader())
        .await
        .map_err(AppError::Database)?
        .ok_or(AppError::UserNotFound)
    }

    pub(crate) async fn auth_find_user_by_email(&self, email: &str) -> Result<User, AppError> {
        sqlx::query_as::<_, User>(
            "SELECT u.id, u.username, u.display_name, u.avatar_url, u.bio, u.status, u.role, \
             u.visibility, u.is_active, u.is_bot, u.last_login_at, u.created_at, u.updated_at, u.deleted_at \
             FROM \"user\" u \
             INNER JOIN user_mail e ON e.user_id = u.id \
             WHERE lower(e.email) = lower($1) AND e.is_verified = true AND u.is_active = true AND u.status = 'active' AND u.deleted_at IS NULL",
        )
        .bind(email)
        .fetch_optional(self.ctx.db.reader())
        .await
        .map_err(AppError::Database)?
        .ok_or(AppError::UserNotFound)
    }

    pub(crate) async fn auth_find_user_by_uid(&self, uid: uuid::Uuid) -> Result<User, AppError> {
        sqlx::query_as::<_, User>(
            "SELECT id, username, display_name, avatar_url, bio, status, role, visibility, \
             is_active, is_bot, last_login_at, created_at, updated_at, deleted_at \
             FROM \"user\" WHERE id = $1 AND is_active = true AND status = 'active' AND deleted_at IS NULL",
        )
        .bind(uid)
        .fetch_optional(self.ctx.db.reader())
        .await
        .map_err(AppError::Database)?
        .ok_or(AppError::UserNotFound)
    }
}