asklyphe/authservice/src/db/user.rs

/*
 * authservice db/user.rs
 * - pgsql user functions
 *
 * Copyright (C) 2025 Real Microsoft, LLC
 *
 * This program is free software: you can redistribute it and/or modify it under the terms of the GNU Affero General Public License as published by the Free Software Foundation, version 3.
 *
 * This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License along with this program. If not, see <https://www.gnu.org/licenses/>.
*/

use argon2::{Argon2, PasswordHasher, PasswordVerifier};
use argon2::password_hash::rand_core::OsRng;
use argon2::password_hash::SaltString;
use chrono::{Days, Utc};
use futures::FutureExt;
use log::{debug, error};
use once_cell::sync::Lazy;
use regex::Regex;
use sea_orm::{ActiveModelTrait, ColumnTrait, DatabaseConnection, EntityTrait, IntoActiveModel, ModelTrait, NotSet, PaginatorTrait, QueryFilter, QueryOrder, QuerySelect};
use sea_orm::ActiveValue::Set;
use sea_orm::prelude::Expr;
use sea_orm::sea_query::{BinOper, ExprTrait};
use entity::{user, verification_code};
use crate::db::genid;

static EMAIL_REGEX: Lazy<Regex> = Lazy::new(|| Regex::new(r#"^[a-zA-Z0-9.!#$%&'*+\/=?^_`{|}~-]+@[a-zA-Z0-9](?:[a-zA-Z0-9-]{0,61}[a-zA-Z0-9])?(?:\.[a-zA-Z0-9](?:[a-zA-Z0-9-]{0,61}[a-zA-Z0-9])?)*$"#).expect("BAD EMAIL REGEX"));

#[repr(i64)]
pub enum UserFlag {
    Verified = 1 << 0,
    NewEmail = 1 << 1,
    Administrator = 1 << 2,
}

#[derive(Debug)]
pub enum CreateUserError {
    /// user already exists with this username
    UsernameTaken,
    /// user already exists with this email
    EmailTaken,
    /// username contains invalid (not alphanumeric with dashes and underscores allowed) characters, or is arbitrarily too long
    UsernameError,
    /// display name contains invalid (not alphanumeric with dashes and underscores and spaces (but not at the beginning or end, or twice in a row) characters, or is arbitrarily too long
    DisplayNameError,
    /// email doesn't match funny email regex
    EmailError,
    /// password weak (not implemented yet)
    PasswordError,
    /// arbitrary database error, check logs
    DatabaseError,
}

#[derive(Debug)]
pub enum FetchUserError {
    UserNotFound,
    DatabaseError,
}

/// takes a user id and returns Ok if the user exists
pub async fn check_if_user_exists(db: &DatabaseConnection, user_id: String) -> Result<(), FetchUserError> {
    if user::Entity::find_by_id(&user_id).one(db).await.map_err(|e| {
        error!("DATABASE ERROR WHILE USER EXIST CHECK: {e}");
        FetchUserError::DatabaseError
    })?.is_none() {
        return Err(FetchUserError::UserNotFound);
    }
    Ok(())
}

/// takes a username and returns the associated user id, or an error if the user does not exist
pub async fn username_to_userid(db: &DatabaseConnection, username: String) -> Result<String, FetchUserError> {
    let user_id: Option<String> = user::Entity::find().filter(user::Column::Username.eq(&username)).one(db).await.map_err(|e| {
        error!("DATABASE ERROR WHILE USERNAME2USERID: {e}");
        FetchUserError::DatabaseError
    })?.map(|v| v.id);
    if user_id.is_none() {
        return Err(FetchUserError::UserNotFound);
    }
    Ok(user_id.unwrap())
}

/// takes an email and returns the associated user id, or an error if the user does not exist
pub async fn email_to_userid(db: &DatabaseConnection, email: String) -> Result<String, FetchUserError> {
    let user_id: Option<String> = user::Entity::find().filter(user::Column::Email.eq(&email)).one(db).await.map_err(|e| {
        error!("DATABASE ERROR WHILE EMAIL2USERID: {e}");
        FetchUserError::DatabaseError
    })?.map(|v| v.id);
    if user_id.is_none() {
        return Err(FetchUserError::UserNotFound);
    }
    Ok(user_id.unwrap())
}


/// creates a user in the database, you'll probably want to register an email verification code after this
/// returns id on success
pub async fn create_user(db: &DatabaseConnection, username: String, display_name: String, email: String, password: String) -> Result<String, CreateUserError> {
    // rustrover gives some warnings here but don't change the regex because it's directly
    // taken from the html standard
    if !EMAIL_REGEX.is_match(&email) {
        return Err(CreateUserError::EmailError);
    }
    
    if username.len() > 32 {
        return Err(CreateUserError::UsernameError);
    }

    if user::Entity::find().filter(user::Column::Username.eq(&username)).one(db).await.map_err(|e| {
        error!("DATABASE ERROR WHILE CREATE USER (username taken check (database)): {e}");
        CreateUserError::DatabaseError
    })?.is_some() {
        return Err(CreateUserError::UsernameTaken);
    }
    if user::Entity::find().filter(user::Column::Email.eq(&email)).one(db).await.map_err(|e| {
        error!("DATABASE ERROR WHILE CREATE USER (email taken check (database)): {e}");
        CreateUserError::DatabaseError
    })?.is_some() {
        return Err(CreateUserError::EmailTaken);
    }

    let id = genid();
    let salt = SaltString::generate(&mut OsRng);
    let argon2 = Argon2::default();
    let password_hash = argon2.hash_password(password.as_bytes(), &salt)
        .expect("argon2 error").to_string();
    let now = Utc::now();

    let record = user::ActiveModel {
        id: Set(id.clone()),
        username: Set(username),
        display_name: Set(Some(display_name)),
        email: Set(email),
        new_email: Set(None),
        password: Set(password_hash),
        flags: Set(0),
        created_on: Set(now.naive_utc()),
        last_updated: Set(now.naive_utc()),
    };

    user::Entity::insert(record).exec(db).await.map_err(|e| {
        error!("DATABASE ERROR WHILE CREATE USER (insert user): {e}");
        CreateUserError::DatabaseError
    })?;

    Ok(id)
}

#[derive(Debug)]
pub enum DeleteUserError {
    UserDoesntExist,
    DatabaseError,
}

/// deletes the given user from the user table, doesn't do any other cleanup so be careful!
/// returns an error if the user doesn't exist
pub async fn delete_user(db: &DatabaseConnection, id: String) -> Result<(), DeleteUserError> {
    let record = user::Entity::find_by_id(&id).one(db).await.map_err(|e| {
        error!("DATABASE ERROR WHILE DELETEUSER (exist check (database)): {e}");
        DeleteUserError::DatabaseError
    })?;
    if record.is_none() {
        return Err(DeleteUserError::UserDoesntExist);
    }

    record.unwrap().delete(db).await.map_err(|e| {
        error!("DATABASE ERROR WHILE DELETEUSER (delete (database)): {e}");
        DeleteUserError::DatabaseError
    })?;

    Ok(())
}

#[derive(Debug)]
pub enum RegisterVerificationCodeError {
    /// user is already verified, and thus we do not need a new verification code
    /// note that for email changes, we use the new_verified field which is true until a new email is set
    UserIsVerifiedAndNotChangingTheirEmail,
    /// user seemingly doesn't exist in the database, this should be
    /// shown to the user as an invalid session error, or an internal server error, depending
    /// on the situation
    UserDoesntExist,
    DatabaseError,
}

/// inserts a new email verification code into the user flags for this user,
/// user is verified after you verify_verification_code
pub async fn register_verification_code(db: &DatabaseConnection, id: String, verification_code: String) -> Result<(), RegisterVerificationCodeError> {
    let record = user::Entity::find_by_id(&id).one(db).await.map_err(|e| {
        error!("DATABASE ERROR WHILE REGISTERVERIFICATIONCODE (record for flags check (query)): {e}");
        RegisterVerificationCodeError::DatabaseError
    })?.ok_or(RegisterVerificationCodeError::UserDoesntExist)?;

    // if verified and not asking for new email
    if record.flags & UserFlag::Verified as i64 != 0 && record.flags & UserFlag::NewEmail as i64 == 0 {
        return Err(RegisterVerificationCodeError::UserIsVerifiedAndNotChangingTheirEmail);
    }

    let vc_record = verification_code::ActiveModel {
        id: Set(genid()),
        user: Set(id.clone()),
        code: Set(verification_code),
        set_new_email: Set(false),
        created_on: Set(Utc::now().naive_utc()),
        death_date: Set(Utc::now().checked_add_days(Days::new(10)).unwrap().naive_utc()),
        used: Set(false),
        used_on: NotSet,
    };

    verification_code::Entity::insert(vc_record).exec(db).await.map_err(|e| {
        error!("DATABASE ERROR WHILE REGISTERVERIFICATIONCODE (insert verification code): {e}");
        RegisterVerificationCodeError::DatabaseError
    })?;

    Ok(())
}

#[derive(Debug)]
pub enum VerifyVerificationCodeError {
    UserAlreadyVerified,
    VerificationCodeInvalid,
    UserDoesntExist,
    DatabaseError,
}

/// verifies that the verification code is correct, and sets the user as verified if so
pub async fn verify_verification_code(db: &DatabaseConnection, id: String, verification_code: String) -> Result<(), VerifyVerificationCodeError> {
    let vc_record = verification_code::Entity::find().filter(verification_code::Column::Code.eq(&verification_code)).one(db).await.map_err(|e| {
        error!("DATABASE ERROR WHILE VERIFYVERIFICATIONCODE (code record select): {e}");
        VerifyVerificationCodeError::DatabaseError
    })?.ok_or(VerifyVerificationCodeError::VerificationCodeInvalid)?;

    if vc_record.user != id {
        return Err(VerifyVerificationCodeError::UserDoesntExist);
    }

    let user_record = user::Entity::find_by_id(&vc_record.user).one(db).await.map_err(|e| {
        error!("DATABASE ERROR WHILE VERIFYVERIFICATIONCODE (user record select): {e}");
        VerifyVerificationCodeError::DatabaseError
    })?.ok_or(VerifyVerificationCodeError::UserDoesntExist)?;

    // if verified and not asking for new email
    if user_record.flags & UserFlag::Verified as i64 != 0 && user_record.flags & UserFlag::NewEmail as i64 == 0 {
        return Err(VerifyVerificationCodeError::UserAlreadyVerified);
    }

    let mut user_record = user_record.into_active_model();

    user_record.flags = Set((*user_record.flags.as_ref() | UserFlag::Verified as i64) & !(UserFlag::NewEmail as i64));

    if vc_record.set_new_email {
        if let Some(email) = user_record.new_email.as_ref() {
            user_record.email = Set(email.clone());
            user_record.new_email = Set(None);
        } else {
            error!("vc record is SetNewEmail, but no new email exists for this user {id}!!! INVESTIGATE ASAP");
        }
    }

    let mut vc_record = vc_record.into_active_model();

    vc_record.used = Set(true);
    vc_record.used_on = Set(Some(Utc::now().naive_utc()));

    user_record.update(db).await.map_err(|e| {
        error!("DATABASE ERROR WHILE VERIFYVERIFICATIONCODE (user record update): {e}");
        VerifyVerificationCodeError::DatabaseError
    })?;

    vc_record.update(db).await.map_err(|e| {
        error!("DATABASE ERROR WHILE VERIFYVERIFICATIONCODE (vc record update): {e}");
        VerifyVerificationCodeError::DatabaseError
    })?;

    Ok(())
}

/// checks if the user can log-in with their current registration status
pub async fn check_if_verified(db: &DatabaseConnection, id: String) -> Result<bool, FetchUserError> {
    let user_record = user::Entity::find_by_id(&id).one(db).await.map_err(|e| {
        error!("DATABASE ERROR WHILE CHECKVERIFICATION (user record select): {e}");
        FetchUserError::DatabaseError
    })?.ok_or(FetchUserError::UserNotFound)?;
    
    Ok(user_record.flags & UserFlag::Verified as i64 != 0)
}

#[derive(Debug)]
pub enum VerifyEmailPassComboError {
    ComboInvalid,
    UserDoesntExist,
    DatabaseError,
}

/// verifies that the given email/pass combo are correct for the account associated with the email
/// password must be unhashed
pub async fn verify_email_pass_combo(db: &DatabaseConnection, email: String, password: String) -> Result<(), VerifyEmailPassComboError> {
    let password_hash = user::Entity::find().filter(user::Column::Email.eq(&email)).one(db).await.map_err(|e| {
        error!("DATABASE ERROR WHILE VERIFYEMAILPASSCOMBO (password hash query): {e}");
        VerifyEmailPassComboError::DatabaseError
    })?.ok_or(VerifyEmailPassComboError::UserDoesntExist)?.password;
    let password_hash = argon2::PasswordHash::new(&password_hash).map_err(|e| {
        error!("invalid password hash found for {email}! check database for corruption! argon2 error: {e}");
        VerifyEmailPassComboError::DatabaseError
    })?;
    if Argon2::default().verify_password(password.as_bytes(), &password_hash).is_err() {
        return Err(VerifyEmailPassComboError::ComboInvalid);
    }
    Ok(())
}

#[derive(Debug)]
pub struct UserInfo {
    pub username: String,
    pub email: String,
    pub new_email: Option<String>,
    pub admin: bool,
    pub verified: bool,
}

/// gets all user info (:
pub async fn user_info(db: &DatabaseConnection, id: String) -> Result<UserInfo, FetchUserError> {
    user::Entity::find_by_id(&id).one(db).await.map_err(|e| {
        error!("DATABASE ERROR WHILE USERINFO: {e}");
        FetchUserError::DatabaseError
    })?.ok_or(FetchUserError::UserNotFound).map(|v| {
        UserInfo {
            username: v.username,
            email: v.email,
            new_email: v.new_email,
            admin: v.flags & UserFlag::Administrator as i64 != 0,
            verified: v.flags & UserFlag::Verified as i64 != 0,
        }
    })
}

#[derive(Debug)]
pub enum EmailChangeError {
    UserNotFound,
    NewEmailIsInvalid,
    DatabaseError,
}

/// starts the process for changing a user's email
pub async fn start_email_change(db: &DatabaseConnection, id: String, new_email: String, verification_code: String) -> Result<(), EmailChangeError> {
    if !EMAIL_REGEX.is_match(&new_email) {
        return Err(EmailChangeError::NewEmailIsInvalid);
    }

    let user_record = user::Entity::find_by_id(&id).one(db).await.map_err(|e| {
        error!("DATABASE ERROR WHILE STARTEMAILCHANGE (user query): {e}");
        EmailChangeError::DatabaseError
    })?.ok_or(EmailChangeError::UserNotFound)?;

    let flags = user_record.flags;
    let mut user_record = user_record.into_active_model();
    user_record.new_email = Set(Some(new_email));
    user_record.flags = Set(flags | UserFlag::NewEmail as i64);
    user_record.update(db).await.map_err(|e| {
        error!("DATABASE ERROR WHILE STARTEMAILCHANGE (user update): {e}");
        EmailChangeError::DatabaseError
    })?;

    let old_vc_records = verification_code::Entity::delete_many()
        .filter(verification_code::Column::User.eq(&id))
        .filter(verification_code::Column::SetNewEmail.eq(true))
        .exec(db).await.map_err(|e| {
        error!("DATABASE ERROR WHILE STARTEMAILCHANGE (delete old email codes): {e}");
        EmailChangeError::DatabaseError
    })?;
    debug!("deleted {} old vc code records during email change", old_vc_records.rows_affected);

    let vc_record = verification_code::ActiveModel {
        id: Set(genid()),
        user: Set(id.clone()),
        code: Set(verification_code),
        set_new_email: Set(true),
        created_on: Set(Utc::now().naive_utc()),
        death_date: Set(Utc::now().checked_add_days(Days::new(10)).unwrap().naive_utc()),
        used: Set(false),
        used_on: NotSet,
    };

    verification_code::Entity::insert(vc_record).exec(db).await.map_err(|e| {
        error!("DATABASE ERROR WHILE STARTEMAILCHANGE (insert verification code): {e}");
        EmailChangeError::DatabaseError
    })?;

    Ok(())
}

/// cancels any in-progress user email change
pub async fn cancel_email_change(db: &DatabaseConnection, id: String) -> Result<(), FetchUserError> {
    verification_code::Entity::delete_many()
        .filter(verification_code::Column::User.eq(&id))
        .filter(verification_code::Column::SetNewEmail.eq(true))
        .exec(db).await.map_err(|e| {
        error!("DATABASE ERROR WHILE CANCELEMAILCHANGE (delete old email codes): {e}");
        FetchUserError::DatabaseError
    })?;

    let mut user_record = user::Entity::find_by_id(&id).one(db).await.map_err(|e| {
        error!("DATABASE ERROR WHILE CANCELEMAILCHANGE (query user): {e}");
        FetchUserError::DatabaseError
    })?.ok_or(FetchUserError::UserNotFound)?.into_active_model();

    let flags = *user_record.flags.as_ref();

    user_record.new_email = Set(None);
    // unset new email
    user_record.flags = Set(flags & !(UserFlag::NewEmail as i64));

    user_record.update(db).await.map_err(|e| {
        error!("DATABASE ERROR WHILE CANCELEMAILCHANGE (user update): {e}");
        FetchUserError::DatabaseError
    })?;

    Ok(())
}

/// returns the last 10 users added to the database as usernames
pub async fn last_ten_usernames(db: &DatabaseConnection) -> Result<Vec<String>, FetchUserError> {
    Ok(user::Entity::find()
        .order_by_desc(user::Column::CreatedOn)
        .limit(Some(10))
        .all(db).await.map_err(|e| {
        error!("DATABASE ERROR WHILE LASTTENUSERNAMES: {e}");
        FetchUserError::DatabaseError
    })?.into_iter().map(|v| v.username).collect())
}

pub async fn all_users(db: &DatabaseConnection) -> Result<Vec<UserInfo>, FetchUserError> {
    Ok(user::Entity::find().all(db).await.map_err(|e| {
        error!("DATABASE ERROR WHILE ALLUSERS: {e}");
        FetchUserError::DatabaseError
    })?.into_iter().map(|v| {
        UserInfo {
            username: v.username,
            email: v.email,
            new_email: v.new_email,
            admin: v.flags & UserFlag::Administrator as i64 != 0,
            verified: v.flags & UserFlag::Verified as i64 != 0,
        }
    }).collect())
}

/// returns the number of users in the database
pub async fn user_count(db: &DatabaseConnection) -> Result<usize, FetchUserError> {
    Ok(user::Entity::find().count(db).await.map_err(|e| {
        error!("DATABASE ERROR WHILE USERCOUNT: {e}");
        FetchUserError::DatabaseError
    })? as usize)
}

/// returns the number of users in the database who are admins
pub async fn admin_count(db: &DatabaseConnection) -> Result<usize, FetchUserError> {
    // dont fucking touch this, i don't know why it works but it does, it's actually evil
    // note: doesn't work
    Ok(user::Entity::find().filter(user::Column::Flags.into_expr().binary(BinOper::LShift, Expr::value(63 - 2)).lt(1 << (63 - 2)))
        .count(db).await.map_err(|e| {
        error!("DATABASE ERROR WHILE ADMINCOUNT: {e}");
        FetchUserError::DatabaseError
    })? as usize)
}

/// returns the number of unverified accounts (this does NOT count email change requests)
pub async fn unverified_count(db: &DatabaseConnection) -> Result<usize, FetchUserError> {
    Ok(verification_code::Entity::find().filter(verification_code::Column::SetNewEmail.eq(false)).filter(verification_code::Column::Used.eq(false))
        .count(db).await.map_err(|e| {
        error!("DATABASE ERROR WHILE UNVERIFIEDCOUNT: {e}");
        FetchUserError::DatabaseError
    })? as usize)
}

/// returns a list of all (username, email) to send announcement emails to
pub async fn email_list(db: &DatabaseConnection) -> Result<Vec<(String, String)>, FetchUserError> {
    // fixme: filter by who wants emails and who doesn't when we implement that
    Ok(user::Entity::find().all(db).await.map_err(|e| {
        error!("DATABASE ERROR WHILE EMAILLIST: {e}");
        FetchUserError::DatabaseError
    })?.into_iter().map(|v| {
        (v.username, v.email)
    }).collect())
}