asklyphe/authservice/src/db/invite_code.rs

/*
 * authservice db/invite_code.rs
 * - pgsql invite code 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 std::time::Instant;
use log::error;
use sea_orm::{ActiveModelTrait, ColumnTrait, DatabaseConnection, EntityTrait, IntoActiveModel, QueryFilter};
use sea_orm::ActiveValue::Set;
use entity::{invite_code, user};
use crate::db::genid;

#[derive(Debug)]
pub enum ConsumeInviteCodeError {
    InviteCodeAlreadyUsed,
    InvalidInviteCode,
    InvalidUser,
    DatabaseError,
}

/// consumes an invite code by relating invite_code_used_by to the given user id
/// and setting the user's invite_code_assigned flag;
/// fails if the invite code is already used, the invite code isn't real, or the given user is invalid
pub async fn consume_invite_code(db: &DatabaseConnection, invite_code: String, user_id: String) -> Result<(), ConsumeInviteCodeError> {

    // is the invite code real?
    let record = invite_code::Entity::find().filter(invite_code::Column::Code.eq(&invite_code)).one(db).await.map_err(|e| {
        error!("DATABASE ERROR WHILE CONSUMEINVITECODE (check if real): {e}");
        ConsumeInviteCodeError::DatabaseError
    })?.ok_or(ConsumeInviteCodeError::InvalidInviteCode)?;
    
    // has it been used?
    
    if record.used == true {
        return Err(ConsumeInviteCodeError::InvalidInviteCode);
    }
    
    // is user real?
    user::Entity::find_by_id(&user_id).one(db).await.map_err(|e| {
        error!("DATABASE ERROR WHILE CONSUMEINVITECODE (check if user good): {e}");
        ConsumeInviteCodeError::DatabaseError
    })?.ok_or(ConsumeInviteCodeError::InvalidUser)?;
    
    // we should be good to go, relate and set flags
    
    let mut record = record.into_active_model();
    record.used = Set(true);
    record.used_by = Set(Some(user_id));
    record.used_on = Set(Some(chrono::Utc::now().naive_utc()));
    
    record.update(db).await.map_err(|e| {
        error!("DATABASE ERROR WHILE CONSUME INVITE CODE (update): {e}");
        ConsumeInviteCodeError::DatabaseError
    })?;

    Ok(())
}

#[derive(Debug)]
pub enum InviteCodeGenerationError {
    CodeAlreadyInserted,
    DatabaseError,
}

pub async fn insert_new_invitecode(db: &DatabaseConnection, invite_code: String, issuer: String, root_issuer: String, comment: String) -> Result<(), InviteCodeGenerationError> {
    let id = genid();
    if invite_code::Entity::find().filter(invite_code::Column::Code.eq(&invite_code)).one(db).await.map_err(|e| {
        error!("DATABASE ERROR WHILE INSERTINVITECODE (check if name free): {e}");
        InviteCodeGenerationError::DatabaseError
    })?.is_some() {
        return Err(InviteCodeGenerationError::CodeAlreadyInserted);
    }

    let record = invite_code::ActiveModel {
        id: Set(id),
        code: Set(invite_code),
        issuer: Set(Some(issuer)),
        root_issuer: Set(root_issuer),
        comment: Set(Some(comment)),
        created_on: Set(chrono::Utc::now().naive_utc()),
        used: Set(false),
        used_on: Default::default(),
        used_by: Default::default(),
        flags: Set(Some(0)),
    };
    
    record.insert(db).await.map_err(|e| {
        error!("DATABASE ERROR WHILE INSERTINVITECODE (insert): {e}");
        InviteCodeGenerationError::DatabaseError
    })?;

    Ok(())
}

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

pub async fn list_invite_codes(db: &DatabaseConnection, show_used: bool) -> Result<Vec<invite_code::Model>, InviteCodeListError> {
    let mut query = invite_code::Entity::find();
    if show_used {
        query = query.filter(invite_code::Column::Used.eq(true));
    } else {
        query = query.filter(invite_code::Column::Used.eq(false));
    }
    query.all(db).await.map_err(|e| {
        error!("DATABASE ERROR WHILE LISTINVITECODES: {e}");
        InviteCodeListError::DatabaseError
    })
}