fix withdraw underflow

crates/sui-framework/packages/sui-system/sources/staking_pool.move

@@ -322,12 +322,16 @@ module sui_system::staking_pool {
         // Check that the stake information matches the pool.
         assert!(staked_sui.pool_id == object::id(pool), EWrongPool);
 
+        // TODO: after fixing the inactive stake issue, change this to always
+        // look forward and find the next available rate. It's guaranteed
+        // the next available rate can be found, otherwise we will never reach
+        // this function, instead we will exit earlier in the inactive stake check.
         let exchange_rate_at_staking_epoch = pool_token_exchange_rate_at_epoch(pool, staked_sui.stake_activation_epoch);
         let principal_withdraw = unwrap_staked_sui(staked_sui);
         let pool_token_withdraw_amount = get_token_amount(
-		&exchange_rate_at_staking_epoch,
-		principal_withdraw.value()
-	);
+            &exchange_rate_at_staking_epoch,
+            principal_withdraw.value()
+        );
 
         (
             pool_token_withdraw_amount,
@@ -371,8 +375,15 @@ module sui_system::staking_pool {
     /// Called at epoch boundaries to process pending stake withdraws requested during the epoch.
     /// Also called immediately upon withdrawal if the pool is inactive.
     fun process_pending_stake_withdraw(pool: &mut StakingPool) {
-        pool.sui_balance = pool.sui_balance - pool.pending_total_sui_withdraw;
-        pool.pool_token_balance = pool.pool_token_balance - pool.pending_pool_token_withdraw;
+        pool.sui_balance =
+            if (pool.sui_balance >= pool.pending_total_sui_withdraw)
+                pool.sui_balance - pool.pending_total_sui_withdraw
+            else 0;
+
+        pool.pool_token_balance =
+            if (pool.pool_token_balance >= pool.pending_pool_token_withdraw)
+                pool.pool_token_balance - pool.pending_pool_token_withdraw
+            else 0;
         pool.pending_total_sui_withdraw = 0;
         pool.pending_pool_token_withdraw = 0;
     }
@@ -676,6 +687,42 @@ module sui_system::staking_pool {
     #[test_only]
     public use fun fungible_staked_sui_data_principal_value as FungibleStakedSuiData.principal_value;
 
+    // NEVER remove `#[test_only]`
+    #[test_only]
+    public(package) fun process_pending_stake_withdraw_test_only(pool: &mut StakingPool) {
+        pool.process_pending_stake_withdraw()
+    }
+
+    // NEVER remove `#[test_only]`
+    #[test_only]
+    public(package) fun increase_pending_pool_token_withdraw_test_only(pool: &mut StakingPool, delta: u64) {
+        pool.pending_pool_token_withdraw = pool.pending_pool_token_withdraw + delta
+    }
+
+    // NEVER remove `#[test_only]`
+    #[test_only]
+    public(package) fun increase_pending_total_sui_withdraw_test_only(pool: &mut StakingPool, delta: u64) {
+        pool.pending_total_sui_withdraw = pool.pending_total_sui_withdraw + delta
+    }
+
+    // NEVER remove `#[test_only]`
+    #[test_only]
+    public(package) fun pending_total_sui_withdraw(pool: &StakingPool): u64 {
+        pool.pending_total_sui_withdraw
+    }
+
+    // NEVER remove `#[test_only]`
+    #[test_only]
+    public(package) fun pool_token_balance(pool: &StakingPool): u64 {
+        pool.pool_token_balance
+    }
+
+    // NEVER remove `#[test_only]`
+    #[test_only]
+    public(package) fun pending_pool_token_withdraw(pool: &StakingPool): u64 {
+        pool.pending_pool_token_withdraw
+    }
+
     #[test_only]
     public(package) fun fungible_staked_sui_data_principal_value(fungible_staked_sui_data: &FungibleStakedSuiData): u64 {
         fungible_staked_sui_data.principal.value()

crates/sui-framework/packages/sui-system/sources/sui_system.move

@@ -698,6 +698,13 @@ module sui_system::sui_system {
         self.set_epoch_for_testing(epoch_num)
     }
 
+    // NEVER remove `#[test_only]`
+    #[test_only]
+    public(package) fun increment_epoch_for_testing(wrapper: &mut SuiSystemState) {
+        let self = load_system_state_mut(wrapper);
+        self.increment_epoch_for_testing()
+    }
+
     #[test_only]
     public fun request_add_validator_for_testing(
         wrapper: &mut SuiSystemState,

crates/sui-framework/packages/sui-system/sources/sui_system_state_inner.move

@@ -1121,11 +1121,18 @@ module sui_system::sui_system_state_inner {
         self.stake_subsidy.get_distribution_counter()
     }
 
+    // NEVER remove `#[test_only]`
     #[test_only]
     public(package) fun set_epoch_for_testing(self: &mut SuiSystemStateInnerV2, epoch_num: u64) {
         self.epoch = epoch_num
     }
 
+    // NEVER remove `#[test_only]`
+    #[test_only]
+    public(package) fun increment_epoch_for_testing(self: &mut SuiSystemStateInnerV2) {
+        self.epoch = self.epoch + 1
+    }
+
     #[test_only]
     public(package) fun request_add_validator_for_testing(
         self: &mut SuiSystemStateInnerV2,

crates/sui-framework/packages/sui-system/tests/rewards_distribution_tests.move

@@ -6,6 +6,7 @@ module sui_system::rewards_distribution_tests {
     use sui::balance;
     use sui::test_scenario::{Self, Scenario};
     use sui_system::sui_system::SuiSystemState;
+    use sui::coin;
     use sui_system::validator_cap::UnverifiedValidatorOperationCap;
     use sui_system::governance_test_utils::{
         advance_epoch,
@@ -19,6 +20,7 @@ module sui_system::rewards_distribution_tests {
         stake_with,
         total_sui_balance, unstake
     };
+    use sui_system::staking_pool::StakedSui;
     use sui::test_utils::assert_eq;
     use sui::address;
 
@@ -628,4 +630,59 @@ module sui_system::rewards_distribution_tests {
         test_scenario::return_shared(sui_system);
         test.end();
     }
+
+    #[test]
+    fun test_process_pending_stake_withdraw_no_underflow_in_safe_mode() {
+        let start_epoch: u64 = 1;
+        let epoch_start_time = 100000000000;
+
+        set_up_sui_system_state();
+        let mut scenario_val = test_scenario::begin(VALIDATOR_ADDR_1); // val 1 has 100 staked sui
+        let mut sui_system = scenario_val.take_shared<SuiSystemState>();
+
+        start_epoch.do!(|_| scenario_val.ctx().increment_epoch_number()); // epoch 1, entering safe mode
+        sui_system.set_epoch_for_testing(start_epoch);
+
+        // staker 1 stakes 101 sui in safe mode
+        scenario_val.next_tx(STAKER_ADDR_1);
+        sui_system.request_add_stake(coin::mint_for_testing(101 * MIST_PER_SUI, scenario_val.ctx()), VALIDATOR_ADDR_1, scenario_val.ctx());
+        scenario_val.next_tx(STAKER_ADDR_1);
+
+        scenario_val.ctx().increment_epoch_number(); // epoch 2: still in safe mode
+        // There is no need to update `sui_system.set_epoch_for_testing` because
+        // it's `ctx.epoch()` being used here
+        sui_system.increment_epoch_for_testing();
+
+        let staked_sui = scenario_val.take_from_address<StakedSui>(STAKER_ADDR_1);
+        // staker 1 unstakes in safe mode
+        sui_system.request_withdraw_stake(staked_sui, scenario_val.ctx());
+        let val = sui_system.active_validator_by_address(VALIDATOR_ADDR_1);
+        let pool = val.get_staking_pool_ref();
+        assert!(pool.sui_balance() == 100 * MIST_PER_SUI, 0);
+        assert!(pool.pool_token_balance() == 100 * MIST_PER_SUI, 0);
+        // FIXME: these 3 values will be fixed in the next PR
+        assert!(pool.pending_stake_amount() == 101 * MIST_PER_SUI, 0);
+        assert!(pool.pending_total_sui_withdraw() == 101 * MIST_PER_SUI, 0);
+        assert!(pool.pending_pool_token_withdraw() == 101 * MIST_PER_SUI, 0);
+
+        // epoch 3: exiting safe mode
+        // There is no underflow here
+        sui_system
+            .inner_mut_for_testing()
+            .advance_epoch(start_epoch + 2, 65, balance::zero(), balance::zero(), 0, 0, 0, 0, epoch_start_time, scenario_val.ctx())
+            .destroy_for_testing(); // balance returned from `advance_epoch`
+
+        scenario_val.next_tx(VALIDATOR_ADDR_1);
+        let val = sui_system.active_validator_by_address(VALIDATOR_ADDR_1);
+        let pool = val.get_staking_pool_ref();
+        assert!(pool.pending_stake_amount() == 0 * MIST_PER_SUI, 0);
+        assert!(pool.pending_total_sui_withdraw() == 0 * MIST_PER_SUI, 0);
+        assert!(pool.pending_pool_token_withdraw() == 0 * MIST_PER_SUI, 0);
+        // FIXME: these 2 values will be fixed in the next PR
+        assert!(pool.sui_balance() == 101 * MIST_PER_SUI, 0);
+        assert!(pool.pool_token_balance() == 101 * MIST_PER_SUI, 0);
+
+        test_scenario::return_shared(sui_system);
+        scenario_val.end();
+    }
 }

crates/sui-framework/packages/sui-system/tests/staking_pool.move

@@ -115,6 +115,31 @@ module sui_system::staking_pool_tests {
         scenario.end();
     }
 
+    #[test]
+    fun test_process_pending_stake_withdraw_no_underflow() {
+        let mut scenario = test_scenario::begin(@0x0);
+        let mut staking_pool = staking_pool::new(scenario.ctx());
+        staking_pool.activate_staking_pool(0);
+
+        let sui = balance::create_for_testing(1_000_000_000);
+        let staked_sui_1 = staking_pool.request_add_stake(sui, scenario.ctx().epoch() + 1, scenario.ctx());
+        assert!(distribute_rewards_and_advance_epoch(&mut staking_pool, &mut scenario, 0) == 1, 0);
+
+        staking_pool.increase_pending_pool_token_withdraw_test_only(1_000_000_000);
+        staking_pool.increase_pending_total_sui_withdraw_test_only(1_000_000_000);
+
+        staking_pool.process_pending_stake_withdraw_test_only();
+
+        assert!(staking_pool.sui_balance() == 0, 0);
+        assert!(staking_pool.pending_total_sui_withdraw() == 0, 0);
+        assert!(staking_pool.pool_token_balance() == 0, 0);
+        assert!(staking_pool.pending_pool_token_withdraw() == 0, 0);
+
+        sui::test_utils::destroy(staking_pool);
+        sui::test_utils::destroy(staked_sui_1);
+        scenario.end();
+    }
+
     #[test]
     fun test_convert_to_fungible_staked_sui_happy() {
         let mut scenario = test_scenario::begin(@0x0);