Adding missing integration tests

programs/steward/src/delegation.rs

@@ -241,7 +241,7 @@ impl UnstakeState {
         })
     }
 
-    fn stake_deposit_unstake(
+    pub fn stake_deposit_unstake(
         &self,
         state: &StewardState,
         index: usize,
@@ -276,7 +276,7 @@ impl UnstakeState {
         Ok(0)
     }
 
-    fn instant_unstake(
+    pub fn instant_unstake(
         &self,
         state: &StewardState,
         index: usize,
@@ -303,7 +303,7 @@ impl UnstakeState {
         Ok(0)
     }
 
-    fn scoring_unstake(&self, current_lamports: u64, target_lamports: u64) -> Result<u64> {
+    pub fn scoring_unstake(&self, current_lamports: u64, target_lamports: u64) -> Result<u64> {
         // If there are additional lamports to unstake on this validator and the total unstaked lamports is below the cap, destake to the target
         if self.scoring_unstake_total < self.scoring_unstake_cap {
             let lamports_above_target = current_lamports

programs/steward/src/instructions/auto_remove_validator_from_pool.rs

@@ -6,11 +6,11 @@ use crate::events::AutoRemoveValidatorEvent;
 use crate::state::Config;
 use crate::utils::{
     deserialize_stake_pool, get_stake_pool_address, get_validator_stake_info_at_index,
-    remove_validator_check,
+    remove_validator_check, stake_is_inactive_without_history, stake_is_usable_by_pool,
 };
 use crate::StewardStateAccount;
 use anchor_lang::prelude::*;
-use anchor_lang::solana_program::{clock::Epoch, program::invoke_signed, stake, sysvar, vote};
+use anchor_lang::solana_program::{program::invoke_signed, stake, sysvar, vote};
 use spl_pod::solana_program::borsh1::try_from_slice_unchecked;
 use spl_pod::solana_program::stake::state::StakeStateV2;
 use spl_stake_pool::state::StakeStatus;
@@ -243,12 +243,14 @@ pub fn handler(ctx: Context<AutoRemoveValidator>, validator_list_index: usize) -
                 }
             }
             StakeStatus::ReadyForRemoval => {
+                // Should never happen but this is logical action
                 marked_for_immediate_removal = true;
                 state_account
                     .state
                     .mark_validator_for_immediate_removal(validator_list_index)?;
             }
             StakeStatus::DeactivatingAll | StakeStatus::DeactivatingTransient => {
+                // DeactivatingTransient should not be possible but this is the logical action
                 marked_for_immediate_removal = false;
                 state_account
                     .state
@@ -267,23 +269,3 @@ pub fn handler(ctx: Context<AutoRemoveValidator>, validator_list_index: usize) -
 
     Ok(())
 }
-
-// CHECKS FROM spl_stake_pool::processor::update_validator_list_balance
-
-/// Checks if a stake account can be managed by the pool
-fn stake_is_usable_by_pool(
-    meta: &stake::state::Meta,
-    expected_authority: &Pubkey,
-    expected_lockup: &stake::state::Lockup,
-) -> bool {
-    meta.authorized.staker == *expected_authority
-        && meta.authorized.withdrawer == *expected_authority
-        && meta.lockup == *expected_lockup
-}
-
-/// Checks if a stake account is active, without taking into account cooldowns
-fn stake_is_inactive_without_history(stake: &stake::state::Stake, epoch: Epoch) -> bool {
-    stake.delegation.deactivation_epoch < epoch
-        || (stake.delegation.activation_epoch == epoch
-            && stake.delegation.deactivation_epoch == epoch)
-}

programs/steward/src/instructions/epoch_maintenance.rs

@@ -85,7 +85,7 @@ pub fn handler(
         if let Some(validator_index_to_remove) = validator_index_to_remove {
             state_account
                 .state
-                .remove_validator(validator_index_to_remove, validators_in_list)?;
+                .remove_validator(validator_index_to_remove)?;
         }
     }
 

programs/steward/src/instructions/instant_remove_validator.rs

@@ -88,7 +88,7 @@ pub fn handler(
 
     state_account
         .state
-        .remove_validator(validator_index_to_remove, validators_in_list)?;
+        .remove_validator(validator_index_to_remove)?;
 
     Ok(())
 }

programs/steward/src/instructions/spl_passthrough.rs

@@ -9,11 +9,13 @@ use crate::errors::StewardError;
 use crate::state::Config;
 use crate::utils::{
     deserialize_stake_pool, get_config_admin, get_stake_pool_address,
-    get_validator_stake_info_at_index,
+    get_validator_stake_info_at_index, stake_is_inactive_without_history, stake_is_usable_by_pool,
 };
 use crate::StewardStateAccount;
 use anchor_lang::prelude::*;
 use anchor_lang::solana_program::{program::invoke_signed, stake, sysvar, vote};
+use spl_pod::solana_program::borsh1::try_from_slice_unchecked;
+use spl_pod::solana_program::stake::state::StakeStateV2;
 use spl_stake_pool::find_stake_program_address;
 use spl_stake_pool::instruction::PreferredValidatorType;
 use spl_stake_pool::state::{StakeStatus, ValidatorListHeader};
@@ -180,9 +182,9 @@ pub fn remove_validator_from_pool_handler(
     ctx: Context<RemoveValidatorFromPool>,
     validator_list_index: usize,
 ) -> Result<()> {
+    let epoch = Clock::get()?.epoch;
     {
         let state_account = ctx.accounts.state_account.load_mut()?;
-        let epoch = Clock::get()?.epoch;
 
         // Should not be able to remove a validator if update is not complete
         require!(
@@ -245,12 +247,37 @@ pub fn remove_validator_from_pool_handler(
 
         let stake_status = StakeStatus::try_from(validator_stake_info.status)?;
 
+        let stake_pool = deserialize_stake_pool(&ctx.accounts.stake_pool)?;
+
         match stake_status {
             StakeStatus::Active => {
                 // Should never happen
                 return Err(StewardError::ValidatorMarkedActive.into());
             }
-            StakeStatus::DeactivatingValidator | StakeStatus::ReadyForRemoval => {
+            StakeStatus::DeactivatingValidator => {
+                let stake_account_data = &mut ctx.accounts.stake_account.data.borrow_mut();
+                let (meta, stake) =
+                    match try_from_slice_unchecked::<StakeStateV2>(stake_account_data)? {
+                        StakeStateV2::Stake(meta, stake, _stake_flags) => (meta, stake),
+                        _ => return Err(StewardError::StakeStateIsNotStake.into()),
+                    };
+
+                if stake_is_usable_by_pool(
+                    &meta,
+                    ctx.accounts.withdraw_authority.key,
+                    &stake_pool.lockup,
+                ) && stake_is_inactive_without_history(&stake, epoch)
+                {
+                    state_account
+                        .state
+                        .mark_validator_for_immediate_removal(validator_list_index)?;
+                } else {
+                    state_account
+                        .state
+                        .mark_validator_for_removal(validator_list_index)?;
+                }
+            }
+            StakeStatus::ReadyForRemoval => {
                 state_account
                     .state
                     .mark_validator_for_immediate_removal(validator_list_index)?;

programs/steward/src/state/steward_state.rs

@@ -460,7 +460,7 @@ impl StewardState {
     }
 
     /// Update internal state when a validator is removed from the pool
-    pub fn remove_validator(&mut self, index: usize, validator_list_len: usize) -> Result<()> {
+    pub fn remove_validator(&mut self, index: usize) -> Result<()> {
         let marked_for_regular_removal = self.validators_to_remove.get(index)?;
         let marked_for_immediate_removal = self.validators_for_immediate_removal.get(index)?;
 
@@ -469,8 +469,16 @@ impl StewardState {
             StewardError::ValidatorNotMarkedForRemoval
         );
 
+        let num_pool_validators = self.num_pool_validators as usize;
+        let num_pool_validators_plus_added = num_pool_validators + self.validators_added as usize;
+
+        require!(
+            index < num_pool_validators_plus_added,
+            StewardError::ValidatorIndexOutOfBounds
+        );
+
         // If the validator was marked for removal in the current cycle, decrement validators_added
-        if index >= self.num_pool_validators as usize {
+        if index >= num_pool_validators {
             self.validators_added = self
                 .validators_added
                 .checked_sub(1)
@@ -482,8 +490,6 @@ impl StewardState {
                 .ok_or(StewardError::ArithmeticError)?;
         }
 
-        let num_pool_validators = self.num_pool_validators as usize;
-
         // Shift all validator state to the left
         for i in index..num_pool_validators {
             let next_i = i.checked_add(1).ok_or(StewardError::ArithmeticError)?;
@@ -497,7 +503,7 @@ impl StewardState {
         }
 
         // For state that can be valid past num_pool_validators, we still need to shift the values
-        for i in index..validator_list_len {
+        for i in index..num_pool_validators_plus_added {
             let next_i = i.checked_add(1).ok_or(StewardError::ArithmeticError)?;
             self.validators_to_remove
                 .set(i, self.validators_to_remove.get(next_i)?)?;
@@ -551,9 +557,10 @@ impl StewardState {
         self.delegations[num_pool_validators] = Delegation::default();
         self.instant_unstake.set(num_pool_validators, false)?;
         self.progress.set(num_pool_validators, false)?;
-        self.validators_to_remove.set(validator_list_len, false)?;
+        self.validators_to_remove
+            .set(num_pool_validators_plus_added, false)?;
         self.validators_for_immediate_removal
-            .set(validator_list_len, false)?;
+            .set(num_pool_validators_plus_added, false)?;
 
         Ok(())
     }

programs/steward/src/utils.rs

@@ -1,7 +1,9 @@
 use std::ops::{Deref, Not};
 
+use anchor_lang::solana_program::stake;
 use anchor_lang::{idl::types::*, prelude::*};
 use borsh::{BorshDeserialize, BorshSerialize};
+use spl_pod::solana_program::clock::Epoch;
 use spl_pod::{bytemuck::pod_from_bytes, primitives::PodU64, solana_program::program_pack::Pack};
 use spl_stake_pool::{
     big_vec::BigVec,
@@ -310,6 +312,26 @@ pub fn check_validator_list_has_stake_status_other_than(
     Ok(false)
 }
 
+/// Checks if a stake account can be managed by the pool
+/// FROM spl_stake_pool::processor::update_validator_list_balance
+pub fn stake_is_usable_by_pool(
+    meta: &stake::state::Meta,
+    expected_authority: &Pubkey,
+    expected_lockup: &stake::state::Lockup,
+) -> bool {
+    meta.authorized.staker == *expected_authority
+        && meta.authorized.withdrawer == *expected_authority
+        && meta.lockup == *expected_lockup
+}
+
+/// Checks if a stake account is active, without taking into account cooldowns
+/// FROM spl_stake_pool::processor::update_validator_list_balance
+pub fn stake_is_inactive_without_history(stake: &stake::state::Stake, epoch: Epoch) -> bool {
+    stake.delegation.deactivation_epoch < epoch
+        || (stake.delegation.activation_epoch == epoch
+            && stake.delegation.deactivation_epoch == epoch)
+}
+
 pub fn get_validator_list_length(validator_list_account_info: &AccountInfo) -> Result<usize> {
     let mut validator_list_data = validator_list_account_info.try_borrow_mut_data()?;
     let (header, validator_list) = ValidatorListHeader::deserialize_vec(&mut validator_list_data)?;

run_tests.sh

@@ -5,5 +5,5 @@ cargo build-sbf --manifest-path programs/steward/Cargo.toml;
 cargo build-sbf --manifest-path programs/validator-history/Cargo.toml;
 
 # Run all tests
-SBF_OUT_DIR=$(pwd)/target/deploy RUST_MIN_STACK=5000000 cargo test --package tests --all-features --color auto
+SBF_OUT_DIR=$(pwd)/target/deploy RUST_MIN_STACK=10000000 cargo test --package tests --all-features --color auto