diff --git a/Cargo.lock b/Cargo.lock index d3ec163..aa7f5ee 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6684,6 +6684,7 @@ dependencies = [ "solana-sdk-ids", "solana-signature", "solana-signer", + "solana-stake-client", "solana-stake-interface 2.0.1", "solana-svm-log-collector", "solana-system-interface", diff --git a/program/Cargo.toml b/program/Cargo.toml index 97c5ae4..21d326e 100644 --- a/program/Cargo.toml +++ b/program/Cargo.toml @@ -43,6 +43,7 @@ solana-sdk-ids = "3.0.0" solana-signature = "3.0.0" solana-signer = "3.0.0" solana-svm-log-collector = "3.0.0" +solana-stake-client = { path = "../clients/rust" } solana-system-interface = { version = "2.0.0", features = ["bincode"] } solana-transaction = "3.0.0" test-case = "3.3.1" diff --git a/program/tests/helpers/context.rs b/program/tests/helpers/context.rs new file mode 100644 index 0000000..400a8e8 --- /dev/null +++ b/program/tests/helpers/context.rs @@ -0,0 +1,182 @@ +use { + super::{ + execution::ExecutionWithChecks, + lifecycle::StakeLifecycle, + utils::{add_sysvars, STAKE_RENT_EXEMPTION}, + }, + mollusk_svm::{result::Check, Mollusk}, + solana_account::AccountSharedData, + solana_instruction::Instruction, + solana_program_error::ProgramError, + solana_pubkey::Pubkey, + solana_stake_interface::state::Lockup, + solana_stake_program::id, +}; + +/// Builder for creating stake accounts with customizable parameters +pub struct StakeAccountBuilder<'a> { + ctx: &'a mut StakeTestContext, + lifecycle: StakeLifecycle, + staked_amount: u64, + stake_authority: Option, + withdraw_authority: Option, + lockup: Option, + vote_account: Option, + stake_pubkey: Option, +} + +impl StakeAccountBuilder<'_> { + /// Set the staked amount (lamports delegated to validator) + pub fn staked_amount(mut self, amount: u64) -> Self { + self.staked_amount = amount; + self + } + + /// Set a custom stake authority (defaults to ctx.staker) + pub fn stake_authority(mut self, authority: &Pubkey) -> Self { + self.stake_authority = Some(*authority); + self + } + + /// Set a custom withdraw authority (defaults to ctx.withdrawer) + pub fn withdraw_authority(mut self, authority: &Pubkey) -> Self { + self.withdraw_authority = Some(*authority); + self + } + + /// Set a custom lockup (defaults to Lockup::default()) + pub fn lockup(mut self, lockup: &Lockup) -> Self { + self.lockup = Some(*lockup); + self + } + + /// Set a custom vote account (defaults to ctx.vote_account) + pub fn vote_account(mut self, vote_account: &Pubkey) -> Self { + self.vote_account = Some(*vote_account); + self + } + + /// Set a specific stake account pubkey (defaults to Pubkey::new_unique()) + pub fn stake_pubkey(mut self, pubkey: &Pubkey) -> Self { + self.stake_pubkey = Some(*pubkey); + self + } + + /// Build the stake account and return (pubkey, account_data) + pub fn build(self) -> (Pubkey, AccountSharedData) { + let stake_pubkey = self.stake_pubkey.unwrap_or_else(Pubkey::new_unique); + let account = self.lifecycle.create_uninitialized_account(); + (stake_pubkey, account) + } +} + +pub struct StakeTestContext { + pub mollusk: Mollusk, + pub rent_exempt_reserve: u64, + pub staker: Pubkey, + pub withdrawer: Pubkey, +} + +impl StakeTestContext { + pub fn new() -> Self { + let mollusk = Mollusk::new(&id(), "solana_stake_program"); + Self { + mollusk, + rent_exempt_reserve: STAKE_RENT_EXEMPTION, + staker: Pubkey::new_unique(), + withdrawer: Pubkey::new_unique(), + } + } + + /// Create a stake account builder for the specified lifecycle stage + /// + /// Example: + /// ``` + /// let (stake, account) = ctx + /// .stake_account(StakeLifecycle::Active) + /// .staked_amount(1_000_000) + /// .build(); + /// ``` + pub fn stake_account(&mut self, lifecycle: StakeLifecycle) -> StakeAccountBuilder<'_> { + StakeAccountBuilder { + ctx: self, + lifecycle, + staked_amount: 0, + stake_authority: None, + withdraw_authority: None, + lockup: None, + vote_account: None, + stake_pubkey: None, + } + } + + /// Configure execution with specific checks, then call .execute(instruction, accounts) + /// + /// Usage: `ctx.checks(&checks).execute(instruction, accounts)` + pub fn checks<'a, 'b>(&'a mut self, checks: &'b [Check<'b>]) -> ExecutionWithChecks<'a, 'b> { + ExecutionWithChecks::new(self, checks) + } + + /// Execute an instruction with default success checks and missing signer testing + /// + /// Usage: `ctx.execute(instruction, accounts)` + pub fn execute( + &mut self, + instruction: Instruction, + accounts: &[(&Pubkey, &AccountSharedData)], + ) -> mollusk_svm::result::InstructionResult { + self.execute_internal(instruction, accounts, &[Check::success()], true) + } + + /// Internal: execute with given checks and current config + pub(crate) fn execute_internal( + &mut self, + instruction: Instruction, + accounts: &[(&Pubkey, &AccountSharedData)], + checks: &[Check], + test_missing_signers: bool, + ) -> mollusk_svm::result::InstructionResult { + let accounts_vec: Vec<(Pubkey, AccountSharedData)> = accounts + .iter() + .map(|(pk, data)| (**pk, (*data).clone())) + .collect(); + + if test_missing_signers { + verify_all_signers_required(&self.mollusk, &instruction, &accounts_vec); + } + + // Process with all signers present + let accounts_with_sysvars = add_sysvars(&self.mollusk, &instruction, accounts_vec); + self.mollusk + .process_and_validate_instruction(&instruction, &accounts_with_sysvars, checks) + } +} + +impl Default for StakeTestContext { + fn default() -> Self { + Self::new() + } +} + +/// Verify that removing any signer from the instruction causes MissingRequiredSignature error +fn verify_all_signers_required( + mollusk: &Mollusk, + instruction: &Instruction, + accounts: &[(Pubkey, AccountSharedData)], +) { + for i in 0..instruction.accounts.len() { + if instruction.accounts[i].is_signer { + let mut modified_instruction = instruction.clone(); + modified_instruction.accounts[i].is_signer = false; + + let accounts_with_sysvars = + add_sysvars(mollusk, &modified_instruction, accounts.to_vec()); + + mollusk.process_and_validate_instruction( + &modified_instruction, + &accounts_with_sysvars, + &[Check::err(ProgramError::MissingRequiredSignature)], + ); + } + } +} diff --git a/program/tests/helpers/execution.rs b/program/tests/helpers/execution.rs new file mode 100644 index 0000000..83e0781 --- /dev/null +++ b/program/tests/helpers/execution.rs @@ -0,0 +1,41 @@ +use { + super::context::StakeTestContext, mollusk_svm::result::Check, + solana_account::AccountSharedData, solana_instruction::Instruction, solana_pubkey::Pubkey, +}; + +/// Wrapper for executing with specific checks +/// +/// Usage: `ctx.checks(&checks).test_missing_signers(false).execute(instruction, accounts)` +pub struct ExecutionWithChecks<'a, 'b> { + pub(crate) ctx: &'a mut StakeTestContext, + pub(crate) checks: &'b [Check<'b>], + pub(crate) test_missing_signers: bool, +} + +impl<'a, 'b> ExecutionWithChecks<'a, 'b> { + pub fn new(ctx: &'a mut StakeTestContext, checks: &'b [Check<'b>]) -> Self { + Self { + ctx, + checks, + test_missing_signers: true, // default: test missing signers + } + } + + pub fn test_missing_signers(mut self, test: bool) -> Self { + self.test_missing_signers = test; + self + } + + pub fn execute( + self, + instruction: Instruction, + accounts: &[(&Pubkey, &AccountSharedData)], + ) -> mollusk_svm::result::InstructionResult { + self.ctx.execute_internal( + instruction, + accounts, + self.checks, + self.test_missing_signers, + ) + } +} diff --git a/program/tests/helpers/lifecycle.rs b/program/tests/helpers/lifecycle.rs new file mode 100644 index 0000000..74e5528 --- /dev/null +++ b/program/tests/helpers/lifecycle.rs @@ -0,0 +1,29 @@ +use { + super::utils::STAKE_RENT_EXEMPTION, solana_account::AccountSharedData, + solana_stake_interface::state::StakeStateV2, solana_stake_program::id, +}; + +/// Lifecycle states for stake accounts in tests +#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord)] +pub enum StakeLifecycle { + Uninitialized = 0, + Initialized, + Activating, + Active, + Deactivating, + Deactive, + Closed, +} + +impl StakeLifecycle { + /// Create an uninitialized stake account + pub fn create_uninitialized_account(self) -> AccountSharedData { + AccountSharedData::new_data_with_space( + STAKE_RENT_EXEMPTION, + &StakeStateV2::Uninitialized, + StakeStateV2::size_of(), + &id(), + ) + .unwrap() + } +} diff --git a/program/tests/helpers/mod.rs b/program/tests/helpers/mod.rs new file mode 100644 index 0000000..adbdb63 --- /dev/null +++ b/program/tests/helpers/mod.rs @@ -0,0 +1,7 @@ +#![allow(clippy::arithmetic_side_effects)] +#![allow(dead_code)] + +pub mod context; +pub mod execution; +pub mod lifecycle; +pub mod utils; diff --git a/program/tests/helpers/utils.rs b/program/tests/helpers/utils.rs new file mode 100644 index 0000000..3bc0659 --- /dev/null +++ b/program/tests/helpers/utils.rs @@ -0,0 +1,65 @@ +use { + mollusk_svm::Mollusk, + solana_account::{Account, AccountSharedData}, + solana_instruction::Instruction, + solana_pubkey::Pubkey, + solana_rent::Rent, + solana_stake_interface::{stake_history::StakeHistory, state::StakeStateV2}, + solana_sysvar_id::SysvarId, + std::collections::HashMap, +}; + +// hardcoded for convenience +pub const STAKE_RENT_EXEMPTION: u64 = 2_282_880; + +#[test] +fn assert_stake_rent_exemption() { + assert_eq!( + Rent::default().minimum_balance(StakeStateV2::size_of()), + STAKE_RENT_EXEMPTION + ); +} + +/// Resolve all accounts for an instruction, including sysvars and instruction accounts +/// +/// This function re-serializes the stake history sysvar from mollusk.sysvars.stake_history +/// every time it's called, ensuring that any updates to the stake history are reflected in the accounts. +pub fn add_sysvars( + mollusk: &Mollusk, + instruction: &Instruction, + accounts: Vec<(Pubkey, AccountSharedData)>, +) -> Vec<(Pubkey, Account)> { + // Build a map of provided accounts + let mut account_map: HashMap = accounts + .into_iter() + .map(|(pk, acc)| (pk, acc.into())) + .collect(); + + // Now resolve all accounts from the instruction + let mut result = Vec::new(); + for account_meta in &instruction.accounts { + let key = account_meta.pubkey; + let account = if let Some(acc) = account_map.remove(&key) { + // Use the provided account + acc + } else if Rent::check_id(&key) { + mollusk.sysvars.keyed_account_for_rent_sysvar().1 + } else if solana_clock::Clock::check_id(&key) { + mollusk.sysvars.keyed_account_for_clock_sysvar().1 + } else if solana_epoch_schedule::EpochSchedule::check_id(&key) { + mollusk.sysvars.keyed_account_for_epoch_schedule_sysvar().1 + } else if solana_epoch_rewards::EpochRewards::check_id(&key) { + mollusk.sysvars.keyed_account_for_epoch_rewards_sysvar().1 + } else if StakeHistory::check_id(&key) { + // Re-serialize stake history from mollusk.sysvars.stake_history + mollusk.sysvars.keyed_account_for_stake_history_sysvar().1 + } else { + // Default empty account + Account::default() + }; + + result.push((key, account)); + } + + result +} diff --git a/program/tests/initialize.rs b/program/tests/initialize.rs new file mode 100644 index 0000000..9920d80 --- /dev/null +++ b/program/tests/initialize.rs @@ -0,0 +1,220 @@ +#![allow(clippy::arithmetic_side_effects)] + +mod helpers; + +use { + helpers::{context::StakeTestContext, lifecycle::StakeLifecycle}, + mollusk_svm::result::Check, + solana_account::{AccountSharedData, ReadableAccount}, + solana_instruction::Instruction, + solana_program_error::ProgramError, + solana_pubkey::Pubkey, + solana_rent::Rent, + solana_sdk_ids::{stake::id, system_program::id as system_program_id}, + solana_stake_client::instructions::{InitializeBuilder, InitializeCheckedBuilder}, + solana_stake_interface::state::{Lockup, StakeStateV2}, + test_case::test_case, +}; + +#[derive(Debug, Clone, Copy)] +enum InitializeVariant { + Initialize, + InitializeChecked, +} + +fn lockup_for(variant: InitializeVariant, custodian: Pubkey) -> Lockup { + match variant { + InitializeVariant::Initialize => Lockup { + epoch: 1, + unix_timestamp: 0, + custodian, + }, + InitializeVariant::InitializeChecked => Lockup::default(), + } +} + +fn build_initialize_instruction( + variant: InitializeVariant, + stake: Pubkey, + staker: Pubkey, + withdrawer: Pubkey, + lockup: Lockup, +) -> Instruction { + match variant { + InitializeVariant::Initialize => InitializeBuilder::new() + .stake(stake) + .arg0(solana_stake_client::types::Authorized { staker, withdrawer }) + .arg1(solana_stake_client::types::Lockup { + unix_timestamp: lockup.unix_timestamp, + epoch: lockup.epoch, + custodian: lockup.custodian, + }) + .instruction(), + InitializeVariant::InitializeChecked => InitializeCheckedBuilder::new() + .stake(stake) + .stake_authority(staker) + .withdraw_authority(withdrawer) + .instruction(), + } +} + +#[test_case(InitializeVariant::Initialize; "initialize")] +#[test_case(InitializeVariant::InitializeChecked; "initialize_checked")] +fn test_initialize(variant: InitializeVariant) { + let mut ctx = StakeTestContext::default(); + + let staker = ctx.staker; + let withdrawer = ctx.withdrawer; + let lockup = lockup_for(variant, Pubkey::new_unique()); + + let (stake, stake_account) = ctx.stake_account(StakeLifecycle::Uninitialized).build(); + + let result = { + let program_id = id(); + let checks = [ + Check::success(), + Check::all_rent_exempt(), + Check::account(&stake) + .lamports(ctx.rent_exempt_reserve) + .owner(&program_id) + .space(StakeStateV2::size_of()) + .build(), + ]; + + ctx.checks(&checks).execute( + build_initialize_instruction(variant, stake, staker, withdrawer, lockup), + &[(&stake, &stake_account)], + ) + }; + + let resulting_account: AccountSharedData = result.resulting_accounts[0].1.clone().into(); + let stake_state: StakeStateV2 = bincode::deserialize(resulting_account.data()).unwrap(); + assert_eq!( + stake_state, + StakeStateV2::Initialized(solana_stake_interface::state::Meta { + authorized: solana_stake_interface::state::Authorized { staker, withdrawer }, + rent_exempt_reserve: ctx.rent_exempt_reserve, + lockup, + }), + ); + + // Re-initialize should fail + ctx.checks(&[Check::err(ProgramError::InvalidAccountData)]) + .test_missing_signers(false) + .execute( + build_initialize_instruction(variant, stake, staker, withdrawer, lockup), + &[(&stake, &resulting_account)], + ); +} + +#[test_case(InitializeVariant::Initialize; "initialize")] +#[test_case(InitializeVariant::InitializeChecked; "initialize_checked")] +fn test_initialize_insufficient_funds(variant: InitializeVariant) { + let mut ctx = StakeTestContext::default(); + + let staker = ctx.staker; + let withdrawer = ctx.withdrawer; + let lockup = lockup_for(variant, Pubkey::new_unique()); + + // Account has insufficient lamports + let stake = Pubkey::new_unique(); + let stake_account = AccountSharedData::new_data_with_space( + ctx.rent_exempt_reserve / 2, + &StakeStateV2::Uninitialized, + StakeStateV2::size_of(), + &id(), + ) + .unwrap(); + + ctx.checks(&[Check::err(ProgramError::InsufficientFunds)]) + .test_missing_signers(false) + .execute( + build_initialize_instruction(variant, stake, staker, withdrawer, lockup), + &[(&stake, &stake_account)], + ); +} + +#[test_case(InitializeVariant::Initialize; "initialize")] +#[test_case(InitializeVariant::InitializeChecked; "initialize_checked")] +fn test_initialize_incorrect_size_larger(variant: InitializeVariant) { + let mut ctx = StakeTestContext::default(); + + let rent_exempt_reserve = Rent::default().minimum_balance(StakeStateV2::size_of() * 2); + + let staker = ctx.staker; + let withdrawer = ctx.withdrawer; + let lockup = lockup_for(variant, Pubkey::new_unique()); + + // Account data length too large + let stake = Pubkey::new_unique(); + let stake_account = AccountSharedData::new_data_with_space( + rent_exempt_reserve, + &StakeStateV2::Uninitialized, + StakeStateV2::size_of() + 1, + &id(), + ) + .unwrap(); + + ctx.checks(&[Check::err(ProgramError::InvalidAccountData)]) + .test_missing_signers(false) + .execute( + build_initialize_instruction(variant, stake, staker, withdrawer, lockup), + &[(&stake, &stake_account)], + ); +} + +#[test_case(InitializeVariant::Initialize; "initialize")] +#[test_case(InitializeVariant::InitializeChecked; "initialize_checked")] +fn test_initialize_incorrect_size_smaller(variant: InitializeVariant) { + let mut ctx = StakeTestContext::default(); + + let rent_exempt_reserve = Rent::default().minimum_balance(StakeStateV2::size_of()); + + let staker = ctx.staker; + let withdrawer = ctx.withdrawer; + let lockup = lockup_for(variant, Pubkey::new_unique()); + + // Account data length too small + let stake = Pubkey::new_unique(); + let stake_account = AccountSharedData::new_data_with_space( + rent_exempt_reserve, + &StakeStateV2::Uninitialized, + StakeStateV2::size_of() - 1, + &id(), + ) + .unwrap(); + + ctx.checks(&[Check::err(ProgramError::InvalidAccountData)]) + .test_missing_signers(false) + .execute( + build_initialize_instruction(variant, stake, staker, withdrawer, lockup), + &[(&stake, &stake_account)], + ); +} + +#[test_case(InitializeVariant::Initialize; "initialize")] +#[test_case(InitializeVariant::InitializeChecked; "initialize_checked")] +fn test_initialize_wrong_owner(variant: InitializeVariant) { + let mut ctx = StakeTestContext::default(); + + let staker = ctx.staker; + let withdrawer = ctx.withdrawer; + let lockup = lockup_for(variant, Pubkey::new_unique()); + + // Owner is not the stake program + let stake = Pubkey::new_unique(); + let stake_account = AccountSharedData::new_data_with_space( + Rent::default().minimum_balance(StakeStateV2::size_of()), + &StakeStateV2::Uninitialized, + StakeStateV2::size_of(), + &system_program_id(), + ) + .unwrap(); + + ctx.checks(&[Check::err(ProgramError::InvalidAccountOwner)]) + .test_missing_signers(false) + .execute( + build_initialize_instruction(variant, stake, staker, withdrawer, lockup), + &[(&stake, &stake_account)], + ); +} diff --git a/program/tests/program_test.rs b/program/tests/program_test.rs index e1988c7..4f14680 100644 --- a/program/tests/program_test.rs +++ b/program/tests/program_test.rs @@ -396,17 +396,6 @@ async fn program_test_stake_checked_instructions() { let seed = "test seed"; let seeded_address = Pubkey::create_with_seed(&seed_base, seed, &system_program::id()).unwrap(); - // Test InitializeChecked with non-signing withdrawer - let stake = create_blank_stake_account(&mut context).await; - let instruction = ixn::initialize_checked(&stake, &Authorized { staker, withdrawer }); - - process_instruction_test_missing_signers( - &mut context, - &instruction, - &vec![&withdrawer_keypair], - ) - .await; - // Test AuthorizeChecked with non-signing staker let stake = create_independent_stake_account(&mut context, &Authorized { staker, withdrawer }, 0).await; @@ -482,116 +471,6 @@ async fn program_test_stake_checked_instructions() { .await; } -#[tokio::test] -async fn program_test_stake_initialize() { - let mut context = program_test().start_with_context().await; - let accounts = Accounts::default(); - accounts.initialize(&mut context).await; - - let rent_exempt_reserve = get_stake_account_rent(&mut context.banks_client).await; - - let staker_keypair = Keypair::new(); - let withdrawer_keypair = Keypair::new(); - let custodian_keypair = Keypair::new(); - - let staker = staker_keypair.pubkey(); - let withdrawer = withdrawer_keypair.pubkey(); - let custodian = custodian_keypair.pubkey(); - - let authorized = Authorized { staker, withdrawer }; - - let lockup = Lockup { - epoch: 1, - unix_timestamp: 0, - custodian, - }; - - let stake = create_blank_stake_account(&mut context).await; - let instruction = ixn::initialize(&stake, &authorized, &lockup); - - // should pass - process_instruction(&mut context, &instruction, NO_SIGNERS) - .await - .unwrap(); - - // check that we see what we expect - let account = get_account(&mut context.banks_client, &stake).await; - let stake_state: StakeStateV2 = bincode::deserialize(&account.data).unwrap(); - assert_eq!( - stake_state, - StakeStateV2::Initialized(Meta { - authorized, - rent_exempt_reserve, - lockup, - }), - ); - - // 2nd time fails, can't move it from anything other than uninit->init - refresh_blockhash(&mut context).await; - let e = process_instruction(&mut context, &instruction, NO_SIGNERS) - .await - .unwrap_err(); - assert_eq!(e, ProgramError::InvalidAccountData); - - // not enough balance for rent - let stake = Pubkey::new_unique(); - let account = SolanaAccount { - lamports: rent_exempt_reserve / 2, - data: vec![0; StakeStateV2::size_of()], - owner: id(), - executable: false, - rent_epoch: 1000, - }; - context.set_account(&stake, &account.into()); - - let instruction = ixn::initialize(&stake, &authorized, &lockup); - let e = process_instruction(&mut context, &instruction, NO_SIGNERS) - .await - .unwrap_err(); - assert_eq!(e, ProgramError::InsufficientFunds); - - // incorrect account sizes - let stake_keypair = Keypair::new(); - let stake = stake_keypair.pubkey(); - - let instruction = system_instruction::create_account( - &context.payer.pubkey(), - &stake, - rent_exempt_reserve * 2, - StakeStateV2::size_of() as u64 + 1, - &id(), - ); - process_instruction(&mut context, &instruction, &vec![&stake_keypair]) - .await - .unwrap(); - - let instruction = ixn::initialize(&stake, &authorized, &lockup); - let e = process_instruction(&mut context, &instruction, NO_SIGNERS) - .await - .unwrap_err(); - assert_eq!(e, ProgramError::InvalidAccountData); - - let stake_keypair = Keypair::new(); - let stake = stake_keypair.pubkey(); - - let instruction = system_instruction::create_account( - &context.payer.pubkey(), - &stake, - rent_exempt_reserve, - StakeStateV2::size_of() as u64 - 1, - &id(), - ); - process_instruction(&mut context, &instruction, &vec![&stake_keypair]) - .await - .unwrap(); - - let instruction = ixn::initialize(&stake, &authorized, &lockup); - let e = process_instruction(&mut context, &instruction, NO_SIGNERS) - .await - .unwrap_err(); - assert_eq!(e, ProgramError::InvalidAccountData); -} - #[tokio::test] async fn program_test_authorize() { let mut context = program_test().start_with_context().await;