diff --git a/program/src/processor.rs b/program/src/processor.rs index d64371ea..c533e1f6 100644 --- a/program/src/processor.rs +++ b/program/src/processor.rs @@ -580,9 +580,9 @@ impl Processor { _ => return Err(ProgramError::InvalidAccountData), } - // Deinitialize state upon zero balance + // Truncate state upon zero balance if split_lamports == source_lamport_balance { - set_stake_state(source_stake_account_info, &StakeStateV2::Uninitialized)?; + source_stake_account_info.realloc(0, false)?; } relocate_lamports( @@ -615,8 +615,8 @@ impl Processor { let (signers, custodian) = collect_signers_checked(Some(withdraw_authority_info), option_lockup_authority_info)?; - let (lockup, reserve, is_staked) = match get_stake_state(source_stake_account_info)? { - StakeStateV2::Stake(meta, stake, _stake_flag) => { + let (lockup, reserve, is_staked) = match get_stake_state(source_stake_account_info) { + Ok(StakeStateV2::Stake(meta, stake, _stake_flag)) => { meta.authorized .check(&signers, StakeAuthorize::Withdrawer) .map_err(to_program_error)?; @@ -637,20 +637,30 @@ impl Processor { let staked_and_reserve = checked_add(staked, meta.rent_exempt_reserve)?; (meta.lockup, staked_and_reserve, staked != 0) } - StakeStateV2::Initialized(meta) => { + Ok(StakeStateV2::Initialized(meta)) => { meta.authorized .check(&signers, StakeAuthorize::Withdrawer) .map_err(to_program_error)?; // stake accounts must have a balance >= rent_exempt_reserve (meta.lockup, meta.rent_exempt_reserve, false) } - StakeStateV2::Uninitialized => { + Ok(StakeStateV2::Uninitialized) => { if !signers.contains(source_stake_account_info.key) { return Err(ProgramError::MissingRequiredSignature); } (Lockup::default(), 0, false) // no lockup, no restrictions } - _ => return Err(ProgramError::InvalidAccountData), + Err(e) + if e == ProgramError::InvalidAccountData + && source_stake_account_info.data_len() == 0 => + { + if !signers.contains(source_stake_account_info.key) { + return Err(ProgramError::MissingRequiredSignature); + } + (Lockup::default(), 0, false) // no lockup, no restrictions + } + Ok(StakeStateV2::RewardsPool) => return Err(ProgramError::InvalidAccountData), + Err(e) => return Err(e), }; // verify that lockup has expired or that the withdrawal is signed by the @@ -666,8 +676,8 @@ impl Processor { return Err(ProgramError::InsufficientFunds); } - // Deinitialize state upon zero balance - set_stake_state(source_stake_account_info, &StakeStateV2::Uninitialized)?; + // Truncate state upon zero balance + source_stake_account_info.realloc(0, false)?; } else { // a partial withdrawal must not deplete the reserve let withdraw_lamports_and_reserve = checked_add(withdraw_lamports, reserve)?; @@ -783,8 +793,8 @@ impl Processor { set_stake_state(destination_stake_account_info, &merged_state)?; } - // Source is about to be drained, deinitialize its state - set_stake_state(source_stake_account_info, &StakeStateV2::Uninitialized)?; + // Source is about to be drained, truncate its state + source_stake_account_info.realloc(0, false)?; // Drain the source stake account relocate_lamports( diff --git a/program/tests/interface.rs b/program/tests/interface.rs index d3e36ba2..2a911521 100644 --- a/program/tests/interface.rs +++ b/program/tests/interface.rs @@ -952,3 +952,45 @@ fn test_epoch_rewards_period() { let instruction = instruction::get_minimum_delegation(); env.process_success(&instruction); } + +// other than Withdraw, instructions should fail with a zero-length stake account +// account data is truncated for instructions that intend to deallocate the account +// we want to ensure topping up such an account mid-transaction does not allow reuse +#[test] +fn test_no_use_dealloc() { + let mut env = Env::init(); + + for declaration in &*INSTRUCTION_DECLARATIONS { + let is_withdraw = matches!(declaration, StakeInterface::Withdraw { .. }); + let mut instruction = declaration.to_instruction(&mut env); + + for stake_address in [STAKE_ACCOUNT_BLACK, STAKE_ACCOUNT_WHITE] { + if instruction + .accounts + .iter() + .any(|account| account.pubkey == stake_address) + { + // replace stake account, if we use it, with its zero-length same-lamports equivalent + let lamports = env + .override_accounts + .get(&stake_address) + .unwrap_or_else(|| env.base_accounts.get(&stake_address).unwrap()) + .lamports(); + let stake_account = Account::create(lamports, vec![], id(), false, u64::MAX); + env.override_accounts.insert(stake_address, stake_account); + + if is_withdraw { + // truncating source account data makes this a self-signed withdraw + if instruction.accounts[0].pubkey == stake_address { + instruction.accounts[4].pubkey = stake_address; + } + env.process_success(&instruction); + } else { + env.process_fail(&instruction); + } + + env.reset(); + } + } + } +} diff --git a/program/tests/program_test.rs b/program/tests/program_test.rs index df570f1a..3b835184 100644 --- a/program/tests/program_test.rs +++ b/program/tests/program_test.rs @@ -10,6 +10,7 @@ use { pubkey::Pubkey, signature::{Keypair, Signer}, signers::Signers, + sysvar::rent::Rent, transaction::{Transaction, TransactionError}, vote::{ instruction as vote_instruction, @@ -280,13 +281,20 @@ pub async fn create_independent_stake_account_with_lockup( pub async fn create_blank_stake_account(context: &mut ProgramTestContext) -> Pubkey { let stake = Keypair::new(); - create_blank_stake_account_from_keypair(context, &stake).await + create_blank_stake_account_from_keypair(context, &stake, false).await +} + +pub async fn create_closed_stake_account(context: &mut ProgramTestContext) -> Pubkey { + let stake = Keypair::new(); + create_blank_stake_account_from_keypair(context, &stake, true).await } pub async fn create_blank_stake_account_from_keypair( context: &mut ProgramTestContext, stake: &Keypair, + is_closed: bool, ) -> Pubkey { + // lamports in a "closed" account is arbitrary, a real one via split/merge/withdraw would have 0 let lamports = get_stake_account_rent(&mut context.banks_client).await; let transaction = Transaction::new_signed_with_payer( @@ -294,7 +302,11 @@ pub async fn create_blank_stake_account_from_keypair( &context.payer.pubkey(), &stake.pubkey(), lamports, - StakeStateV2::size_of() as u64, + if is_closed { + 0 + } else { + StakeStateV2::size_of() as u64 + }, &id(), )], Some(&context.payer.pubkey()), @@ -594,7 +606,7 @@ async fn program_test_authorize() { let withdrawers: [_; 3] = std::array::from_fn(|_| Keypair::new()); let stake_keypair = Keypair::new(); - let stake = create_blank_stake_account_from_keypair(&mut context, &stake_keypair).await; + let stake = create_blank_stake_account_from_keypair(&mut context, &stake_keypair, false).await; // authorize uninitialized fails for (authority, authority_type) in [ @@ -880,6 +892,7 @@ pub enum StakeLifecycle { Active, Deactivating, Deactive, + Closed, } impl StakeLifecycle { // (stake, staker, withdrawer) @@ -918,16 +931,23 @@ impl StakeLifecycle { withdrawer_keypair: &Keypair, lockup: &Lockup, ) { - let authorized = Authorized { - staker: staker_keypair.pubkey(), - withdrawer: withdrawer_keypair.pubkey(), - }; + let is_closed = self == StakeLifecycle::Closed; - let stake = create_blank_stake_account_from_keypair(context, stake_keypair).await; + let stake = + create_blank_stake_account_from_keypair(context, stake_keypair, is_closed).await; if staked_amount > 0 { transfer(context, &stake, staked_amount).await; } + if is_closed { + return; + } + + let authorized = Authorized { + staker: staker_keypair.pubkey(), + withdrawer: withdrawer_keypair.pubkey(), + }; + if self >= StakeLifecycle::Initialized { let instruction = ixn::initialize(&stake, &authorized, lockup); process_instruction(context, &instruction, NO_SIGNERS) @@ -973,14 +993,14 @@ impl StakeLifecycle { pub fn split_minimum_enforced(&self) -> bool { match self { Self::Activating | Self::Active | Self::Deactivating | Self::Deactive => true, - Self::Uninitialized | Self::Initialized => false, + Self::Uninitialized | Self::Initialized | Self::Closed => false, } } pub fn withdraw_minimum_enforced(&self) -> bool { match self { Self::Activating | Self::Active | Self::Deactivating => true, - Self::Uninitialized | Self::Initialized | Self::Deactive => false, + Self::Uninitialized | Self::Initialized | Self::Deactive | Self::Closed => false, } } } @@ -1141,6 +1161,7 @@ async fn program_test_split(split_source_type: StakeLifecycle) { #[test_case(StakeLifecycle::Active; "active")] #[test_case(StakeLifecycle::Deactivating; "deactivating")] #[test_case(StakeLifecycle::Deactive; "deactive")] +#[test_case(StakeLifecycle::Closed; "closed")] #[tokio::test] async fn program_test_withdraw_stake(withdraw_source_type: StakeLifecycle) { let mut context = program_test().start_with_context().await; @@ -1167,16 +1188,22 @@ async fn program_test_withdraw_stake(withdraw_source_type: StakeLifecycle) { transfer(&mut context, &recipient, wallet_rent_exempt_reserve).await; let signers = match withdraw_source_type { - StakeLifecycle::Uninitialized => vec![&withdraw_source_keypair], + StakeLifecycle::Uninitialized | StakeLifecycle::Closed => vec![&withdraw_source_keypair], _ => vec![&withdrawer_keypair], }; // withdraw that would end rent-exemption always fails + let rent_spillover = if withdraw_source_type == StakeLifecycle::Closed { + stake_rent_exempt_reserve - Rent::default().minimum_balance(0) + 1 + } else { + 1 + }; + let instruction = ixn::withdraw( &withdraw_source, &signers[0].pubkey(), &recipient, - staked_amount + 1, + staked_amount + rent_spillover, None, ); let e = process_instruction(&mut context, &instruction, &signers) diff --git a/program/tests/stake_instruction.rs b/program/tests/stake_instruction.rs index d7f56cfa..e465f307 100644 --- a/program/tests/stake_instruction.rs +++ b/program/tests/stake_instruction.rs @@ -214,6 +214,10 @@ fn just_stake(meta: Meta, stake: u64) -> StakeStateV2 { ) } +fn is_closed(account: &AccountSharedData) -> bool { + account.lamports() == 0 && *account.owner() == id() && account.data().len() == 0 +} + fn get_active_stake_for_tests( stake_accounts: &[AccountSharedData], clock: &Clock, @@ -221,7 +225,7 @@ fn get_active_stake_for_tests( ) -> u64 { let mut active_stake = 0; for account in stake_accounts { - if let StakeStateV2::Stake(_meta, stake, _stake_flags) = account.state().unwrap() { + if let Ok(StakeStateV2::Stake(_meta, stake, _stake_flags)) = account.state() { let stake_status = stake.delegation.stake_activating_and_deactivating( clock.epoch, stake_history, @@ -376,7 +380,7 @@ fn test_stake_process_instruction() { 100, None, ), - Err(ProgramError::InvalidAccountData), + Err(ProgramError::MissingRequiredSignature), ); process_instruction_as_one_arg( &mollusk, @@ -1354,7 +1358,7 @@ fn test_authorize() { instruction_accounts.clone(), Ok(()), ); - assert_eq!(from(&accounts[0]).unwrap(), StakeStateV2::Uninitialized); + assert!(is_closed(&accounts[0])); // Test that withdrawal to account fails without authorized withdrawer instruction_accounts[4].is_signer = false; @@ -2488,7 +2492,7 @@ fn test_withdraw_stake() { Ok(()), ); assert_eq!(accounts[0].lamports(), 0); - assert_eq!(from(&accounts[0]).unwrap(), StakeStateV2::Uninitialized); + assert!(is_closed(&accounts[0])); // initialize stake let lockup = Lockup { @@ -2633,7 +2637,7 @@ fn test_withdraw_stake() { Ok(()), ); assert_eq!(accounts[0].lamports(), 0); - assert_eq!(from(&accounts[0]).unwrap(), StakeStateV2::Uninitialized); + assert!(is_closed(&accounts[0])); // overflow let rent = Rent::default(); @@ -2904,7 +2908,7 @@ fn test_withdraw_lockup() { instruction_accounts.clone(), Ok(()), ); - assert_eq!(from(&accounts[0]).unwrap(), StakeStateV2::Uninitialized); + assert!(is_closed(&accounts[0])); // should pass, custodian is the same as the withdraw authority instruction_accounts[5].pubkey = stake_address; @@ -2924,7 +2928,7 @@ fn test_withdraw_lockup() { instruction_accounts.clone(), Ok(()), ); - assert_eq!(from(&accounts[0]).unwrap(), StakeStateV2::Uninitialized); + assert!(is_closed(&accounts[0])); transaction_accounts[0] = (stake_address, stake_account); // should pass, lockup has expired @@ -2938,7 +2942,7 @@ fn test_withdraw_lockup() { instruction_accounts, Ok(()), ); - assert_eq!(from(&accounts[0]).unwrap(), StakeStateV2::Uninitialized); + assert!(is_closed(&accounts[0])); } #[test] @@ -5260,7 +5264,7 @@ fn test_split_100_percent_of_source() { match state { StakeStateV2::Initialized(_) => { assert_eq!(Ok(*state), accounts[1].state()); - assert_eq!(Ok(StakeStateV2::Uninitialized), accounts[0].state()); + assert!(is_closed(&accounts[0])); } StakeStateV2::Stake(meta, stake, stake_flags) => { assert_eq!( @@ -5277,7 +5281,7 @@ fn test_split_100_percent_of_source() { )), accounts[1].state() ); - assert_eq!(Ok(StakeStateV2::Uninitialized), accounts[0].state()); + assert!(is_closed(&accounts[0])); } _ => unreachable!(), } @@ -5399,7 +5403,7 @@ fn test_split_100_percent_of_source_to_account_with_lamports() { )), accounts[1].state() ); - assert_eq!(Ok(StakeStateV2::Uninitialized), accounts[0].state()); + assert!(is_closed(&accounts[0])); } } } @@ -5545,7 +5549,7 @@ fn test_split_rent_exemptness() { Ok(StakeStateV2::Initialized(expected_split_meta)), accounts[1].state() ); - assert_eq!(Ok(StakeStateV2::Uninitialized), accounts[0].state()); + assert!(is_closed(&accounts[0])); } StakeStateV2::Stake(_meta, stake, stake_flags) => { // Expected stake should reflect original stake amount so that extra lamports @@ -5570,7 +5574,7 @@ fn test_split_rent_exemptness() { accounts[1].lamports(), expected_stake + source_rent_exempt_reserve, ); - assert_eq!(Ok(StakeStateV2::Uninitialized), accounts[0].state()); + assert!(is_closed(&accounts[0])); } _ => unreachable!(), } @@ -5719,7 +5723,7 @@ fn test_split_require_rent_exempt_destination() { if let StakeStateV2::Stake(meta, stake, stake_flags) = state { // split entire source account, including rent-exempt reserve if accounts[0].lamports() == 0 { - assert_eq!(Ok(StakeStateV2::Uninitialized), accounts[0].state()); + assert!(is_closed(&accounts[0])); assert_eq!( Ok(StakeStateV2::Stake( *meta, @@ -5903,7 +5907,7 @@ fn test_merge() { } _ => unreachable!(), } - assert_eq!(accounts[1].state(), Ok(StakeStateV2::Uninitialized)); + assert!(is_closed(&accounts[1])); } } } @@ -6365,10 +6369,7 @@ fn test_merge_active_stake() { expected_result.clone(), ); if expected_result.is_ok() { - assert_eq!( - accounts[1 - iteration].state(), - Ok(StakeStateV2::Uninitialized) - ); + assert!(is_closed(&accounts[1 - iteration])); } } }