/* * 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 . */ 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, 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 }) }