stake-pool: Option to use transfer + allocate + assign for PDA (#2920)

Author: joncinque

stake-pool/program/src/processor.rs

@@ -160,6 +160,59 @@ fn create_transient_stake_account<'a>(
     )
 }
 
+/// Create an account on a program-derived address
+fn create_pda_account<'a>(
+    payer: &AccountInfo<'a>,
+    required_lamports: u64,
+    space: usize,
+    owner: &Pubkey,
+    system_program: &AccountInfo<'a>,
+    new_pda_account: &AccountInfo<'a>,
+    new_pda_signer_seeds: &[&[u8]],
+) -> ProgramResult {
+    if new_pda_account.lamports() > 0 {
+        let required_lamports = required_lamports.saturating_sub(new_pda_account.lamports());
+        if required_lamports > 0 {
+            invoke(
+                &system_instruction::transfer(payer.key, new_pda_account.key, required_lamports),
+                &[
+                    payer.clone(),
+                    new_pda_account.clone(),
+                    system_program.clone(),
+                ],
+            )?;
+        }
+
+        invoke_signed(
+            &system_instruction::allocate(new_pda_account.key, space as u64),
+            &[new_pda_account.clone(), system_program.clone()],
+            &[new_pda_signer_seeds],
+        )?;
+
+        invoke_signed(
+            &system_instruction::assign(new_pda_account.key, owner),
+            &[new_pda_account.clone(), system_program.clone()],
+            &[new_pda_signer_seeds],
+        )
+    } else {
+        invoke_signed(
+            &system_instruction::create_account(
+                payer.key,
+                new_pda_account.key,
+                required_lamports,
+                space as u64,
+                owner,
+            ),
+            &[
+                payer.clone(),
+                new_pda_account.clone(),
+                system_program.clone(),
+            ],
+            &[new_pda_signer_seeds],
+        )
+    }
+}
+
 /// Program state handler.
 pub struct Processor {}
 impl Processor {
@@ -798,22 +851,19 @@ impl Processor {
         ];
 
         // Fund the stake account with the minimum + rent-exempt balance
-        let required_lamports = MINIMUM_ACTIVE_STAKE
-            + rent.minimum_balance(std::mem::size_of::<stake::state::StakeState>());
+        let space = std::mem::size_of::<stake::state::StakeState>();
+        let required_lamports = MINIMUM_ACTIVE_STAKE + rent.minimum_balance(space);
 
         // Create new stake account
-        invoke_signed(
-            &system_instruction::create_account(
-                funder_info.key,
-                stake_info.key,
-                required_lamports,
-                std::mem::size_of::<stake::state::StakeState>() as u64,
-                &stake::program::id(),
-            ),
-            &[funder_info.clone(), stake_info.clone()],
-            &[stake_account_signer_seeds],
+        create_pda_account(
+            funder_info,
+            required_lamports,
+            space,
+            &stake::program::id(),
+            system_program_info,
+            stake_info,
+            stake_account_signer_seeds,
         )?;
-
         invoke(
             &stake::instruction::initialize(
                 stake_info.key,

stake-pool/program/tests/vsa_add.rs

@@ -500,3 +500,46 @@ async fn fail_on_incorrectly_derived_stake_account() {
         )
     );
 }
+
+#[tokio::test]
+async fn success_with_lamports_in_account() {
+    let (mut banks_client, payer, recent_blockhash, stake_pool_accounts, validator_stake) =
+        setup().await;
+
+    transfer(
+        &mut banks_client,
+        &payer,
+        &recent_blockhash,
+        &validator_stake.stake_account,
+        1_000_000,
+    )
+    .await;
+
+    let error = stake_pool_accounts
+        .add_validator_to_pool(
+            &mut banks_client,
+            &payer,
+            &recent_blockhash,
+            &validator_stake.stake_account,
+            &validator_stake.vote.pubkey(),
+        )
+        .await;
+    assert!(error.is_none());
+
+    // Check stake account existence and authority
+    let stake = get_account(&mut banks_client, &validator_stake.stake_account).await;
+    let stake_state = deserialize::<stake::state::StakeState>(&stake.data).unwrap();
+    match stake_state {
+        stake::state::StakeState::Stake(meta, _) => {
+            assert_eq!(
+                &meta.authorized.staker,
+                &stake_pool_accounts.withdraw_authority
+            );
+            assert_eq!(
+                &meta.authorized.withdrawer,
+                &stake_pool_accounts.withdraw_authority
+            );
+        }
+        _ => panic!(),
+    }
+}