use axum::{Json, extract::State}; use serde::{Deserialize, Serialize}; use tracing::error; use tranquil_pds::api::ApiError; use tranquil_pds::api::error::DbResultExt; use tranquil_pds::auth::{Admin, Auth, NotTakendown}; use tranquil_pds::state::AppState; use tranquil_pds::types::Did; use tranquil_pds::util::gen_invite_code; use tranquil_types::InviteCode as InviteCodeValue; #[derive(Deserialize)] #[serde(rename_all = "camelCase")] pub struct CreateInviteCodeInput { pub use_count: i32, pub for_account: Option, } #[derive(Serialize)] pub struct CreateInviteCodeOutput { pub code: InviteCodeValue, } pub async fn create_invite_code( State(state): State, auth: Auth, Json(input): Json, ) -> Result, ApiError> { if input.use_count < 1 { return Err(ApiError::InvalidRequest( "useCount must be at least 1".into(), )); } let for_account: Did = match &input.for_account { Some(acct) => acct .parse() .map_err(|_| ApiError::InvalidDid("Invalid DID format".into()))?, None => auth.did.clone(), }; let code = gen_invite_code(); match state .repos .infra .create_invite_code(&code, input.use_count, &for_account) .await { Ok(true) => Ok(Json(CreateInviteCodeOutput { code })), Ok(false) => { error!("No admin user found to create invite code"); Err(ApiError::InternalError(None)) } Err(e) => { error!("DB error creating invite code: {:?}", e); Err(ApiError::InternalError(None)) } } } #[derive(Deserialize)] #[serde(rename_all = "camelCase")] pub struct CreateInviteCodesInput { pub code_count: Option, pub use_count: i32, pub for_accounts: Option>, } #[derive(Serialize)] pub struct CreateInviteCodesOutput { pub codes: Vec, } #[derive(Serialize)] pub struct AccountCodes { pub account: Did, pub codes: Vec, } pub async fn create_invite_codes( State(state): State, auth: Auth, Json(input): Json, ) -> Result, ApiError> { if input.use_count < 1 { return Err(ApiError::InvalidRequest( "useCount must be at least 1".into(), )); } let code_count = input.code_count.unwrap_or(1).max(1); let for_accounts: Vec = match &input.for_accounts { Some(accounts) if !accounts.is_empty() => accounts .iter() .map(|a| a.parse()) .collect::, _>>() .map_err(|_| ApiError::InvalidDid("Invalid DID format".into()))?, _ => vec![auth.did.clone()], }; let admin_user_id = state .repos .user .get_any_admin_user_id() .await .log_db_err("looking up admin user")? .ok_or_else(|| { error!("No admin user found to create invite codes"); ApiError::InternalError(None) })?; let result = futures::future::try_join_all(for_accounts.into_iter().map(|account| { let infra_repo = state.repos.infra.clone(); let use_count = input.use_count; async move { let codes: Vec = (0..code_count).map(|_| gen_invite_code()).collect(); infra_repo .create_invite_codes_batch(&codes, use_count, admin_user_id, &account) .await .map(|_| AccountCodes { account, codes }) } })) .await; match result { Ok(result_codes) => Ok(Json(CreateInviteCodesOutput { codes: result_codes, })), Err(e) => { error!("DB error creating invite codes: {:?}", e); Err(ApiError::InternalError(None)) } } } #[derive(Deserialize)] #[serde(rename_all = "camelCase")] pub struct GetAccountInviteCodesParams { pub include_used: Option, pub create_available: Option, } #[derive(Serialize)] #[serde(rename_all = "camelCase")] pub struct InviteCode { pub code: InviteCodeValue, pub available: i32, pub disabled: bool, pub for_account: String, pub created_by: String, pub created_at: String, pub uses: Vec, } #[derive(Serialize)] #[serde(rename_all = "camelCase")] pub struct InviteCodeUse { pub used_by: String, #[serde(skip_serializing_if = "Option::is_none")] pub used_by_handle: Option, pub used_at: String, } #[derive(Serialize)] pub struct GetAccountInviteCodesOutput { pub codes: Vec, } pub async fn get_account_invite_codes( State(state): State, auth: Auth, axum::extract::Query(params): axum::extract::Query, ) -> Result, ApiError> { let include_used = params.include_used.unwrap_or(true); let codes_info = state .repos .infra .get_invite_codes_for_account(&auth.did) .await .log_db_err("fetching invite codes")?; let filtered_codes: Vec<_> = codes_info .into_iter() .filter(|info| info.state.is_active()) .collect(); let codes = futures::future::join_all(filtered_codes.into_iter().map(|info| { let infra_repo = state.repos.infra.clone(); async move { let uses: Vec = infra_repo .get_invite_code_uses(&info.code) .await .log_db_err("fetching invite code uses")? .into_iter() .map(|u| InviteCodeUse { used_by: u.used_by_did.to_string(), used_by_handle: u.used_by_handle.map(|h| h.to_string()), used_at: u.used_at.to_rfc3339(), }) .collect(); let use_count = i32::try_from(uses.len()).unwrap_or(i32::MAX); if !include_used && use_count >= info.available_uses { return Ok(None); } Ok(Some(InviteCode { code: info.code, available: info.available_uses, disabled: false, for_account: info.for_account.map(|d| d.to_string()).unwrap_or_default(), created_by: info .created_by .map(|d| d.to_string()) .unwrap_or_else(|| "admin".to_string()), created_at: info.created_at.to_rfc3339(), uses, })) } })) .await; let codes: Vec = codes .into_iter() .collect::>, ApiError>>()? .into_iter() .flatten() .collect(); Ok(Json(GetAccountInviteCodesOutput { codes })) }