appks/service/auth/register.rs

use argon2::password_hash::SaltString;
use argon2::{Argon2, password_hash::PasswordHasher};
use rand::Rng;
use serde::{Deserialize, Serialize};
use std::time::Duration;

use crate::error::AppError;
use crate::models::users::User;
use crate::pb::email::{EmailAddress, SendEmailRequest};
use crate::service::AuthService;
use crate::session::Session;

#[derive(Deserialize, Serialize, Clone, Debug, utoipa::ToSchema)]
pub struct RegisterParams {
    pub username: String,
    pub email: String,
    pub password: String,
    pub captcha: String,
    pub email_code: String,
}

#[derive(Deserialize, Serialize, Clone, Debug, utoipa::ToSchema)]
pub struct RegisterEmailCodeParams {
    pub email: String,
    pub captcha: String,
}

#[derive(Serialize, Clone, Debug, utoipa::ToSchema)]
pub struct RegisterEmailCodeResponse {
    pub expires_in_secs: u64,
}

impl AuthService {
    const REGISTER_EMAIL_CODE_PREFIX: &'static str = "auth:register_email:";
    const REGISTER_EMAIL_CODE_TTL_SECS: u64 = 600;
    const REGISTER_EMAIL_CODE_COOLDOWN_SECS: u64 = 60;

    pub async fn auth_register_email_code(
        &self,
        params: RegisterEmailCodeParams,
        context: &Session,
    ) -> Result<RegisterEmailCodeResponse, AppError> {
        self.auth_check_captcha(context, params.captcha).await?;
        let email = params.email.trim().to_lowercase();
        if email.is_empty() {
            return Err(AppError::BadRequest("email is required".into()));
        }

        if self.auth_verified_email_exists(&email).await? {
            return Err(AppError::EmailExists);
        }

        let cooldown_key = format!("{}cooldown:{}", Self::REGISTER_EMAIL_CODE_PREFIX, email);
        if self.ctx.cache.exists(&cooldown_key).await {
            return Err(AppError::BadRequest(
                "verification code was sent recently; please try again later".into(),
            ));
        }

        let code = Self::generate_register_email_code();
        let cache_key = Self::register_email_code_key(&email);
        self.ctx
            .cache
            .set(
                &cache_key,
                &code,
                Some(Duration::from_secs(Self::REGISTER_EMAIL_CODE_TTL_SECS)),
            )
            .await
            .map_err(|e| AppError::InternalServerError(e.to_string()))?;
        self.ctx
            .cache
            .set(
                &cooldown_key,
                &true,
                Some(Duration::from_secs(Self::REGISTER_EMAIL_CODE_COOLDOWN_SECS)),
            )
            .await
            .map_err(|e| AppError::InternalServerError(e.to_string()))?;

        let mut mail = self
            .ctx
            .registry
            .get_email_client()
            .ok_or(AppError::Config("mail service not available".into()))?;
        mail.send_email(tonic::Request::new(SendEmailRequest {
            to: vec![EmailAddress {
                email: email.clone(),
                name: String::new(),
            }],
            subject: "Register Email Verification".into(),
            text_body: format!(
                "Your registration verification code is: {}\n\nThis code expires in 10 minutes.",
                code
            ),
            ..Default::default()
        }))
        .await
        .map_err(|e| AppError::InternalServerError(e.to_string()))?;

        Ok(RegisterEmailCodeResponse {
            expires_in_secs: Self::REGISTER_EMAIL_CODE_TTL_SECS,
        })
    }

    async fn auth_check_register_email_code(&self, email: &str, code: &str) -> Result<(), AppError> {
        let cache_key = Self::register_email_code_key(email);
        let stored = self
            .ctx
            .cache
            .get::<String>(&cache_key)
            .await
            .ok_or(AppError::InvalidEmailCode)?;
        if !crate::service::util::constant_time_eq(stored.trim(), code.trim()) {
            return Err(AppError::InvalidEmailCode);
        }
        let _ = self.ctx.cache.delete(&cache_key).await;
        Ok(())
    }

