From f8b8b00d26ee85c2cc4c1dee51e0846910848a65 Mon Sep 17 00:00:00 2001
From: 0xripleys <0xripleys@solend.fi>
Date: Fri, 18 Oct 2024 06:57:18 -0400
Subject: [PATCH] remove inactive validators on storage refresh

---
 contracts/sources/storage.move     | 36 +++++++++++++++++++-----------
 contracts/tests/storage_tests.move | 34 ++++++++++++++++++++++++++++
 2 files changed, 57 insertions(+), 13 deletions(-)

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<SuiSystemState>();
+        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<SuiSystemState>();
+        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<SuiSystemState>();
+        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() {