@@ -7,9 +7,15 @@ module liquid_staking::storage {
77 use sui::bag::{Self , Bag };
88
99 /* Errors */
10+ const ENotEnoughSuiInSuiPool : u64 = 0 ;
11+ const ENotActiveValidator : u64 = 1 ;
12+ const ETooManyValidators : u64 = 2 ;
13+ const EValidatorAlreadyExists : u64 = 3 ;
1014
1115 /* Constants */
1216 const MIN_STAKE_THRESHOLD : u64 = 1_000_000_000 ;
17+ const MAX_SUI_SUPPLY : u64 = 10_000_000_000 * 1_000_000_000 ;
18+ const MAX_VALIDATORS : u64 = 50 ;
1319
1420 /// The Storage struct holds all stake for the LST.
1521 public struct Storage has store {
@@ -129,32 +135,41 @@ module liquid_staking::storage {
129135 return false
130136 };
131137
138+ let active_validator_addresses = system_state.active_validator_addresses ();
139+
132140 let mut i = self.validator_infos.length ();
133141 while (i > 0 ) {
134142 i = i - 1 ;
135143
136- // update pool token exchange rates
137- let validator_info = &mut self.validator_infos[i];
144+ // if validator is inactive, withdraw all stake.
145+ if (!active_validator_addresses.contains (&self.validator_infos[i].validator_address)) {
146+ // technically this is using a stale exchange rate, but it doesn't matter because we're unstaking everything.
147+ // this is done before fetching the exchange rate because i don't want the function to abort if an epoch is skipped.
148+ self.unstake_approx_n_sui_from_validator (system_state, i, MAX_SUI_SUPPLY , ctx);
149+ };
138150
139- let exchange_rates = system_state.pool_exchange_rates (&validator_info.staking_pool_id);
151+ if (self.validator_infos[i].is_empty ()) {
152+ let ValidatorInfo { active_stake, inactive_stake, extra_fields, .. } = self.validator_infos.remove (i);
153+ active_stake.destroy_none ();
154+ inactive_stake.destroy_none ();
155+ extra_fields.destroy_empty ();
156+
157+ continue
158+ };
159+
160+ // update pool token exchange rates
161+ let exchange_rates = system_state.pool_exchange_rates (&self.validator_infos[i].staking_pool_id);
140162 let latest_exchange_rate = exchange_rates.borrow (ctx.epoch ());
141163
142- validator_info.exchange_rate = *latest_exchange_rate;
164+ self.validator_infos[i].exchange_rate = *latest_exchange_rate;
165+ self.refresh_validator_info (i);
143166
144- if (validator_info .inactive_stake.is_some ()) {
167+ if (self.validator_infos[i] .inactive_stake.is_some ()) {
145168 let inactive_stake = self.take_from_inactive_stake (i);
146169 let fungible_staked_sui = system_state.convert_to_fungible_staked_sui (inactive_stake, ctx);
147170 self.join_fungible_staked_sui_to_validator (i, fungible_staked_sui);
148171 };
149172
150- refresh_validator_info (self, i);
151-
152- if (self.validator_infos[i].is_empty ()) {
153- let ValidatorInfo { active_stake, inactive_stake, extra_fields, .. } = self.validator_infos.remove (i);
154- active_stake.destroy_none ();
155- inactive_stake.destroy_none ();
156- extra_fields.destroy_empty ();
157- };
158173 };
159174
160175 self.last_refresh_epoch = ctx.epoch ();
@@ -373,7 +388,7 @@ module liquid_staking::storage {
373388 let target_unstake_sui_amount = max (target_unstake_sui_amount, MIN_STAKE_THRESHOLD );
374389
375390 let staked_sui_amount = validator_info.inactive_stake.borrow ().staked_sui_amount ();
376- let staked_sui = if (staked_sui_amount <= target_unstake_sui_amount + MIN_STAKE_THRESHOLD ) {
391+ let staked_sui = if (staked_sui_amount < target_unstake_sui_amount + MIN_STAKE_THRESHOLD ) {
377392 self.take_from_inactive_stake (validator_index)
378393 }
379394 else {
@@ -424,6 +439,7 @@ module liquid_staking::storage {
424439 };
425440 };
426441
442+ assert ! (self.sui_pool.value () >= max_sui_amount_out, ENotEnoughSuiInSuiPool );
427443 self.split_from_sui_pool (max_sui_amount_out)
428444 }
429445
@@ -495,16 +511,31 @@ module liquid_staking::storage {
495511 staking_pool_id: ID ,
496512 ctx: &mut TxContext
497513 ): u64 {
514+ let mut current_validator_addresses = vector [];
515+
498516 let mut i = 0 ;
499517 while (i < self.validator_infos.length ()) {
500518 if (self.validator_infos[i].staking_pool_id == staking_pool_id) {
501519 return i
502520 };
503521
522+ current_validator_addresses.push_back (self.validator_infos[i].validator_address);
504523 i = i + 1 ;
505524 };
506525
507526 let validator_address = system_state.validator_address_by_pool_id (&staking_pool_id);
527+
528+ assert ! (
529+ !current_validator_addresses.contains (&validator_address),
530+ EValidatorAlreadyExists
531+ );
532+
533+ let active_validator_addresses = system_state.active_validator_addresses ();
534+ assert ! (
535+ active_validator_addresses.contains (&validator_address),
536+ ENotActiveValidator
537+ );
538+
508539 let exchange_rates = system_state.pool_exchange_rates (&staking_pool_id);
509540 let latest_exchange_rate = exchange_rates.borrow (ctx.epoch ());
510541
@@ -518,6 +549,8 @@ module liquid_staking::storage {
518549 extra_fields: bag::new (ctx)
519550 });
520551
552+ assert ! (self.validator_infos.length () <= MAX_VALIDATORS , ETooManyValidators );
553+
521554 i
522555 }
523556
0 commit comments