    fn register_email_code_key(email: &str) -> String {
        format!(
            "{}{}",
            Self::REGISTER_EMAIL_CODE_PREFIX,
            email.trim().to_lowercase()
        )
    }

    fn generate_register_email_code() -> String {
        let mut rng = rand::thread_rng();
        format!("{:06}", rng.gen_range(0..1_000_000))
    }

    async fn auth_username_exists(&self, username: &str) -> Result<bool, AppError> {
        sqlx::query_scalar::<_, bool>(
            "SELECT EXISTS(SELECT 1 FROM \"user\" WHERE lower(username) = lower($1) AND deleted_at IS NULL)",
        )
        .bind(username)
        .fetch_one(self.ctx.db.reader())
        .await
        .map_err(AppError::Database)
    }

    async fn auth_verified_email_exists(&self, email: &str) -> Result<bool, AppError> {
        sqlx::query_scalar::<_, bool>(
            "SELECT EXISTS(SELECT 1 FROM user_mail WHERE lower(email) = lower($1) AND is_verified = true)",
        )
        .bind(email)
        .fetch_one(self.ctx.db.reader())
        .await
        .map_err(AppError::Database)
    }

    #[tracing::instrument(skip(self, params, context), fields(username = %params.username))]
    pub async fn auth_register(
        &self,
        params: RegisterParams,
        context: &Session,
    ) -> Result<User, AppError> {
        self.auth_check_captcha(context, params.captcha).await?;
        let username = params.username.trim().to_string();
        let email = params.email.trim().to_lowercase();
        if username.is_empty() || email.is_empty() {
            return Err(AppError::BadRequest(
                "username and email are required".into(),
            ));
        }
        let password = self.auth_rsa_decode(context, params.password).await?;
        crate::service::util::validate_password_strength(&password)?;

        let username_exists = self.auth_username_exists(&username).await?;
        let email_exists = self.auth_verified_email_exists(&email).await?;
        if username_exists || email_exists {
            return Err(AppError::AccountAlreadyExists);
        }

        self.auth_check_register_email_code(&email, &params.email_code).await?;

        let user_id = uuid::Uuid::now_v7();
        let now = chrono::Utc::now();
        let salt = SaltString::generate(&mut rand::thread_rng());
        let password_hash = Argon2::default()
            .hash_password(password.as_bytes(), &salt)
            .map_err(|e| AppError::PasswordHashError(e.to_string()))?
            .to_string();

        let mut txn = self
            .ctx
            .db
            .writer()
            .begin()
            .await
            .map_err(|_| AppError::TxnError)?;

        let user = sqlx::query_as::<_, User>(
            "INSERT INTO \"user\" \
             (id, username, display_name, status, role, visibility, is_active, is_bot, \
              last_login_at, created_at, updated_at) \
             VALUES ($1, $2, $2, 'active', 'user', 'public', true, false, NULL, $3, $3) \
             RETURNING id, username, display_name, avatar_url, bio, status, role, visibility, \
                       is_active, is_bot, last_login_at, created_at, updated_at, deleted_at",
        )
        .bind(user_id)
        .bind(&username)
        .bind(now)
        .fetch_one(&mut *txn)
        .await
        .map_err(AppError::Database)?;

        sqlx::query(
            "INSERT INTO user_mail (id, user_id, email, is_primary, is_verified, created_at, updated_at) \
             VALUES ($1, $2, $3, true, true, $4, $4)",
        )
        .bind(uuid::Uuid::now_v7())
        .bind(user_id)
        .bind(&email)
        .bind(now)
        .execute(&mut *txn)
        .await
        .map_err(AppError::Database)?;

        sqlx::query(
            "INSERT INTO user_password (user_id, password_hash, password_algo, password_salt, \
             must_change_password, password_updated_at, created_at, updated_at) \
             VALUES ($1, $2, 'argon2id', '', false, $3, $3, $3)",
        )
        .bind(user_id)
        .bind(&password_hash)
        .bind(now)
        .execute(&mut *txn)
        .await
        .map_err(AppError::Database)?;

        txn.commit().await.map_err(|_| AppError::TxnError)?;

        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 registered");
        Ok(user)
    }
}