diff --git a/contracts/sources/storage.move b/contracts/sources/storage.move index b992309..75ddc22 100644 --- a/contracts/sources/storage.move +++ b/contracts/sources/storage.move @@ -12,6 +12,7 @@ module liquid_staking::storage { /* Constants */ const MIN_STAKE_THRESHOLD: u64 = 1_000_000_000; + const MAX_SUI_SUPPLY: u64 = 10_000_000_000 * 1_000_000_000; /// The Storage struct holds all stake for the LST. public struct Storage has store { @@ -131,32 +132,41 @@ module liquid_staking::storage { return false }; + let active_validator_addresses = system_state.active_validator_addresses(); + let mut i = self.validator_infos.length(); while (i > 0) { i = i - 1; - // update pool token exchange rates - let validator_info = &mut self.validator_infos[i]; + // if validator is inactive, withdraw all stake. + if (!active_validator_addresses.contains(&self.validator_infos[i].validator_address)) { + // technically this is using a stale exchange rate, but it doesn't matter because we're unstaking everything. + // this is done before fetching the exchange rate because i don't want the function to abort if an epoch is skipped. + self.unstake_approx_n_sui_from_validator(system_state, i, MAX_SUI_SUPPLY, ctx); + }; + + if (self.validator_infos[i].is_empty()) { + let ValidatorInfo { active_stake, inactive_stake, extra_fields, .. } = self.validator_infos.remove(i); + active_stake.destroy_none(); + inactive_stake.destroy_none(); + extra_fields.destroy_empty(); - let exchange_rates = system_state.pool_exchange_rates(&validator_info.staking_pool_id); + continue + }; + + // update pool token exchange rates + let exchange_rates = system_state.pool_exchange_rates(&self.validator_infos[i].staking_pool_id); let latest_exchange_rate = exchange_rates.borrow(ctx.epoch()); - validator_info.exchange_rate = *latest_exchange_rate; + self.validator_infos[i].exchange_rate = *latest_exchange_rate; + self.refresh_validator_info(i); - if (validator_info.inactive_stake.is_some()) { + if (self.validator_infos[i].inactive_stake.is_some()) { let inactive_stake = self.take_from_inactive_stake(i); let fungible_staked_sui = system_state.convert_to_fungible_staked_sui(inactive_stake, ctx); self.join_fungible_staked_sui_to_validator(i, fungible_staked_sui); }; - refresh_validator_info(self, i); - - if (self.validator_infos[i].is_empty()) { - let ValidatorInfo { active_stake, inactive_stake, extra_fields, .. } = self.validator_infos.remove(i); - active_stake.destroy_none(); - inactive_stake.destroy_none(); - extra_fields.destroy_empty(); - }; }; self.last_refresh_epoch = ctx.epoch(); diff --git a/contracts/tests/storage_tests.move b/contracts/tests/storage_tests.move index 1e53a19..9d96186 100644 --- a/contracts/tests/storage_tests.move +++ b/contracts/tests/storage_tests.move @@ -310,6 +310,40 @@ module liquid_staking::storage_tests { scenario.end(); } + #[test] + fun test_refresh_with_inactive_stake() { + let mut scenario = test_scenario::begin(@0x0); + + setup_sui_system(&mut scenario, vector[100, 100]); + + let staked_sui_1 = stake_with(1, 100, &mut scenario); + + let mut system_state = scenario.take_shared(); + let mut storage = new(scenario.ctx()); + storage.join_stake(&mut system_state, staked_sui_1, scenario.ctx()); + test_scenario::return_shared(system_state); + + scenario.next_tx(@0x1); + let mut system_state = scenario.take_shared(); + system_state.request_remove_validator(scenario.ctx()); + test_scenario::return_shared(system_state); + + advance_epoch_with_reward_amounts(0, 0, &mut scenario); + advance_epoch_with_reward_amounts(0, 0, &mut scenario); + + let mut system_state = scenario.take_shared(); + assert!(!system_state.active_validator_addresses().contains(&@0x1), 0); + + storage.refresh(&mut system_state, scenario.ctx()); + assert!(storage.validators().length() == 0, 0); // got removed + assert!(storage.total_sui_supply() == 100 * MIST_PER_SUI, 0); + + test_scenario::return_shared(system_state); + sui::test_utils::destroy(storage); + + scenario.end(); + } + #[test] #[expected_failure(abort_code = 1, location = liquid_staking::storage)] fun test_join_active_stake_from_non_active_validator() {