program: reallocate stake account on zero lamports

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(

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();
+            }
+        }
+    }
+}

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,21 +281,32 @@ 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(
         &[system_instruction::create_account(
             &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)