gitdataai/lib/service/auth/login.rs

use argon2::{
    Argon2, PasswordHash,
    password_hash::{PasswordHasher, PasswordVerifier},
};
use db::sqlx;
use model::users::{UserModel, user_pass::UserPasswordModel};
use serde::{Deserialize, Serialize};
use session::Session;

use crate::{AppService, error::AppError};

#[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 AppService {
    pub const TOTP_KEY: &'static str = "totp_key";
    #[tracing::instrument(skip(self, params, context), fields(username = %params.username, ip = ?context.ip_address()))]
    pub async fn auth_login(
        &self,
        params: LoginParams,
        context: Session,
    ) -> Result<(), AppError> {
        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(&params.username).await
        {
            Ok(user) => user,
            Err(_) => {
                match self.auth_find_user_by_email(&params.username).await {
                    Ok(user) => user,
                    Err(_) => {
                        let _ = Argon2::default()
                            .hash_password(password.as_bytes());
                        self.metrics
                            .auth_login_total
                            .with_label_values(&["user_not_found"])
                            .inc();
                        return Err(AppError::UserNotFound);
                    }
                }
            }
        };

        let user_password = sqlx::query_as::<_, UserPasswordModel>(
            "SELECT \"user\", hash, salt, is_active, reason, created_at, updated_at \
             FROM user_password WHERE \"user\" = $1 AND is_active = true",
        )
        .bind(user.id)
        .fetch_optional(self.db.reader())
        .await
        .map_err(|e| AppError::DatabaseError(e.to_string()))?
        .ok_or(AppError::UserNotFound)?;

        let password_hash = PasswordHash::new(&user_password.hash)
            .map_err(|_| AppError::UserNotFound)?;
        if Argon2::default()
            .verify_password(password.as_bytes(), &password_hash)
            .is_err()
        {
            tracing::warn!(username = %params.username, ip = ?context.ip_address(), "Login failed: invalid password");
            self.metrics
                .auth_login_total
                .with_label_values(&["invalid_password"])
                .inc();
            return Err(AppError::UserNotFound);
        }

        if context
            .get::<String>(Self::TOTP_KEY)
            .ok()
            .flatten()
            .is_some()
        {
            if let Some(ref totp_code) = params.totp_code {
                if !self.auth_2fa_verify_login(&context, totp_code).await? {
                    return Err(AppError::InvalidTwoFactorCode);
                }
            } else {
                return Err(AppError::InvalidTwoFactorCode);
            }
        } else if self.auth_2fa_status_by_uid(user.id).await?.is_enabled {
            let totp_session_key = uuid::Uuid::new_v4().to_string();
            context
                .insert(Self::TOTP_KEY, totp_session_key.clone())
                .map_err(|_| AppError::InternalError)?;
            self.cache
                .set(&totp_session_key, &user.id)
                .await
                .map_err(|e| AppError::InternalServerError(e.to_string()))?;
            tracing::info!(username = %params.username, ip = ?context.ip_address(), "Login 2FA triggered");
            self.metrics
                .auth_login_total
                .with_label_values(&["2fa_required"])
                .inc();
            self.metrics
                .auth_2fa_triggered_total
                .with_label_values(&[])
                .inc();
            return Err(AppError::TwoFactorRequired);
        }

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

        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, ip = ?context.ip_address(), "User logged in successfully");
        self.metrics
            .auth_login_total
            .with_label_values(&["success"])
            .inc();
        Ok(())
    }

    pub async fn auth_find_user_by_username(
        &self,
        username: &str,
    ) -> Result<UserModel, AppError> {
        sqlx::query_as::<_, UserModel>(
            "SELECT id, username, display_name, avatar_url, website_url, allow_use, can_search, \
                    last_sign_in_at, created_at, updated_at \
             FROM \"user\" WHERE username = $1",
        )
        .bind(username)
        .fetch_optional(self.db.reader())
        .await
        .map_err(|e| AppError::DatabaseError(e.to_string()))?
        .ok_or(AppError::UserNotFound)
    }

    pub async fn auth_find_user_by_email(
        &self,
        email: &str,
    ) -> Result<UserModel, AppError> {
        sqlx::query_as::<_, UserModel>(
            "SELECT u.id, u.username, u.display_name, u.avatar_url, u.website_url, u.allow_use, \
                    u.can_search, u.last_sign_in_at, u.created_at, u.updated_at \
             FROM \"user\" u \
             INNER JOIN user_email e ON e.\"user\" = u.id \
             WHERE e.email = $1 AND e.active = true",
        )
        .bind(email)
        .fetch_optional(self.db.reader())
        .await
        .map_err(|e| AppError::DatabaseError(e.to_string()))?
        .ok_or(AppError::UserNotFound)
    }

    pub async fn auth_find_user_by_uid(
        &self,
        uid: uuid::Uuid,
    ) -> Result<UserModel, AppError> {
        sqlx::query_as::<_, UserModel>(
            "SELECT id, username, display_name, avatar_url, website_url, allow_use, can_search, \
                    last_sign_in_at, created_at, updated_at \
             FROM \"user\" WHERE id = $1",
        )
        .bind(uid)
        .fetch_optional(self.db.reader())
        .await
        .map_err(|e| AppError::DatabaseError(e.to_string()))?
        .ok_or(AppError::UserNotFound)
    }

    pub fn validate_password_strength(password: &str) -> Result<(), AppError> {
        if password.len() < 8 {
            return Err(AppError::PasswordTooWeak);
        }

        let has_uppercase = password.chars().any(|c| c.is_uppercase());
        let has_lowercase = password.chars().any(|c| c.is_lowercase());
        let has_digit = password.chars().any(|c| c.is_numeric());

        if !has_uppercase || !has_lowercase || !has_digit {
            return Err(AppError::PasswordTooWeak);
        }
        Ok(())
    }
}