audit fixes (2)

multiversx · Sep 20, 2023 · bbff637 · bbff637
1 parent 1e9ec92
commit bbff637
Show file tree

Hide file tree

Showing 8 changed files with 123 additions and 32 deletions.
diff --git a/common/modules/farm/config/src/config.rs b/common/modules/farm/config/src/config.rs
@@ -54,21 +54,17 @@ pub trait ConfigModule: pausable::PausableModule + permissions_module::Permissio
     }
 
     fn is_old_farm_position(&self, token_nonce: Nonce) -> bool {
-        let farm_position_migration_block_nonce = self.farm_position_migration_block_nonce().get();
-        token_nonce > 0 && token_nonce < farm_position_migration_block_nonce
+        let farm_position_migration_nonce = self.farm_position_migration_nonce().get();
+        token_nonce > 0 && token_nonce < farm_position_migration_nonce
     }
 
     #[endpoint(allowExternalClaimBoostedRewards)]
     fn allow_external_claim_boosted_rewards(&self, allow_external_claim: bool) {
         let caller = self.blockchain().get_caller();
-        let user_total_farm_position_mapper = self.user_total_farm_position(&caller);
-        require!(
-            !user_total_farm_position_mapper.is_empty(),
-            "User must have a farm position"
-        );
-        user_total_farm_position_mapper.update(|user_total_farm_position| {
-            user_total_farm_position.allow_external_claim_boosted_rewards = allow_external_claim;
-        });
+        let mut user_total_farm_position = self.get_user_total_farm_position(&caller);
+        user_total_farm_position.allow_external_claim_boosted_rewards = allow_external_claim;
+        self.user_total_farm_position(&caller)
+            .set(user_total_farm_position);
     }
 
     #[view(getFarmingTokenId)]
@@ -101,7 +97,7 @@ pub trait ConfigModule: pausable::PausableModule + permissions_module::Permissio
         user: &ManagedAddress,
     ) -> SingleValueMapper<UserTotalFarmPosition<Self::Api>>;
 
-    #[view(getFarmPositionMigrationBlockNonce)]
-    #[storage_mapper("farm_position_migration_block_nonce")]
-    fn farm_position_migration_block_nonce(&self) -> SingleValueMapper<Nonce>;
+    #[view(getFarmPositionMigrationNonce)]
+    #[storage_mapper("farm_position_migration_nonce")]
+    fn farm_position_migration_nonce(&self) -> SingleValueMapper<Nonce>;
 }
