FEATURE: Adding re-usable crank_checks function (#58)

The new `crank_checks` function does all of the checks needed across each crank. Now with the new addition of looking for mismatched indices. --------- Co-authored-by: Christian Krueger <[email protected]> Co-authored-by: Evan Batsell <[email protected]>
jito-foundation · Jul 23, 2024 · a7dccae · a7dccae
1 parent 0425db0
commit a7dccae
Show file tree

Hide file tree

Showing 13 changed files with 217 additions and 156 deletions.
diff --git a/programs/steward/idl/steward.json b/programs/steward/idl/steward.json
@@ -330,6 +330,9 @@
         {
           "name": "state_account",
           "writable": true
+        },
+        {
+          "name": "validator_list"
         }
       ],
       "args": []
@@ -649,6 +652,9 @@
         {
           "name": "state_account",
           "writable": true
+        },
+        {
+          "name": "validator_list"
         }
       ],
       "args": []

diff --git a/programs/steward/src/instructions/compute_delegations.rs b/programs/steward/src/instructions/compute_delegations.rs
@@ -1,4 +1,4 @@
-use crate::errors::StewardError;
+use crate::utils::{get_validator_list, state_checks};
 use crate::{maybe_transition_and_emit, Config, StewardStateAccount, StewardStateEnum};
 use anchor_lang::prelude::*;
 
@@ -12,6 +12,10 @@ pub struct ComputeDelegations<'info> {
         bump
     )]
     pub state_account: AccountLoader<'info, StewardStateAccount>,
