Audit fixes: harden pool init against PDA pre-funding griefing

idl.json

@@ -67,7 +67,7 @@
       "docs": [
         "Stake tokens into the pool. Creates user stake account if needed.",
         "Preserves maturity percentage when adding to existing stake.",
-        "Optional trailing account: PoolMetadata PDA to increment member_count on new stake."
+        "Requires the PoolMetadata PDA (increments member_count on a new stake; an uninitialized account is tolerated for pools without metadata)."
       ],
       "accounts": [
         {
@@ -108,6 +108,11 @@
         {
           "name": "tokenProgram",
           "docs": ["Token program (SPL Token or Token 2022)"]
+        },
+        {
+          "name": "metadata",
+          "writable": true,
+          "docs": ["Metadata PDA (required; member_count is incremented on a new stake; an uninitialized account is tolerated for pools without metadata)"]
         }
       ],
       "args": [
@@ -163,15 +168,13 @@
         },
         {
           "name": "systemProgram",
-          "optional": true,
           "address": "11111111111111111111111111111111",
-          "docs": ["System program (optional, for legacy account reallocation)"]
+          "docs": ["System program (required; for legacy account reallocation)"]
         },
         {
           "name": "metadata",
           "writable": true,
-          "optional": true,
-          "docs": ["Metadata PDA (optional, to decrement member_count on full-unstake close)"]
+          "docs": ["Metadata PDA (required; member_count is decremented on a full-unstake close; an uninitialized account is tolerated for pools without metadata)"]
         }
       ],
       "args": [
@@ -426,15 +429,13 @@
         },
         {
           "name": "systemProgram",
-          "optional": true,
           "address": "11111111111111111111111111111111",
-          "docs": ["System program (optional)"]
+          "docs": ["System program (required)"]
         },
         {
           "name": "metadata",
           "writable": true,
-          "optional": true,
-          "docs": ["Metadata PDA (optional, to decrement member_count on close)"]
+          "docs": ["Metadata PDA (required; member_count is decremented on close; an uninitialized account is tolerated for pools without metadata)"]
         }
       ],
       "args": []
@@ -494,6 +495,11 @@
           "writable": true,
           "signer": true,
           "docs": ["User / owner (receives rent refund)"]
+        },
+        {
+          "name": "metadata",
+          "writable": true,
+          "docs": ["Metadata PDA (required; member_count is decremented on close; an uninitialized account is tolerated for pools without metadata)"]
         }
       ],
       "args": []
@@ -651,7 +657,7 @@
         "Stake tokens on behalf of another user (beneficiary).",
         "Staker provides tokens and pays rent; beneficiary receives the staking position.",
         "Beneficiary does NOT need to sign.",