diff --git a/common/modules/farm/farm_base_impl/src/base_traits_impl.rs b/common/modules/farm/farm_base_impl/src/base_traits_impl.rs
@@ -196,12 +196,12 @@ pub trait FarmContract {
     ) {
         let farm_token_mapper = sc.farm_token();
         for farm_position in farm_positions {
+            farm_token_mapper.require_same_token(&farm_position.token_identifier);
+
             if sc.is_old_farm_position(farm_position.token_nonce) {
                 continue;
             }
 
-            farm_token_mapper.require_same_token(&farm_position.token_identifier);
-
             let token_attributes: FarmTokenAttributes<<Self::FarmSc as ContractBase>::Api> =
                 farm_token_mapper.get_token_attributes(farm_position.token_nonce);
 

diff --git a/dex/farm-with-locked-rewards/src/lib.rs b/dex/farm-with-locked-rewards/src/lib.rs
@@ -75,6 +75,9 @@ pub trait Farm:
 
         let current_epoch = self.blockchain().get_block_epoch();
         self.first_week_start_epoch().set_if_empty(current_epoch);
+
+        // Farm position migration code
+        self.try_set_farm_position_migration_nonce();
     }
 
     #[payable("*")]
@@ -155,11 +158,13 @@ pub trait Farm:
 
         let payment = self.call_value().single_esdt();
 
-        self.migrate_old_farm_positions(&orig_caller);
+        let migrated_amount = self.migrate_old_farm_positions(&orig_caller);
 
         let exit_farm_result = self.exit_farm::<NoMintWrapper<Self>>(orig_caller.clone(), payment);
-        let rewards = exit_farm_result.rewards;
 
+        self.decrease_old_farm_positions(migrated_amount, &orig_caller);
+
+        let rewards = exit_farm_result.rewards;
         self.send_payment_non_zero(&caller, &exit_farm_result.farming_tokens);
 
         let locked_rewards_payment = self.send_to_lock_contract_non_zero(

diff --git a/dex/farm/src/base_functions.rs b/dex/farm/src/base_functions.rs
@@ -204,20 +204,63 @@ pub trait BaseFunctionsModule:
         reward
     }
 
-    fn migrate_old_farm_positions(&self, caller: &ManagedAddress) {
+    fn migrate_old_farm_positions(&self, caller: &ManagedAddress) -> BigUint {
         let payments = self.get_non_empty_payments();
         let farm_token_mapper = self.farm_token();
         let farm_token_id = farm_token_mapper.get_token_id();
+        let mut migrated_amount = BigUint::zero();
         for farm_position in &payments {
             if farm_position.token_identifier == farm_token_id
                 && self.is_old_farm_position(farm_position.token_nonce)
             {
-                let mut user_total_farm_position = self.get_user_total_farm_position(caller);
-                user_total_farm_position.total_farm_position += farm_position.amount;
-                self.user_total_farm_position(caller)
-                    .set(user_total_farm_position);
+                migrated_amount += farm_position.amount;
             }
         }
+
+        if migrated_amount > 0 {
+            let mut user_total_farm_position = self.get_user_total_farm_position(caller);
+            user_total_farm_position.total_farm_position += &migrated_amount;
+            self.user_total_farm_position(caller)
+                .set(user_total_farm_position);
+        }
+
+        return migrated_amount;
+    }
+
+    fn decrease_old_farm_positions(&self, migrated_amount: BigUint, caller: &ManagedAddress) {
+        if migrated_amount == BigUint::zero() {
+            return;
+        }
+        self.user_total_farm_position(caller)
+            .update(|user_total_farm_position| {
+                user_total_farm_position.total_farm_position -= migrated_amount;
+            });
+    }
+
+    fn try_set_farm_position_migration_nonce(&self) {
+        if !self.farm_position_migration_nonce().is_empty() {
+            return;
+        }
+
+        let farm_token_mapper = self.farm_token();
+
+        let attributes = FarmTokenAttributes {
+            reward_per_share: BigUint::zero(),
+            entering_epoch: 0,
+            compounded_reward: BigUint::zero(),
+            current_farm_amount: BigUint::zero(),
+            original_owner: self.blockchain().get_sc_address(),
+        };
+
+        let migration_farm_token = farm_token_mapper.nft_create(BigUint::from(1u64), &attributes);
+
+        self.farm_position_migration_nonce()
+            .set(migration_farm_token.token_nonce);
+
+        farm_token_mapper.nft_burn(
+            migration_farm_token.token_nonce,
+            &migration_farm_token.amount,
+        )
     }
 
     fn end_produce_rewards<FC: FarmContract<FarmSc = Self>>(&self) {

diff --git a/dex/farm/src/lib.rs b/dex/farm/src/lib.rs
@@ -78,9 +78,7 @@ pub trait Farm:
         self.first_week_start_epoch().set_if_empty(current_epoch);
 
         // Farm position migration code
-        let block_nonce = self.blockchain().get_block_nonce();
-        self.farm_position_migration_block_nonce()
-            .set_if_empty(block_nonce);
+        self.try_set_farm_position_migration_nonce();
     }
 
     #[payable("*")]
@@ -159,10 +157,12 @@ pub trait Farm:
 
         let payment = self.call_value().single_esdt();
 
-        self.migrate_old_farm_positions(&orig_caller);
+        let migrated_amount = self.migrate_old_farm_positions(&orig_caller);
 
         let exit_farm_result = self.exit_farm::<Wrapper<Self>>(orig_caller.clone(), payment);
 
+        self.decrease_old_farm_positions(migrated_amount, &orig_caller);
+
         self.send_payment_non_zero(&caller, &exit_farm_result.farming_tokens);
         self.send_payment_non_zero(&caller, &exit_farm_result.rewards);
 

diff --git a/farm-staking/farm-staking/src/claim_only_boosted_staking_rewards.rs b/farm-staking/farm-staking/src/claim_only_boosted_staking_rewards.rs
@@ -49,20 +49,37 @@ pub trait ClaimOnlyBoostedStakingRewardsModule:
         boosted_rewards_payment
     }
 
-    fn migrate_old_farm_positions(&self, caller: &ManagedAddress) {
+    fn migrate_old_farm_positions(&self, caller: &ManagedAddress) -> BigUint {
         let payments = self.call_value().all_esdt_transfers().clone_value();
         let farm_token_mapper = self.farm_token();
         let farm_token_id = farm_token_mapper.get_token_id();
+        let mut migrated_amount = BigUint::zero();
         for farm_position in &payments {
             if farm_position.token_identifier == farm_token_id
                 && self.is_old_farm_position(farm_position.token_nonce)
             {
-                let mut user_total_farm_position = self.get_user_total_farm_position(caller);
-                user_total_farm_position.total_farm_position += farm_position.amount;
-                self.user_total_farm_position(caller)
-                    .set(user_total_farm_position);
+                migrated_amount += farm_position.amount;
             }
         }
+
+        if migrated_amount > 0 {
+            let mut user_total_farm_position = self.get_user_total_farm_position(caller);
+            user_total_farm_position.total_farm_position += &migrated_amount;
+            self.user_total_farm_position(caller)
+                .set(user_total_farm_position);
+        }
+
+        return migrated_amount;
+    }
+
+    fn decrease_old_farm_positions(&self, migrated_amount: BigUint, caller: &ManagedAddress) {
+        if migrated_amount == BigUint::zero() {
+            return;
+        }
+        self.user_total_farm_position(caller)
+            .update(|user_total_farm_position| {
+                user_total_farm_position.total_farm_position -= migrated_amount;
+            });
     }
 
     // Cannot import the one from farm, as the Wrapper struct has different dependencies

diff --git a/farm-staking/farm-staking/src/lib.rs b/farm-staking/farm-staking/src/lib.rs
@@ -88,6 +88,9 @@ pub trait FarmStaking:
             "Invalid min unbond epochs"
         );
         self.min_unbond_epochs().set_if_empty(min_unbond_epochs);
+
+        // Farm position migration code
+        self.try_set_farm_position_migration_nonce();
     }
 
     #[payable("*")]
