Merge pull request #827 from multiversx/fix/metastaking-withdraw-fix-…

…after-audit

Fix/metastaking withdraw fix after audit

farm-staking/farm-staking/src/custom_rewards.rs

@@ -10,6 +10,7 @@ use crate::base_impl_wrapper::FarmStakingWrapper;
 pub const MAX_PERCENT: u64 = 10_000;
 pub const BLOCKS_IN_YEAR: u64 = 31_536_000 / 6; // seconds_in_year / 6_seconds_per_block
 pub const MAX_MIN_UNBOND_EPOCHS: u64 = 30;
+pub const WITHDRAW_AMOUNT_TOO_HIGH: &str = "Withdraw amount is higher than the remaining uncollected rewards!";
 
 #[multiversx_sc::module]
 pub trait CustomRewardsModule:
@@ -48,7 +49,15 @@ pub trait CustomRewardsModule:
     fn withdraw_rewards(&self, withdraw_amount: BigUint) {
         self.require_caller_has_admin_permissions();
 
-        self.reward_capacity().update(|rewards| {
+        let mut storage_cache = StorageCache::new(self);
+        FarmStakingWrapper::<Self>::generate_aggregated_rewards(self, &mut storage_cache);
+
+        let reward_capacity_mapper = self.reward_capacity();
+        let accumulated_rewards_mapper = self.accumulated_rewards();
+        let remaining_rewards = reward_capacity_mapper.get() - accumulated_rewards_mapper.get();
+        require!(withdraw_amount <= remaining_rewards, WITHDRAW_AMOUNT_TOO_HIGH);
+
+        reward_capacity_mapper.update(|rewards| {
             require!(
                 *rewards >= withdraw_amount,
                 "Not enough rewards to withdraw"

farm-staking/farm-staking/tests/farm_staking_setup/mod.rs

@@ -43,6 +43,8 @@ pub const USER_REWARDS_ENERGY_CONST: u64 = 3;
 pub const USER_REWARDS_FARM_CONST: u64 = 2;
 pub const MIN_ENERGY_AMOUNT_FOR_BOOSTED_YIELDS: u64 = 1;
 pub const MIN_FARM_AMOUNT_FOR_BOOSTED_YIELDS: u64 = 1;
+pub const WITHDRAW_AMOUNT_TOO_HIGH: &str = "Withdraw amount is higher than the remaining uncollected rewards!";
+
 
 pub struct FarmStakingSetup<FarmObjBuilder, EnergyFactoryBuilder>
 where
@@ -515,4 +517,17 @@ where
             )
             .assert_ok();
     }
+
+    pub fn withdraw_rewards_with_error(&mut self, withdraw_amount: &RustBigUint, expected_status: u64, expected_message: &str) {
+        self.b_mock
+            .execute_tx(
+                &self.owner_address,
+                &self.farm_wrapper,
+                &rust_biguint!(0),
+                |sc| {
+                    sc.withdraw_rewards(withdraw_amount.into());
+                },
+            )
+            .assert_error(expected_status, expected_message)
+    }
 }

farm-staking/farm-staking/tests/farm_staking_test.rs

@@ -1,8 +1,6 @@
 #![allow(deprecated)]
 
-use multiversx_sc_scenario::{
-    rust_biguint, whitebox_legacy::TxTokenTransfer, DebugApi,
-};
+use multiversx_sc_scenario::{rust_biguint, whitebox_legacy::TxTokenTransfer, DebugApi};
 
 pub mod farm_staking_setup;
 use farm_staking::{
@@ -249,3 +247,33 @@ fn test_withdraw_rewards() {
     let final_rewards_capacity = 0u64;
     farm_setup.check_rewards_capacity(final_rewards_capacity);
 }
+
+#[test]
+fn test_withdraw_after_produced_rewards() {
+    DebugApi::dummy();
+    let mut farm_setup =
+        FarmStakingSetup::new(farm_staking::contract_obj, energy_factory::contract_obj);
+
+    let initial_rewards_capacity = 1_000_000_000_000u64;
+    farm_setup.check_rewards_capacity(initial_rewards_capacity);
+
+    let farm_in_amount = 100_000_000;
+    let expected_farm_token_nonce = 1;
+    farm_setup.stake_farm(farm_in_amount, &[], expected_farm_token_nonce, 0, 0);
+    farm_setup.check_farm_token_supply(farm_in_amount);
+
+    farm_setup.set_block_epoch(5);
+    farm_setup.set_block_nonce(10);
+
+    let withdraw_amount = rust_biguint!(TOTAL_REWARDS_AMOUNT);
+    farm_setup.withdraw_rewards_with_error(&withdraw_amount, 4, WITHDRAW_AMOUNT_TOO_HIGH);
+
+    let expected_reward_token_out = 40;
+
+    let withdraw_amount = rust_biguint!(TOTAL_REWARDS_AMOUNT) - rust_biguint!(expected_reward_token_out);
+    farm_setup.withdraw_rewards(&withdraw_amount);
+
+    // Only the user's rewards will remain
+    let final_rewards_capacity = expected_reward_token_out;
+    farm_setup.check_rewards_capacity(final_rewards_capacity);
+}

farm-staking/farm-staking/wasm/src/lib.rs

@@ -5,9 +5,9 @@
 ////////////////////////////////////////////////////
 
 // Init:                                 1
-// Endpoints:                           70
+// Endpoints:                           71
 // Async Callback:                       1
-// Total number of exported functions:  72
+// Total number of exported functions:  73
 
 #![no_std]
 #![allow(internal_features)]
@@ -24,6 +24,7 @@ multiversx_sc_wasm_adapter::endpoints! {
         mergeFarmTokens => merge_farm_tokens_endpoint
         calculateRewardsForGivenPosition => calculate_rewards_for_given_position
         topUpRewards => top_up_rewards
+        withdrawRewards => withdraw_rewards
         endProduceRewards => end_produce_rewards
         setPerBlockRewardAmount => set_per_block_rewards
         setMaxApr => set_max_apr