Audit recommendations

contracts/sources/liquid_staking.move

@@ -327,6 +327,9 @@ module liquid_staking::liquid_staking {
         };
         self.fees.join(sui.split(redeem_fee_amount as u64));
 
+        // Abort if fees consumed entire redemption amount
+        assert!(sui.value() > 0, ERedeemInvariantViolated);
+
         emit_event(RedeemEvent {
             typename: type_name::get<P>(),
             lst_amount_in: lst.value(),

contracts/sources/storage.move

@@ -148,6 +148,8 @@ module liquid_staking::storage {
                 // 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);
+                // Reconcile total_sui_amount accounting after full unstake
+                self.refresh_validator_info(i);
             };
 
             if (self.validator_infos[i].is_empty()) {
@@ -373,8 +375,6 @@ module liquid_staking::storage {
             return 0
         };
 
-        let target_unstake_sui_amount = max(target_unstake_sui_amount, MIN_STAKE_THRESHOLD);
-
         let fungible_staked_sui_amount = validator_info.active_stake.borrow().value();
         let total_sui_amount = get_sui_amount(
             &validator_info.exchange_rate, 
@@ -495,9 +495,7 @@ module liquid_staking::storage {
             .borrow_mut()
             .split_fungible_staked_sui(fungible_staked_sui_amount, ctx);
 
-        self.refresh_validator_info(validator_index);
-
-        system_state.redeem_fungible_staked_sui(stake, ctx)
+        self.redeem_and_update_accounting(system_state, validator_index, stake, ctx)
     }
 
     fun take_active_stake(
@@ -509,9 +507,7 @@ module liquid_staking::storage {
         let validator_info = &mut self.validator_infos[validator_index];
         let fungible_staked_sui = validator_info.active_stake.extract();
 
-        self.refresh_validator_info(validator_index);
-
-        system_state.redeem_fungible_staked_sui(fungible_staked_sui, ctx)
+        self.redeem_and_update_accounting(system_state, validator_index, fungible_staked_sui, ctx)
     }
 
     fun split_from_inactive_stake(
@@ -542,6 +538,30 @@ module liquid_staking::storage {
         stake
     }
 
+    /// Redeems fungible staked SUI and updates accounting based on actual redeemed amount.
+    fun redeem_and_update_accounting(
+        self: &mut Storage,
+        system_state: &mut SuiSystemState,
+        validator_index: u64,
+        fungible_staked_sui: FungibleStakedSui,
+        ctx: &TxContext
+    ): Balance<SUI> {
+        let redeemed_sui = system_state.redeem_fungible_staked_sui(fungible_staked_sui, ctx);
+        let redeemed_amount = redeemed_sui.value();
+
+        // Update accounting based on actual redeemed amount.
+        // Use min() to handle the case where the redeemed amount exceeds the tracked
+        // amount due to a stale exchange rate (e.g. inactive validator cleanup).
+        // Any excess is untracked reward income that enters total_sui_supply via join_to_sui_pool.
+        // The validator's total_sui_amount will be reconciled during the next refresh.
+        let validator_info = &mut self.validator_infos[validator_index];
+        let tracked_decrease = redeemed_amount.min(validator_info.total_sui_amount);
+        self.total_sui_supply = self.total_sui_supply - tracked_decrease;
+        validator_info.total_sui_amount = validator_info.total_sui_amount - tracked_decrease;
+
+        redeemed_sui
+    }
+
     /* Private functions */
     fun get_or_add_validator_index_by_staking_pool_id_mut(
         self: &mut Storage, 
@@ -575,7 +595,13 @@ module liquid_staking::storage {
         );
 
         let exchange_rates = system_state.pool_exchange_rates(&staking_pool_id);
-        let latest_exchange_rate = exchange_rates.borrow(ctx.epoch());
+        // Search backwards to handle newly activated validators that may only
+        // have an exchange rate at their activation epoch, not the current epoch.
+        let mut cur_epoch = ctx.epoch();
+        while (!exchange_rates.contains(cur_epoch)) {
+            cur_epoch = cur_epoch - 1;
+        };
+        let latest_exchange_rate = exchange_rates.borrow(cur_epoch);
 
         self.validator_infos.push_back(ValidatorInfo {
             staking_pool_id: copy staking_pool_id,

contracts/tests/liquid_staking_tests.move

@@ -811,4 +811,50 @@ module liquid_staking::liquid_staking_tests {
 
         scenario.end(); 
     }
+
+    #[test]
+    fun test_exchange_rate_gap() {
+        use sui_system::sui_system;
+        use std::unit_test;
+
+        let mut scenario = test_scenario::begin(@0x0);
+
+        // activate validator
+        setup_sui_system(&mut scenario, vector[100]);
+
+        scenario.next_tx(@0x0);
+        let mut system_state = scenario.take_shared<SuiSystemState>();
+        sui_system::set_epoch_for_testing(&mut system_state, scenario.ctx().epoch() + 1);
+        scenario.next_epoch(@0x0);
+
+        let (admin_cap, mut lst_info) = create_lst<TEST>(
+            fees::new_builder(scenario.ctx())
+                .set_sui_mint_fee_bps(100)
+                .set_redeem_fee_bps(100)
+                .to_fee_config(),
+            coin::create_treasury_cap_for_testing(scenario.ctx()),
+            scenario.ctx()
+        );
+
+        assert!(lst_info.total_lst_supply() == 0, 0);
+        assert!(lst_info.storage().total_sui_supply() == 0, 0);
+
+        let sui = coin::mint_for_testing<SUI>(200 * MIST_PER_SUI, scenario.ctx());
+        let lst = lst_info.mint(&mut system_state, sui, scenario.ctx());
+
+        lst_info.increase_validator_stake(
+            &admin_cap, 
+            &mut system_state, 
+            @0x0,
+            200 * MIST_PER_SUI, 
+            scenario.ctx()
+        );
+
+        test_scenario::return_shared(system_state);
+        unit_test::destroy(admin_cap);
+        unit_test::destroy(lst);
+        unit_test::destroy(lst_info);
+
+        scenario.end();
+    }
 }

contracts/tests/storage_tests.move

@@ -943,12 +943,12 @@ module liquid_staking::storage_tests {
             scenario.ctx()
         );
 
-        assert!(amount == MIST_PER_SUI, 0);
+        assert!(amount == 2, 0);
         assert!(storage.validators().length() == 1, 0);
         assert!(storage.total_sui_supply() == 200 * MIST_PER_SUI, 0);
-        assert!(storage.sui_pool().value() == MIST_PER_SUI, 0);
-        assert!(storage.validators()[0].total_sui_amount() == 199 * MIST_PER_SUI, 0);
-        assert!(storage.validators()[0].active_stake().borrow().value() == 99_500_000_000, 0);
+        assert!(storage.sui_pool().value() == 2, 0);
+        assert!(storage.validators()[0].total_sui_amount() == 200 * MIST_PER_SUI - 2, 0);
+        assert!(storage.validators()[0].active_stake().borrow().value() == 100 * MIST_PER_SUI - 1, 0);
         assert!(storage.validators()[0].inactive_stake().is_none(), 0);
 
         sui::test_utils::destroy(storage);
@@ -1456,4 +1456,103 @@ module liquid_staking::storage_tests {
         sui::test_utils::destroy(storage);
         scenario.end();
     }
-}   
+
+    #[test]
+    fun test_unstake_rounding_loss_regression() {
+        let mut scenario = test_scenario::begin(@0x0);
+
+        // Single validator with 500 SUI
+        setup_sui_system(&mut scenario, vector[500]);
+
+        // User stakes 200 SUI → validator has 700 tokens total
+        let staked_sui = stake_with(0, 200, &mut scenario);
+        advance_epoch_with_reward_amounts(0, 0, &mut scenario);
+        // Add 400 SUI rewards → rate becomes 1100:700 = 11:7
+        advance_epoch_with_reward_amounts(0, 400, &mut scenario);
+
+        scenario.next_tx(@0x0);
+
+        let mut system_state = scenario.take_shared<SuiSystemState>();
+        let mut storage = new(scenario.ctx());
+
+        storage.refresh(&mut system_state, scenario.ctx());
+        storage.join_stake(&mut system_state, staked_sui, scenario.ctx());
+
+        let initial_total = storage.total_sui_supply();
+        let initial_validator_total = storage.validators()[0].total_sui_amount();
+        let initial_sui_pool = storage.sui_pool().value();
+
+        // Unstake minimum (1 SUI due to MIN_STAKE_THRESHOLD)
+        let redeemed = storage.unstake_approx_n_sui_from_active_stake(
+            &mut system_state,
+            0,
+            MIST_PER_SUI,
+            scenario.ctx()
+        );
+
+        let final_total = storage.total_sui_supply();
+        let final_validator_total = storage.validators()[0].total_sui_amount();
+        let final_sui_pool = storage.sui_pool().value();
+
+        // total_sui_supply should not change during unstaking.
+        // We're just moving SUI from staked form to liquid form.
+        assert!(final_total == initial_total);
+
+        // The redeemed amount should equal the increase in sui_pool
+        assert!(final_sui_pool == initial_sui_pool + redeemed);
+
+        // validator_total decreases by exactly the redeemed amount
+        assert!(final_validator_total == initial_validator_total - redeemed);
+
+        // total_sui_supply = sui_pool + validator totals
+        assert!(final_total == final_sui_pool + final_validator_total);
+
+        test_scenario::return_shared(system_state);
+        sui::test_utils::destroy(storage);
+        scenario.end();
+    }
+
+    #[test]
+    fun test_refresh_inactive_validator_with_rewards() {
+        let mut scenario = test_scenario::begin(@0x0);
+
+        setup_sui_system(&mut scenario, vector[100, 100]);
+
+        let staked_sui = stake_with(1, 100, &mut scenario);
+
+        // Activate the stake
+        advance_epoch_with_reward_amounts(0, 0, &mut scenario);
+
+        let mut system_state = scenario.take_shared<SuiSystemState>();
+        let mut storage = new(scenario.ctx());
+        storage.join_stake(&mut system_state, staked_sui, scenario.ctx());
+        storage.refresh(&mut system_state, scenario.ctx());
+        test_scenario::return_shared(system_state);
+
+        // Earn rewards, then remove validator.
+        // Storage's exchange rate is now stale — it doesn't include these rewards.
+        advance_epoch_with_reward_amounts(0, 400, &mut scenario);
+
+        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);
+
+        // refresh() detects the inactive validator and unstakes everything.
+        // redeem_and_update_accounting subtracts the actual redeemed amount
+        // (at the current rate, which includes rewards) from
+        // total_sui_amount (which was computed at the stale rate)
+        scenario.next_tx(@0x0);
+        let mut system_state = scenario.take_shared<SuiSystemState>();
+        assert!(!system_state.active_validator_addresses().contains(&@0x1));
+
+        storage.refresh(&mut system_state, scenario.ctx());
+        assert!(storage.validators().length() == 0);
+
+        test_scenario::return_shared(system_state);
+        sui::test_utils::destroy(storage);
+        scenario.end();
+    }
+}