fixes

src/canister/individual_user_template/src/api/cdao/airdrop.rs

@@ -1,46 +1,76 @@
 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};
 
 use crate::CANISTER_DATA;
 
 #[update]
-async fn request_airdrop(token_root: Principal, memo: Option<Memo>, amount: Nat, user_canister: Principal) -> Result<(), AirdropError> {
+async fn request_airdrop(
+    token_root: Principal,
+    memo: Option<Memo>,
+    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, &current_caller)? {// assertion for token owner is checked here will return err if deployed sns cans not found
+    if !is_airdrop_unclaimed(token_root, &current_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<Memo>, amount: Nat) -> Result<(), AirdropError> {
     let ledger = get_ledger(token_root).await?;
     let balance = get_balance(ledger.into()).await?;

src/lib/shared_utils/src/canister_specific/individual_user_template/types/cdao/mod.rs

@@ -12,46 +12,47 @@ 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
     pub principals_who_successfully_claimed: HashMap<Principal, ClaimStatus>,
 }
 
 impl AirdropInfo {
-    pub fn get_claim_status(&self, user_id: &Principal) -> Result<ClaimStatus, String> {
+    pub fn get_claim_status(&self, user_principal_id: &Principal) -> Result<ClaimStatus, String> {
         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<bool, String> {
-        match self.get_claim_status(user_id)? {
+    pub fn is_airdrop_claimed(&self, user_principal_id: &Principal) -> Result<bool, String> {
+        match self.get_claim_status(user_principal_id)? {
             ClaimStatus::Claimed => Ok(true),
             _ => Ok(false),
         }
     }
 
-    pub fn is_airdrop_claiming(&self, user_id: &Principal) -> Result<bool, String> {
-        match self.get_claim_status(user_id)? {
+    pub fn is_airdrop_claiming(&self, user_principal_id: &Principal) -> Result<bool, String> {
+        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]