diff --git a/crates/sui-framework/packages/sui-system/sources/staking_pool.move b/crates/sui-framework/packages/sui-system/sources/staking_pool.move index ba10693c3dccf..0274e8e061435 100644 --- a/crates/sui-framework/packages/sui-system/sources/staking_pool.move +++ b/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() diff --git a/crates/sui-framework/packages/sui-system/sources/sui_system.move b/crates/sui-framework/packages/sui-system/sources/sui_system.move index 641cc3f9dd86e..a3d594222d28d 100644 --- a/crates/sui-framework/packages/sui-system/sources/sui_system.move +++ b/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, diff --git a/crates/sui-framework/packages/sui-system/sources/sui_system_state_inner.move b/crates/sui-framework/packages/sui-system/sources/sui_system_state_inner.move index 4384f3fcd2599..b79bf42e38339 100644 --- a/crates/sui-framework/packages/sui-system/sources/sui_system_state_inner.move +++ b/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, diff --git a/crates/sui-framework/packages/sui-system/tests/rewards_distribution_tests.move b/crates/sui-framework/packages/sui-system/tests/rewards_distribution_tests.move index ec94cbf81a1bf..a809c749b1cc5 100644 --- a/crates/sui-framework/packages/sui-system/tests/rewards_distribution_tests.move +++ b/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(); + + 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(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(); + } } diff --git a/crates/sui-framework/packages/sui-system/tests/staking_pool.move b/crates/sui-framework/packages/sui-system/tests/staking_pool.move index 2a38e5f249aa8..242f8a265bfd9 100644 --- a/crates/sui-framework/packages/sui-system/tests/staking_pool.move +++ b/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);