-        "Optional trailing account: PoolMetadata PDA to increment member_count on new stake."
+        "Requires the PoolMetadata PDA (increments member_count on a new stake; an uninitialized account is tolerated for pools without metadata)."
       ],
       "accounts": [
         {
@@ -697,6 +703,11 @@
         {
           "name": "tokenProgram",
           "docs": ["Token program (SPL Token or Token 2022)"]
+        },
+        {
+          "name": "metadata",
+          "writable": true,
+          "docs": ["Metadata PDA (required; member_count is incremented on a new stake; an uninitialized account is tolerated for pools without metadata)"]
         }
       ],
       "args": [
@@ -1032,7 +1043,7 @@
           {
             "name": "memberCount",
             "type": "u64",
-            "docs": ["Active staker count"]
+            "docs": ["Active staker count. The metadata PDA is a required account on stake/unstake/close, so this stays exact (incremented on a new stake, decremented on close)."]
           },
           {
             "name": "bump",

programs/chiefstaker/src/instructions/close_stake.rs

@@ -1,6 +1,6 @@
 //! Close an empty user stake account to reclaim rent
 
-use borsh::{BorshDeserialize, BorshSerialize};
+use borsh::BorshDeserialize;
 use solana_program::{
     account_info::{next_account_info, AccountInfo},
     entrypoint::ProgramResult,
@@ -11,7 +11,7 @@ use solana_program::{
 use crate::{
     error::StakingError,
     math::WAD,
-    state::{PoolMetadata, StakingPool, UserStake},
+    state::{update_pool_member_count, StakingPool, UserStake},
 };
 
 /// Close a zero-balance user stake account, returning rent to the user.
@@ -29,6 +29,9 @@ pub fn process_close_stake_account(
     let pool_info = next_account_info(account_info_iter)?;
     let user_stake_info = next_account_info(account_info_iter)?;
     let user_info = next_account_info(account_info_iter)?;
+    // Metadata PDA (required). Tolerates an uninitialized account for pools that
+    // never created metadata (see update_pool_member_count).
+    let metadata_info = next_account_info(account_info_iter)?;
 
     // Validate user is signer
     if !user_info.is_signer {
@@ -95,22 +98,8 @@ pub fn process_close_stake_account(
     let mut stake_data = user_stake_info.try_borrow_mut_data()?;
     stake_data.fill(0);
 
-    // Optional metadata account: decrement member_count on close
-    if let Some(metadata_info) = account_info_iter.next() {
-        if metadata_info.owner == program_id && !metadata_info.data_is_empty() {
-            let (expected_metadata, _) =
-                PoolMetadata::derive_pda(pool_info.key, program_id);
-            if *metadata_info.key == expected_metadata {
-                let mut metadata =
-                    PoolMetadata::try_from_slice(&metadata_info.try_borrow_data()?)?;
-                if metadata.is_initialized() && metadata.pool == *pool_info.key {
-                    metadata.member_count = metadata.member_count.saturating_sub(1);
-                    let mut metadata_data = metadata_info.try_borrow_mut_data()?;
-                    metadata.serialize(&mut &mut metadata_data[..])?;
-                }
-            }
-        }
-    }
+    // Decrement member_count on close (metadata PDA is a required account).
+    update_pool_member_count(program_id, pool_info, metadata_info, -1)?;
 
     msg!("Closed user stake account, returned {} lamports", stake_lamports);
 

programs/chiefstaker/src/instructions/complete_unstake.rs

@@ -138,9 +138,10 @@ pub fn process_complete_unstake(
 
     let withdraw_amount = user_stake.unstake_request_amount;
 
-    // Optional trailing accounts: system program (legacy realloc) then metadata
-    let system_program_info = account_info_iter.next();
-    let metadata_info = account_info_iter.next();
+    // System program (for legacy account realloc) and the metadata PDA (for the
+    // member_count decrement on close) are both required accounts.
+    let system_program_info = next_account_info(account_info_iter)?;
+    let metadata_info = next_account_info(account_info_iter)?;
 
     if user_stake.unstake_request_settled == 0 {
         // ── Legacy in-flight request ──────────────────────────────────────
@@ -161,8 +162,8 @@ pub fn process_complete_unstake(
             user_info,
             withdraw_amount,
             current_time,
-            system_program_info,
-            metadata_info,
+            Some(system_program_info),
+            Some(metadata_info),
         );
     }
 
@@ -215,7 +216,7 @@ pub fn process_complete_unstake(
 
     // A full unstake fully resets the position; close the account to reclaim rent.
     if should_close {
-        close_user_stake_account(program_id, pool_info, user_stake_info, user_info, metadata_info)?;
+        close_user_stake_account(program_id, pool_info, user_stake_info, user_info, Some(metadata_info))?;
     }
 
     Ok(())

programs/chiefstaker/src/instructions/initialize.rs

@@ -6,7 +6,7 @@ use solana_program::{
     clock::Clock,
     entrypoint::ProgramResult,
     msg,
-    program::invoke_signed,
+    program::{invoke, invoke_signed},
     pubkey::Pubkey,
     rent::Rent,
     system_instruction,
@@ -288,30 +288,34 @@ pub fn process_initialize_pool(
         return Err(StakingError::InvalidPDA.into());
     }
 
+    // Reject re-initialization. The pool/vault PDAs are deterministic and could
+    // have been pre-funded with lamports by a third party (a griefing vector that
+    // would make plain `create_account` fail), but a third party cannot allocate
+    // or assign them. So an empty (zero-length) data buffer means "not yet
+    // initialized" regardless of any stray lamports; non-empty means already done.
+    if !pool_info.data_is_empty() {
+        return Err(StakingError::AlreadyInitialized.into());
+    }
+    if !token_vault_info.data_is_empty() {
+        return Err(StakingError::AlreadyInitialized.into());
+    }
+
     let rent = Rent::from_account_info(rent_sysvar_info)?;
     let clock = Clock::get()?;
 
-    // Create pool account
+    // Create pool account (tolerates a pre-funded PDA)
     let pool_seeds = &[POOL_SEED, mint_info.key.as_ref(), &[pool_bump]];
-    let pool_rent = rent.minimum_balance(StakingPool::LEN);
-
-    invoke_signed(
-        &system_instruction::create_account(
-            authority_info.key,
-            pool_info.key,
-            pool_rent,
-            StakingPool::LEN as u64,
-            program_id,
-        ),
-        &[
-            authority_info.clone(),
-            pool_info.clone(),
-            system_program_info.clone(),
-        ],
-        &[pool_seeds],
+    create_pda_account(
+        authority_info,
+        pool_info,
+        system_program_info,
+        program_id,
+        StakingPool::LEN,
+        &rent,
+        pool_seeds,
     )?;
 
-    // Create token vault account
+    // Create token vault account (tolerates a pre-funded PDA)
     let vault_seeds = &[TOKEN_VAULT_SEED, pool_info.key.as_ref(), &[vault_bump]];
 
     // Get the size needed for a token account
@@ -322,22 +326,15 @@ pub fn process_initialize_pool(
     } else {
         spl_token_2022::state::Account::LEN
     };
-    let vault_rent = rent.minimum_balance(vault_size);
 
-    invoke_signed(
-        &system_instruction::create_account(
-            authority_info.key,
-            token_vault_info.key,
-            vault_rent,
-            vault_size as u64,
-            token_program_info.key,
-        ),
-        &[
-            authority_info.clone(),
-            token_vault_info.clone(),
-            system_program_info.clone(),
-        ],
-        &[vault_seeds],
+    create_pda_account(
+        authority_info,
+        token_vault_info,
+        system_program_info,
+        token_program_info.key,
+        vault_size,
+        &rent,
+        vault_seeds,
     )?;
 
     // Initialize token vault as token account
@@ -372,3 +369,58 @@ pub fn process_initialize_pool(
 
     Ok(())
 }
+
+/// Create a program-derived account at `target`, tolerating the case where the
+/// target PDA has already been funded with lamports by a third party.
+///
+/// `system_instruction::create_account` requires the destination to hold zero
+/// lamports, so anyone can permanently block creation of a deterministic PDA by
+/// sending it 1 lamport first. To avoid that DoS we, when the account is already
+/// funded, top it up to rent-exemption and then `allocate` + `assign` it under
+/// its own seeds — operations that require the PDA's signature and therefore
+/// cannot be performed by a griefer who only transferred lamports.
+fn create_pda_account<'a>(
+    payer: &AccountInfo<'a>,
+    target: &AccountInfo<'a>,
+    system_program: &AccountInfo<'a>,
+    owner: &Pubkey,
+    space: usize,
+    rent: &Rent,
+    seeds: &[&[u8]],
+) -> ProgramResult {
+    let required = rent.minimum_balance(space);
+
+    if target.lamports() == 0 {
+        // Fast path: empty + unfunded, a single create_account does it all.
+        return invoke_signed(
+            &system_instruction::create_account(
+                payer.key,
+                target.key,
+                required,
+                space as u64,
+                owner,
+            ),
+            &[payer.clone(), target.clone(), system_program.clone()],
+            &[seeds],
+        );
+    }
+
+    // Pre-funded path: fund to rent-exemption, then allocate + assign under seeds.
+    let current = target.lamports();
+    if current < required {
+        invoke(
+            &system_instruction::transfer(payer.key, target.key, required - current),
+            &[payer.clone(), target.clone(), system_program.clone()],
+        )?;
+    }
+    invoke_signed(
+        &system_instruction::allocate(target.key, space as u64),
+        &[target.clone(), system_program.clone()],
+        &[seeds],
+    )?;
+    invoke_signed(
+        &system_instruction::assign(target.key, owner),
+        &[target.clone(), system_program.clone()],
+        &[seeds],
+    )
+}

programs/chiefstaker/src/instructions/stake.rs

@@ -17,7 +17,7 @@ use spl_token_2022::extension::StateWithExtensions;
 use crate::{
     error::StakingError,
     math::{exp_time_ratio, wad_mul, MAX_EXP_INPUT, U256, WAD},
-    state::{is_valid_token_program, PoolMetadata, StakingPool, UserStake, STAKE_SEED},
+    state::{is_valid_token_program, update_pool_member_count, StakingPool, UserStake, STAKE_SEED},
 };
 
 /// Stake tokens into the pool
@@ -50,6 +50,9 @@ pub fn process_stake(
     let user_info = next_account_info(account_info_iter)?;
     let system_program_info = next_account_info(account_info_iter)?;
     let token_program_info = next_account_info(account_info_iter)?;
+    // Metadata PDA (required). Carries member_count; an uninitialized account is
+    // tolerated for pools that never created metadata (see update_pool_member_count).
+    let metadata_info = next_account_info(account_info_iter)?;
 
     // Validate token program (SPL Token or Token 2022)
     if !is_valid_token_program(token_program_info.key) {
@@ -304,23 +307,9 @@ pub fn process_stake(
         ],
     )?;
 
-    // Optional metadata account: increment member_count on new stake
+    // Increment member_count on a new stake (metadata PDA is a required account).
     if is_new_stake {
-        if let Some(metadata_info) = account_info_iter.next() {
-            if metadata_info.owner == program_id && !metadata_info.data_is_empty() {
-                let (expected_metadata, _) =
-                    PoolMetadata::derive_pda(pool_info.key, program_id);
-                if *metadata_info.key == expected_metadata {
-                    let mut metadata =
-                        PoolMetadata::try_from_slice(&metadata_info.try_borrow_data()?)?;
-                    if metadata.is_initialized() && metadata.pool == *pool_info.key {
-                        metadata.member_count = metadata.member_count.saturating_add(1);
-                        let mut metadata_data = metadata_info.try_borrow_mut_data()?;
-                        metadata.serialize(&mut &mut metadata_data[..])?;
-                    }
-                }
-            }
-        }
+        update_pool_member_count(program_id, pool_info, metadata_info, 1)?;
     }
 
     msg!("Staked {} tokens", amount);