diff --git a/src/canister/individual_user_template/src/api/cdao/airdrop.rs b/src/canister/individual_user_template/src/api/cdao/airdrop.rs index 62aca103..81da56c5 100644 --- a/src/canister/individual_user_template/src/api/cdao/airdrop.rs +++ b/src/canister/individual_user_template/src/api/cdao/airdrop.rs @@ -1,6 +1,6 @@ use candid::{Nat, Principal}; use ic_base_types::PrincipalId; -use ic_cdk::update; +use ic_cdk_macros::update; use ic_sns_root::pb::v1::{ListSnsCanistersRequest, ListSnsCanistersResponse}; use icrc_ledger_types::icrc1::{account::Account, transfer::{Memo, TransferArg, TransferError}}; use shared_utils::canister_specific::individual_user_template::types::{error::AirdropError, profile::UserProfileDetailsForFrontendV2}; @@ -8,39 +8,69 @@ use shared_utils::canister_specific::individual_user_template::types::{error::Ai use crate::CANISTER_DATA; #[update] -async fn request_airdrop(token_root: Principal, memo: Option, amount: Nat, user_canister: Principal) -> Result<(), AirdropError> { +async fn request_airdrop( + token_root: Principal, + memo: Option, + amount: Nat, + user_canister: Principal, +) -> Result<(), AirdropError> { let current_caller = ic_cdk::caller(); let profile_info = get_profile_info(user_canister).await?; - + if profile_info.principal_id != current_caller { return Err(AirdropError::CanisterPrincipalDoNotMatch); } - if !is_airdrop_unclaimed(token_root, ¤t_caller)? {// assertion for token owner is checked here will return err if deployed sns cans not found + if !is_airdrop_unclaimed(token_root, ¤t_caller)? { return Err(AirdropError::AlreadyClaimedAirdrop); } let amount = amount.min(1000u32.into()); if amount < 100u32 { - return Err(AirdropError::RequestedAmountTooLow) + return Err(AirdropError::RequestedAmountTooLow); } - set_airdrop_claiming(token_root, current_caller); + let deployed_sns_index = CANISTER_DATA.with_borrow(|cans_data| { + cans_data + .cdao_canisters + .iter() + .position(|cdao| cdao.root == token_root) + }); - request_airdrop_internal(token_root, current_caller, memo, amount).await.inspect(|_|{ - CANISTER_DATA.with_borrow_mut(|cans_data| { - cans_data - .cdao_canisters - .iter_mut() - .find(|cdao| cdao.root == token_root) - .map(|cdao| cdao.airdrop_info.set_airdrop_unclaimed(current_caller)).unwrap(); // can safely unwrap updating the states for the airdrop for the user creates it in place if not exists - }); - })?; // rollback to unclaimed if error + let deployed_sns_index = match deployed_sns_index { + Some(index) => index, + None => return Err(AirdropError::InvalidRoot), + }; - set_airdrop_claimed(token_root, current_caller); - Ok(()) + CANISTER_DATA.with_borrow_mut(|cans_data| { + cans_data.cdao_canisters[deployed_sns_index] + .airdrop_info + .set_airdrop_claiming(current_caller); + }); + + let result = request_airdrop_internal(token_root, current_caller, memo, amount).await; + + match result { + Ok(_) => { + CANISTER_DATA.with_borrow_mut(|cans_data| { + cans_data.cdao_canisters[deployed_sns_index] + .airdrop_info + .set_airdrop_claimed(current_caller); + }); + Ok(()) + } + Err(e) => { + CANISTER_DATA.with_borrow_mut(|cans_data| { + cans_data.cdao_canisters[deployed_sns_index] + .airdrop_info + .set_airdrop_unclaimed(current_caller); + }); + Err(e) + } + } } + async fn request_airdrop_internal(token_root: Principal, current_caller: Principal, memo: Option, amount: Nat) -> Result<(), AirdropError> { let ledger = get_ledger(token_root).await?; let balance = get_balance(ledger.into()).await?; diff --git a/src/lib/shared_utils/src/canister_specific/individual_user_template/types/cdao/mod.rs b/src/lib/shared_utils/src/canister_specific/individual_user_template/types/cdao/mod.rs index 710bc2b2..433b0360 100644 --- a/src/lib/shared_utils/src/canister_specific/individual_user_template/types/cdao/mod.rs +++ b/src/lib/shared_utils/src/canister_specific/individual_user_template/types/cdao/mod.rs @@ -12,6 +12,7 @@ pub struct DeployedCdaoCanisters { pub index: Principal, pub airdrop_info: AirdropInfo, } + #[derive(CandidType, PartialEq, Eq, Debug, Serialize, Deserialize, Clone)] pub struct AirdropInfo { /// Maps each principal to their claim status @@ -19,39 +20,39 @@ pub struct AirdropInfo { } impl AirdropInfo { - pub fn get_claim_status(&self, user_id: &Principal) -> Result { + pub fn get_claim_status(&self, user_principal_id: &Principal) -> Result { self.principals_who_successfully_claimed - .get(user_id) + .get(user_principal_id) .cloned() - .ok_or_else(|| format!("Principal {} not found", user_id)) + .ok_or_else(|| format!("Principal {} not found", user_principal_id)) } - pub fn is_airdrop_claimed(&self, user_id: &Principal) -> Result { - match self.get_claim_status(user_id)? { + pub fn is_airdrop_claimed(&self, user_principal_id: &Principal) -> Result { + match self.get_claim_status(user_principal_id)? { ClaimStatus::Claimed => Ok(true), _ => Ok(false), } } - pub fn is_airdrop_claiming(&self, user_id: &Principal) -> Result { - match self.get_claim_status(user_id)? { + pub fn is_airdrop_claiming(&self, user_principal_id: &Principal) -> Result { + match self.get_claim_status(user_principal_id)? { ClaimStatus::Claiming => Ok(true), _ => Ok(false), } } - pub fn is_airdrop_unclaimed(&self, user_id: &Principal) -> bool{ - matches!(self.get_claim_status(user_id), Ok(ClaimStatus::Unclaimed) | Err(_)) + pub fn is_airdrop_unclaimed(&self, user_principal_id: &Principal) -> bool{ + matches!(self.get_claim_status(user_principal_id), Ok(ClaimStatus::Unclaimed) | Err(_)) } fn set_claim_status_or_insert_with_claim_status_if_not_exist( &mut self, - user_id: &Principal, + user_principal_id: &Principal, status: ClaimStatus, ) { use std::collections::hash_map::Entry; - match self.principals_who_successfully_claimed.entry(*user_id) { + match self.principals_who_successfully_claimed.entry(*user_principal_id) { Entry::Occupied(mut entry) => { *entry.get_mut() = status; } @@ -61,25 +62,19 @@ impl AirdropInfo { } } - pub fn set_airdrop_claimed(&mut self, user_id: Principal) { - self.set_claim_status_or_insert_with_claim_status_if_not_exist(&user_id, ClaimStatus::Claimed) + pub fn set_airdrop_claimed(&mut self, user_principal_id: Principal) { + self.set_claim_status_or_insert_with_claim_status_if_not_exist(&user_principal_id, ClaimStatus::Claimed) } - pub fn set_airdrop_claiming(&mut self, user_id: Principal){ - self.set_claim_status_or_insert_with_claim_status_if_not_exist(&user_id, ClaimStatus::Claiming) + pub fn set_airdrop_claiming(&mut self, user_principal_id: Principal){ + self.set_claim_status_or_insert_with_claim_status_if_not_exist(&user_principal_id, ClaimStatus::Claiming) } - pub fn set_airdrop_unclaimed(&mut self, user_id: Principal) { - self.set_claim_status_or_insert_with_claim_status_if_not_exist(&user_id, ClaimStatus::Unclaimed) + pub fn set_airdrop_unclaimed(&mut self, user_principal_id: Principal) { + self.set_claim_status_or_insert_with_claim_status_if_not_exist(&user_principal_id, ClaimStatus::Unclaimed) } } -#[derive(CandidType, PartialEq, Eq, Debug, Serialize, Deserialize, Clone, Hash)] -pub struct PrincipalEligibleToClaimAirdrop { - pub user_id: Principal, - pub claim_status: ClaimStatus, -} - #[derive(Serialize, Deserialize, CandidType, Clone, Debug, PartialEq, Eq, Default, Hash)] pub enum ClaimStatus { #[default]