+
+    /// CHECK: Account owner checked, account type checked in get_validator_stake_info_at_index
+    #[account(address = get_validator_list(&config)?)]
+    pub validator_list: AccountInfo<'info>,
 }
 
 /*
@@ -21,33 +25,16 @@ It computes a share of the pool for each validator.
 pub fn handler(ctx: Context<ComputeDelegations>) -> Result<()> {
     let config = ctx.accounts.config.load()?;
     let mut state_account = ctx.accounts.state_account.load_mut()?;
-
     let clock = Clock::get()?;
     let epoch_schedule = EpochSchedule::get()?;
 
-    {
-        if config.is_paused() {
-            return Err(StewardError::StateMachinePaused.into());
-        }
-
-        require!(
-            matches!(
-                state_account.state.state_tag,
-                StewardStateEnum::ComputeDelegations
-            ),
-            StewardError::InvalidState
-        );
-
-        require!(
-            clock.epoch == state_account.state.current_epoch,
-            StewardError::EpochMaintenanceNotComplete
-        );
-
-        require!(
-            state_account.state.validators_for_immediate_removal.count() == 0,
-            StewardError::ValidatorsNeedToBeRemoved
-        );
-    }
+    state_checks(
+        &clock,
+        &config,
+        &state_account,
+        &ctx.accounts.validator_list,
+        Some(StewardStateEnum::ComputeDelegations),
+    )?;
 
     state_account
         .state

diff --git a/programs/steward/src/instructions/compute_instant_unstake.rs b/programs/steward/src/instructions/compute_instant_unstake.rs
@@ -1,7 +1,7 @@
 use crate::{
     errors::StewardError,
     maybe_transition_and_emit,
-    utils::{get_validator_list, get_validator_stake_info_at_index},
+    utils::{get_validator_list, get_validator_stake_info_at_index, state_checks},
     Config, StewardStateAccount, StewardStateEnum,
 };
 use anchor_lang::prelude::*;
@@ -41,29 +41,13 @@ pub fn handler(ctx: Context<ComputeInstantUnstake>, validator_list_index: usize)
     let clock = Clock::get()?;
     let epoch_schedule = EpochSchedule::get()?;
 
-    {
-        if config.is_paused() {
-            return Err(StewardError::StateMachinePaused.into());
-        }
-
-        require!(
-            matches!(
-                state_account.state.state_tag,
-                StewardStateEnum::ComputeInstantUnstake
-            ),
-            StewardError::InvalidState
-        );
-
-        require!(
-            clock.epoch == state_account.state.current_epoch,
-            StewardError::EpochMaintenanceNotComplete
-        );
-
-        require!(
-            state_account.state.validators_for_immediate_removal.count() == 0,
-            StewardError::ValidatorsNeedToBeRemoved
-        );
-    }
+    state_checks(
+        &clock,
+        &config,
+        &state_account,
+        &ctx.accounts.validator_list,
+        Some(StewardStateEnum::ComputeInstantUnstake),
+    )?;
 
     let validator_stake_info =
         get_validator_stake_info_at_index(validator_list, validator_list_index)?;

diff --git a/programs/steward/src/instructions/compute_score.rs b/programs/steward/src/instructions/compute_score.rs
@@ -3,7 +3,10 @@ use anchor_lang::prelude::*;
 use crate::{
     errors::StewardError,
     maybe_transition_and_emit,
-    utils::{get_validator_list, get_validator_list_length, get_validator_stake_info_at_index},
+    utils::{
+        get_validator_list, get_validator_list_length, get_validator_stake_info_at_index,
+        state_checks,
+    },
     Config, StewardStateAccount, StewardStateEnum,
 };
 use validator_history::{ClusterHistory, ValidatorHistory};
@@ -42,21 +45,7 @@ pub fn handler(ctx: Context<ComputeScore>, validator_list_index: usize) -> Resul
     let clock: Clock = Clock::get()?;
     let epoch_schedule = EpochSchedule::get()?;
 
-    {
-        if config.is_paused() {
-            return Err(StewardError::StateMachinePaused.into());
-        }
-
-        require!(
-            clock.epoch == state_account.state.current_epoch,
-            StewardError::EpochMaintenanceNotComplete
-        );
-
-        require!(
-            state_account.state.validators_for_immediate_removal.count() == 0,
-            StewardError::ValidatorsNeedToBeRemoved
-        );
-    }
+    state_checks(&clock, &config, &state_account, validator_list, None)?;
 
     let validator_stake_info =
         get_validator_stake_info_at_index(validator_list, validator_list_index)?;
@@ -65,8 +54,6 @@ pub fn handler(ctx: Context<ComputeScore>, validator_list_index: usize) -> Resul
         StewardError::ValidatorNotInList
     );
 
-    let num_pool_validators = get_validator_list_length(validator_list)?;
-
     // May need to force an extra transition here in case cranking got stuck in any previous state
     // and it's now the start of a new scoring cycle
     if !matches!(
@@ -89,6 +76,8 @@ pub fn handler(ctx: Context<ComputeScore>, validator_list_index: usize) -> Resul
         StewardError::InvalidState
     );
 
+    let num_pool_validators = get_validator_list_length(validator_list)?;
+
     if let Some(score) = state_account.state.compute_score(
         &clock,
         &epoch_schedule,

diff --git a/programs/steward/src/instructions/epoch_maintenance.rs b/programs/steward/src/instructions/epoch_maintenance.rs
@@ -5,8 +5,8 @@ use crate::{
         check_validator_list_has_stake_status_other_than, deserialize_stake_pool,
         get_stake_pool_address, get_validator_list, get_validator_list_length,
     },
-    Config, StewardStateAccount, CHECKED_VALIDATORS_REMOVED_FROM_LIST, COMPUTE_INSTANT_UNSTAKES,
-    EPOCH_MAINTENANCE, POST_LOOP_IDLE, PRE_LOOP_IDLE, REBALANCE, RESET_TO_IDLE,
+    Config, StewardStateAccount, COMPUTE_INSTANT_UNSTAKES, EPOCH_MAINTENANCE, POST_LOOP_IDLE,
+    PRE_LOOP_IDLE, REBALANCE, RESET_TO_IDLE,
 };
 use anchor_lang::prelude::*;
 use spl_stake_pool::state::StakeStatus;
@@ -97,13 +97,9 @@ pub fn handler(
             state_account.state.current_epoch = clock.epoch;
 
             // We keep Compute Scores and Compute Delegations to be unset on next epoch cycle
-            state_account.state.unset_flag(
-                CHECKED_VALIDATORS_REMOVED_FROM_LIST
-                    | PRE_LOOP_IDLE
-                    | COMPUTE_INSTANT_UNSTAKES
-                    | REBALANCE
-                    | POST_LOOP_IDLE,
-            );
+            state_account
+                .state
+                .unset_flag(PRE_LOOP_IDLE | COMPUTE_INSTANT_UNSTAKES | REBALANCE | POST_LOOP_IDLE);
             state_account
                 .state
                 .set_flag(RESET_TO_IDLE | EPOCH_MAINTENANCE);

diff --git a/programs/steward/src/instructions/idle.rs b/programs/steward/src/instructions/idle.rs
@@ -1,7 +1,9 @@
 use anchor_lang::prelude::*;
 
 use crate::{
-    errors::StewardError, maybe_transition_and_emit, Config, StewardStateAccount, StewardStateEnum,
+    maybe_transition_and_emit,
+    utils::{get_validator_list, state_checks},
+    Config, StewardStateAccount, StewardStateEnum,
 };
 
 #[derive(Accounts)]
@@ -14,6 +16,10 @@ pub struct Idle<'info> {
         bump
     )]
     pub state_account: AccountLoader<'info, StewardStateAccount>,
+
+    /// CHECK: account type checked in state_checks and address set in config
+    #[account(address = get_validator_list(&config)?)]
+    pub validator_list: AccountInfo<'info>,
 }
 
 /*
@@ -25,26 +31,13 @@ pub fn handler(ctx: Context<Idle>) -> Result<()> {
     let clock = Clock::get()?;
     let epoch_schedule = EpochSchedule::get()?;
 
-    {
-        if config.is_paused() {
-            return Err(StewardError::StateMachinePaused.into());
-        }
-
-        require!(
-            matches!(state_account.state.state_tag, StewardStateEnum::Idle),
-            StewardError::InvalidState
-        );
-
-        require!(
-            clock.epoch == state_account.state.current_epoch,
-            StewardError::EpochMaintenanceNotComplete
-        );
-
-        require!(
-            state_account.state.validators_for_immediate_removal.count() == 0,
-            StewardError::ValidatorsNeedToBeRemoved
-        );
-    }
+    state_checks(
+        &clock,
+        &config,
+        &state_account,
+        &ctx.accounts.validator_list,
+        Some(StewardStateEnum::Idle),
+    )?;
 
     maybe_transition_and_emit(
         &mut state_account.state,

diff --git a/programs/steward/src/instructions/rebalance.rs b/programs/steward/src/instructions/rebalance.rs
@@ -23,7 +23,10 @@ use crate::{
     errors::StewardError,
     events::{DecreaseComponents, RebalanceEvent, RebalanceTypeTag},
     maybe_transition_and_emit,
-    utils::{deserialize_stake_pool, get_stake_pool_address, get_validator_stake_info_at_index},
+    utils::{
+        deserialize_stake_pool, get_stake_pool_address, get_validator_stake_info_at_index,
+        state_checks,
+    },
     Config, StewardStateAccount, StewardStateEnum,
 };
 
@@ -146,26 +149,13 @@ pub fn handler(ctx: Context<Rebalance>, validator_list_index: usize) -> Result<(
     let clock = Clock::get()?;
     let epoch_schedule = EpochSchedule::get()?;
 
-    {
-        if config.is_paused() {
-            return Err(StewardError::StateMachinePaused.into());
-        }
-
-        require!(
-            matches!(state_account.state.state_tag, StewardStateEnum::Rebalance),
-            StewardError::InvalidState
-        );
-
-        require!(
-            clock.epoch == state_account.state.current_epoch,
-            StewardError::EpochMaintenanceNotComplete
-        );
-
-        require!(
-            state_account.state.validators_for_immediate_removal.count() == 0,
-            StewardError::ValidatorsNeedToBeRemoved
-        );
-    }
+    state_checks(
+        &clock,
+        &config,
+        &state_account,
+        &ctx.accounts.validator_list,
+        Some(StewardStateEnum::Rebalance),
+    )?;
 
     let validator_stake_info =
         get_validator_stake_info_at_index(validator_list, validator_list_index)?;

diff --git a/programs/steward/src/state/steward_state.rs b/programs/steward/src/state/steward_state.rs
@@ -127,7 +127,7 @@ pub struct StewardState {
 
 pub const STATE_PADDING_0_SIZE: usize = MAX_VALIDATORS * 8 + 2;
 
-#[derive(Clone, Copy)]
+#[derive(Clone, Copy, PartialEq)]
 #[repr(u64)]
 pub enum StewardStateEnum {
     /// Start state
@@ -245,15 +245,10 @@ pub const REBALANCE: u32 = 1 << 5;
 pub const POST_LOOP_IDLE: u32 = 1 << 6;
 // BITS 8-15 RESERVED FOR FUTURE USE
 // BITS 16-23 OPERATIONAL FLAGS
-/// In epoch maintenance, we only need to check the validator pool
-/// once for any validators that still need to be removed
-/// when there are no validators to remove from the pool, the operation continues
-/// and this condition is not checked again
-pub const CHECKED_VALIDATORS_REMOVED_FROM_LIST: u32 = 1 << 16;
 /// In epoch maintenance, when a new epoch is detected, we need a flag to tell the
 /// state transition layer that it needs to be reset to the IDLE state
 /// this flag is set in in epoch_maintenance and unset in the IDLE state transition
-pub const RESET_TO_IDLE: u32 = 1 << 17;
+pub const RESET_TO_IDLE: u32 = 1 << 16;
 // BITS 24-31 RESERVED FOR FUTURE USE
 
 impl StewardState {

diff --git a/programs/steward/src/utils.rs b/programs/steward/src/utils.rs
@@ -8,7 +8,47 @@ use spl_stake_pool::{
     state::{StakeStatus, ValidatorListHeader, ValidatorStakeInfo},
 };
 
-use crate::{errors::StewardError, Config, Delegation};
+use crate::{errors::StewardError, Config, Delegation, StewardStateAccount, StewardStateEnum};
+
+pub fn state_checks(
+    clock: &Clock,
+    config: &Config,
+    state_account: &StewardStateAccount,
+    validator_list_account_info: &AccountInfo,
+    expected_state: Option<StewardStateEnum>,
+) -> Result<()> {
+    if config.is_paused() {
+        return Err(StewardError::StateMachinePaused.into());
+    }
+
+    if let Some(expected_state) = expected_state {
+        require!(
+            state_account.state.state_tag == expected_state,
+            StewardError::InvalidState
+        );
+    }
+
+    require!(
+        clock.epoch == state_account.state.current_epoch,
+        StewardError::EpochMaintenanceNotComplete
+    );
+
+    require!(
+        state_account.state.validators_for_immediate_removal.count() == 0,
+        StewardError::ValidatorsNeedToBeRemoved
+    );
+
+    // Ensure we have a 1-1 mapping between the number of validators
+    let validators_in_list = get_validator_list_length(validator_list_account_info)?;
+    require!(
+        state_account.state.num_pool_validators as usize
+            + state_account.state.validators_added as usize
+            == validators_in_list,
+        StewardError::ListStateMismatch
+    );
+
+    Ok(())
+}
 
 pub fn get_stake_pool_address(account: &AccountLoader<Config>) -> Result<Pubkey> {
     let config = account.load()?;