@@ -144,4 +147,29 @@ pub trait FarmStaking:
             "May only call this function through VM query"
         );
     }
+
+    fn try_set_farm_position_migration_nonce(&self) {
+        if !self.farm_position_migration_nonce().is_empty() {
+            return;
+        }
+
+        let farm_token_mapper = self.farm_token();
+
+        let attributes = StakingFarmTokenAttributes {
+            reward_per_share: BigUint::zero(),
+            compounded_reward: BigUint::zero(),
+            current_farm_amount: BigUint::zero(),
+            original_owner: self.blockchain().get_sc_address(),
+        };
+
+        let migration_farm_token = farm_token_mapper.nft_create(BigUint::from(1u64), &attributes);
+
+        self.farm_position_migration_nonce()
+            .set(migration_farm_token.token_nonce);
+
+        farm_token_mapper.nft_burn(
+            migration_farm_token.token_nonce,
+            &migration_farm_token.amount,
+        )
+    }
 }
diff --git a/farm-staking/farm-staking/src/unstake_farm.rs b/farm-staking/farm-staking/src/unstake_farm.rs
@@ -72,11 +72,13 @@ pub trait UnstakeFarmModule:
         payment: EsdtTokenPayment,
         opt_unbond_amount: Option<BigUint>,
     ) -> ExitFarmWithPartialPosResultType<Self::Api> {
-        self.migrate_old_farm_positions(&original_caller);
+        let migrated_amount = self.migrate_old_farm_positions(&original_caller);
 
         let exit_result =
             self.exit_farm_base::<FarmStakingWrapper<Self>>(original_caller.clone(), payment);
 
+        self.decrease_old_farm_positions(migrated_amount, &original_caller);
+
         let unbond_token_amount =
             opt_unbond_amount.unwrap_or(exit_result.farming_token_payment.amount);
         let farm_token_id = exit_result.storage_cache.farm_token_id.clone();