From 9342d4a60c944b70c4bb905387d2d143b3b965dd Mon Sep 17 00:00:00 2001 From: Mark Karpeles Date: Sun, 24 May 2026 03:49:30 +0900 Subject: [PATCH 1/3] Audit fixes: harden pool init against PDA pre-funding griefing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses the medium + low findings from the post-merge security audit. Medium — init griefing: the pool (["pool", mint]) and vault (["token_vault", pool]) PDAs are deterministic, so anyone could send 1 lamport to them before the real authority initializes; system create_account fails on a non-zero balance, permanently bricking pool creation for that mint. initialize.rs now uses a create_pda_account helper that tolerates a pre-funded PDA (fund-to-rent + allocate + assign under the PDA seeds — operations a griefer cannot perform) and adds explicit re-init guards. Adds an init_griefing ProgramTest that pre-funds both PDAs and asserts initialization still succeeds. Low — total_residual_unpaid: documented as legacy-only (the always-close change means current code never increments it; kept for layout compat + draining pre-upgrade residual balances). Low — member_count: documented as best-effort (accurate only when the optional metadata account is supplied on stake AND every close path; display-only, never affects funds). Co-Authored-By: Claude Opus 4.7 (1M context) --- .../src/instructions/initialize.rs | 120 ++++++++++++----- programs/chiefstaker/src/state.rs | 15 ++- programs/chiefstaker/tests/init_griefing.rs | 124 ++++++++++++++++++ 3 files changed, 224 insertions(+), 35 deletions(-) create mode 100644 programs/chiefstaker/tests/init_griefing.rs diff --git a/programs/chiefstaker/src/instructions/initialize.rs b/programs/chiefstaker/src/instructions/initialize.rs index b95c75a..f0f5171 100644 --- a/programs/chiefstaker/src/instructions/initialize.rs +++ b/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], + ) +} diff --git a/programs/chiefstaker/src/state.rs b/programs/chiefstaker/src/state.rs index 67d7fbe..2854742 100644 --- a/programs/chiefstaker/src/state.rs +++ b/programs/chiefstaker/src/state.rs @@ -150,6 +150,14 @@ pub struct StakingPool { /// amount=0 (no allocation in `total_staked * acc_rps`), so including their /// debt in `total_reward_debt` would break the FixTotalRewardDebt formula. /// Starts at 0 for existing pools (binary-compatible with old `_reserved3`). + /// + /// LEGACY-ONLY as of the "always close on full unstake" change: full unstakes + /// now pay out in full and redistribute any unpayable remainder, so the + /// current code never *increments* this field. It is only ever decremented, + /// servicing residual balances created by pre-upgrade accounts via the + /// `amount == 0` claim path. The field is retained for binary layout + /// compatibility and to drain those legacy balances; off-chain tooling should + /// not treat it as a live obligation counter for new pools. pub total_residual_unpaid: u64, } @@ -528,7 +536,12 @@ pub struct PoolMetadata { /// UTF-8 URL, zero-padded pub url: [u8; 128], - /// Active staker count + /// Active staker count (best-effort). + /// Incremented on a new stake and decremented on close ONLY when the optional + /// metadata account is supplied to those instructions. Clients that care about + /// an accurate count must pass this metadata PDA on stake AND on every close + /// path (Unstake/CompleteUnstake full unstake, CloseStakeAccount); otherwise + /// the counter can drift upward. It is display-only and never affects funds. pub member_count: u64, /// PDA bump seed diff --git a/programs/chiefstaker/tests/init_griefing.rs b/programs/chiefstaker/tests/init_griefing.rs new file mode 100644 index 0000000..6646725 --- /dev/null +++ b/programs/chiefstaker/tests/init_griefing.rs @@ -0,0 +1,124 @@ +//! Regression test for the pool/vault init griefing vector. +//! +//! The pool (`["pool", mint]`) and vault (`["token_vault", pool]`) PDAs are +//! deterministic, so anyone could pre-fund them with lamports before the real +//! authority initializes. `system_instruction::create_account` fails on a +//! non-zero balance, which would permanently brick pool creation for that mint. +//! `process_initialize_pool` now tolerates a pre-funded PDA (fund-to-rent + +//! allocate + assign under the PDA seeds). This test pre-funds BOTH PDAs and +//! asserts initialization still succeeds. + +use borsh::BorshDeserialize; +use chiefstaker::{ + state::{StakingPool, POOL_SEED, TOKEN_VAULT_SEED}, + StakingInstruction, +}; +use solana_program::{ + instruction::{AccountMeta, Instruction}, + program_option::COption, + program_pack::Pack, + pubkey::Pubkey, + system_program, + sysvar, +}; +use solana_program_test::*; +use solana_sdk::{ + account::Account, + signature::{Keypair, Signer}, + transaction::Transaction, +}; + +fn pack_mint(mint_authority: Pubkey, decimals: u8) -> Vec { + let mut data = vec![0u8; spl_token_2022::state::Mint::LEN]; + spl_token_2022::state::Mint { + mint_authority: COption::Some(mint_authority), + supply: 0, + decimals, + is_initialized: true, + freeze_authority: COption::None, + } + .pack_into_slice(&mut data); + data +} + +#[tokio::test] +async fn initialize_succeeds_when_pdas_are_pre_funded() { + let program_id = chiefstaker::id(); + + // The signer must prove authority over the mint; the simplest valid proof is + // being the mint_authority, so craft a basic Token-2022 mint with this signer + // as the mint authority. + let authority = Keypair::new(); + let mint = Pubkey::new_unique(); + let (pool_pda, _) = Pubkey::find_program_address(&[POOL_SEED, mint.as_ref()], &program_id); + let (vault_pda, _) = Pubkey::find_program_address(&[TOKEN_VAULT_SEED, pool_pda.as_ref()], &program_id); + + let mut pt = ProgramTest::new("chiefstaker", program_id, processor!(chiefstaker::process_instruction)); + pt.add_program( + "spl_token_2022", + spl_token_2022::id(), + processor!(spl_token_2022::processor::Processor::process), + ); + + pt.add_account( + mint, + Account { + lamports: 10_000_000, + data: pack_mint(authority.pubkey(), 9), + owner: spl_token_2022::id(), + executable: false, + rent_epoch: 0, + }, + ); + // Authority funds the rent top-ups + fees. + pt.add_account( + authority.pubkey(), + Account { lamports: 1_000_000_000, data: vec![], owner: system_program::id(), executable: false, rent_epoch: 0 }, + ); + // GRIEF: pre-fund both PDAs with stray lamports, system-owned, empty data. + pt.add_account( + pool_pda, + Account { lamports: 1, data: vec![], owner: system_program::id(), executable: false, rent_epoch: 0 }, + ); + pt.add_account( + vault_pda, + Account { lamports: 1, data: vec![], owner: system_program::id(), executable: false, rent_epoch: 0 }, + ); + + let (mut banks, payer, blockhash) = pt.start().await; + + let ix = Instruction { + program_id, + accounts: vec![ + AccountMeta::new(pool_pda, false), + AccountMeta::new_readonly(mint, false), + AccountMeta::new(vault_pda, false), + AccountMeta::new(authority.pubkey(), true), + AccountMeta::new_readonly(system_program::id(), false), + AccountMeta::new_readonly(spl_token_2022::id(), false), + AccountMeta::new_readonly(sysvar::rent::id(), false), + ], + data: borsh::to_vec(&StakingInstruction::InitializePool { tau_seconds: 60 }).unwrap(), + }; + let tx = Transaction::new_signed_with_payer(&[ix], Some(&payer.pubkey()), &[&payer, &authority], blockhash); + banks + .process_transaction(tx) + .await + .expect("InitializePool must succeed even when the PDAs were pre-funded"); + + // Pool created, owned by the program, with correct state. + let pool_acc = banks.get_account(pool_pda).await.unwrap().unwrap(); + assert_eq!(pool_acc.owner, program_id, "pool must be owned by the program"); + let pool = StakingPool::try_from_slice(&pool_acc.data).unwrap(); + assert!(pool.is_initialized()); + assert_eq!(pool.mint, mint); + assert_eq!(pool.total_staked, 0); + assert_eq!(pool.token_vault, vault_pda); + + // Vault created as a token account owned (token-authority) by the pool PDA. + let vault_acc = banks.get_account(vault_pda).await.unwrap().unwrap(); + assert_eq!(vault_acc.owner, spl_token_2022::id(), "vault must be a token account"); + let vault = spl_token_2022::state::Account::unpack(&vault_acc.data).unwrap(); + assert_eq!(vault.owner, pool_pda, "vault authority must be the pool PDA"); + assert_eq!(vault.mint, mint); +} From 91ce9124e07e056f8aa0f0a9e8ae96745810ac74 Mon Sep 17 00:00:00 2001 From: Mark Karpeles Date: Sun, 24 May 2026 04:05:37 +0900 Subject: [PATCH 2/3] Make member_count exact: require metadata PDA on stake/unstake/close paths MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Previously the metadata account was optional on the count-changing instructions, so member_count could drift (e.g. a full unstake that closed the account without passing metadata never decremented). Make the metadata PDA a REQUIRED account on every member-count-changing instruction — Stake, StakeOnBehalf, Unstake, CompleteUnstake, CloseStakeAccount — and route all updates through a shared state::update_pool_member_count helper that: - rejects a non-canonical metadata account (InvalidPDA), so the correct PDA is always supplied, and - tolerates an uninitialized account (pool that never created metadata) by skipping the update — so metadata-less pools keep working. This makes the counter exact (+1 on a new stake, -1 on close). Unstake and CompleteUnstake also now require the system program positionally (it precedes metadata and is used for legacy realloc). Breaking ABI change: clients must pass the (derived) metadata PDA on these instructions. TS helpers updated to always pass it; lib.rs account tables and idl.json updated; legacy_migration complete test passes the metadata account. Co-Authored-By: Claude Opus 4.7 (1M context) --- idl.json | 33 ++++++++---- .../src/instructions/close_stake.rs | 25 +++------ .../src/instructions/complete_unstake.rs | 13 ++--- .../chiefstaker/src/instructions/stake.rs | 23 +++----- .../src/instructions/stake_on_behalf.rs | 23 +++----- .../chiefstaker/src/instructions/unstake.rs | 29 ++++------ programs/chiefstaker/src/lib.rs | 11 ++-- programs/chiefstaker/src/state.rs | 54 ++++++++++++++++--- .../chiefstaker/tests/legacy_migration.rs | 5 +- tests/typescript/test_staking.ts | 14 +++-- 10 files changed, 127 insertions(+), 103 deletions(-) diff --git a/idl.json b/idl.json index ee14d7b..cec7668 100644 --- a/idl.json +++ b/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", diff --git a/programs/chiefstaker/src/instructions/close_stake.rs b/programs/chiefstaker/src/instructions/close_stake.rs index 66cd91c..7207bc7 100644 --- a/programs/chiefstaker/src/instructions/close_stake.rs +++ b/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); diff --git a/programs/chiefstaker/src/instructions/complete_unstake.rs b/programs/chiefstaker/src/instructions/complete_unstake.rs index 70a9f07..a635e56 100644 --- a/programs/chiefstaker/src/instructions/complete_unstake.rs +++ b/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(()) diff --git a/programs/chiefstaker/src/instructions/stake.rs b/programs/chiefstaker/src/instructions/stake.rs index f0fa32f..f89e9a7 100644 --- a/programs/chiefstaker/src/instructions/stake.rs +++ b/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); diff --git a/programs/chiefstaker/src/instructions/stake_on_behalf.rs b/programs/chiefstaker/src/instructions/stake_on_behalf.rs index 4688069..4313371 100644 --- a/programs/chiefstaker/src/instructions/stake_on_behalf.rs +++ b/programs/chiefstaker/src/instructions/stake_on_behalf.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 on behalf of another user (beneficiary) @@ -52,6 +52,9 @@ pub fn process_stake_on_behalf( let beneficiary_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). See update_pool_member_count for the metadata-less + // pool (uninitialized account) tolerance. + 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_on_behalf( ], )?; - // 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 on behalf of beneficiary", amount); diff --git a/programs/chiefstaker/src/instructions/unstake.rs b/programs/chiefstaker/src/instructions/unstake.rs index dbd1d44..c5be4da 100644 --- a/programs/chiefstaker/src/instructions/unstake.rs +++ b/programs/chiefstaker/src/instructions/unstake.rs @@ -17,7 +17,7 @@ use crate::{ error::StakingError, events::{emit_reward_payout, RewardPayoutType}, math::{calculate_user_weighted_stake, wad_div, wad_mul, U256, WAD}, - state::{is_valid_token_program, PoolMetadata, StakingPool, UserStake, POOL_SEED}, + state::{is_valid_token_program, update_pool_member_count, StakingPool, UserStake, POOL_SEED}, }; /// Accounting-only portion of an unstake, shared by direct `Unstake` @@ -235,20 +235,10 @@ pub fn close_user_stake_account<'a>( stake_data.fill(0); } - // Optional metadata account: decrement member_count on close + // Decrement member_count on close (metadata PDA is a required account on the + // close paths; an uninitialized account is tolerated for metadata-less pools). if let Some(metadata_info) = metadata_info { - 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[..])?; - } - } - } + update_pool_member_count(program_id, pool_info, metadata_info, -1)?; } msg!("Closed user stake account, returned {} lamports", stake_lamports); @@ -475,9 +465,10 @@ pub fn process_unstake( } } - // Optional trailing accounts: system program (legacy realloc) then metadata (close) - 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 a full-unstake close) are both required accounts. + let system_program_info = next_account_info(account_info_iter)?; + let metadata_info = next_account_info(account_info_iter)?; // Execute the shared unstake logic execute_unstake( @@ -492,7 +483,7 @@ pub fn process_unstake( user_info, amount, current_time, - system_program_info, - metadata_info, + Some(system_program_info), + Some(metadata_info), ) } diff --git a/programs/chiefstaker/src/lib.rs b/programs/chiefstaker/src/lib.rs index 45b823f..da563b8 100644 --- a/programs/chiefstaker/src/lib.rs +++ b/programs/chiefstaker/src/lib.rs @@ -55,6 +55,7 @@ pub enum StakingInstruction { /// 5. `[writable, signer]` User/owner /// 6. `[]` System program /// 7. `[]` Token 2022 program + /// 8. `[writable]` Metadata PDA (required; member_count, uninitialized tolerated) Stake { /// Amount of tokens to stake amount: u64, @@ -70,8 +71,8 @@ pub enum StakingInstruction { /// 4. `[]` Token mint /// 5. `[writable, signer]` User/owner /// 6. `[]` Token 2022 program - /// 7. `[]` System program (optional, for legacy account reallocation) - /// 8. `[writable]` Metadata PDA (optional, to decrement member_count on full-unstake close) + /// 7. `[]` System program (required; for legacy account reallocation) + /// 8. `[writable]` Metadata PDA (required; member_count, uninitialized tolerated) Unstake { /// Amount of tokens to unstake amount: u64, @@ -155,8 +156,8 @@ pub enum StakingInstruction { /// 4. `[]` Token mint /// 5. `[writable, signer]` User/owner /// 6. `[]` Token 2022 program - /// 7. `[]` System program (optional) - /// 8. `[writable]` Metadata PDA (optional, to decrement member_count on close) + /// 7. `[]` System program (required) + /// 8. `[writable]` Metadata PDA (required; member_count, uninitialized tolerated) CompleteUnstake, /// Cancel a pending unstake request, restoring the frozen coins to the @@ -175,6 +176,7 @@ pub enum StakingInstruction { /// 0. `[]` Pool account /// 1. `[writable]` User stake account /// 2. `[writable, signer]` User/owner + /// 3. `[writable]` Metadata PDA (required; member_count, uninitialized tolerated) CloseStakeAccount, /// DEPRECATED: Slot 13 reserved for ABI compatibility (was FixTotalRewardDebt). @@ -235,6 +237,7 @@ pub enum StakingInstruction { /// 6. `[writable]` Beneficiary (receives position) /// 7. `[]` System program /// 8. `[]` Token 2022 program + /// 9. `[writable]` Metadata PDA (required; member_count, uninitialized tolerated) StakeOnBehalf { amount: u64, }, diff --git a/programs/chiefstaker/src/state.rs b/programs/chiefstaker/src/state.rs index 2854742..2eec079 100644 --- a/programs/chiefstaker/src/state.rs +++ b/programs/chiefstaker/src/state.rs @@ -536,12 +536,12 @@ pub struct PoolMetadata { /// UTF-8 URL, zero-padded pub url: [u8; 128], - /// Active staker count (best-effort). - /// Incremented on a new stake and decremented on close ONLY when the optional - /// metadata account is supplied to those instructions. Clients that care about - /// an accurate count must pass this metadata PDA on stake AND on every close - /// path (Unstake/CompleteUnstake full unstake, CloseStakeAccount); otherwise - /// the counter can drift upward. It is display-only and never affects funds. + /// Active staker count. The metadata PDA is a REQUIRED account on every + /// member-count-changing instruction (Stake, StakeOnBehalf, Unstake, + /// CompleteUnstake, CloseStakeAccount), so this stays exact: +1 on a new + /// stake, -1 on close. Maintained via `update_pool_member_count`, which + /// tolerates an uninitialized metadata account for pools that never created + /// metadata. Display-only; never affects funds. pub member_count: u64, /// PDA bump seed @@ -573,6 +573,48 @@ impl PoolMetadata { } } +/// Apply `delta` (+1 / -1) to a pool's `member_count`. +/// +/// `metadata_info` must be the canonical metadata PDA for `pool_info`; a +/// non-canonical account is rejected with `InvalidPDA`. This lets callers +/// require the metadata account unconditionally on every member-count-changing +/// instruction (stake / unstake / close) so the counter stays exact. A pool that +/// never created metadata is supported transparently: the canonical PDA is then +/// an uninitialized (system-owned, empty) account, which is detected and the +/// update is skipped. +pub fn update_pool_member_count<'a>( + program_id: &Pubkey, + pool_info: &AccountInfo<'a>, + metadata_info: &AccountInfo<'a>, + delta: i64, +) -> Result<(), solana_program::program_error::ProgramError> { + // The supplied account must be the canonical metadata PDA for this pool. + let (expected_metadata, _) = PoolMetadata::derive_pda(pool_info.key, program_id); + if *metadata_info.key != expected_metadata { + return Err(StakingError::InvalidPDA.into()); + } + + // Metadata-less pool: the canonical PDA is uninitialized — nothing to update. + if metadata_info.owner != program_id || metadata_info.data_is_empty() { + return Ok(()); + } + + let mut metadata = PoolMetadata::try_from_slice(&metadata_info.try_borrow_data()?)?; + if !metadata.is_initialized() || metadata.pool != *pool_info.key { + return Ok(()); + } + + metadata.member_count = if delta >= 0 { + metadata.member_count.saturating_add(delta as u64) + } else { + metadata.member_count.saturating_sub(delta.unsigned_abs()) + }; + + let mut data = metadata_info.try_borrow_mut_data()?; + metadata.serialize(&mut &mut data[..])?; + Ok(()) +} + #[cfg(test)] mod tests { use super::*; diff --git a/programs/chiefstaker/tests/legacy_migration.rs b/programs/chiefstaker/tests/legacy_migration.rs index e029466..9f8451e 100644 --- a/programs/chiefstaker/tests/legacy_migration.rs +++ b/programs/chiefstaker/tests/legacy_migration.rs @@ -11,7 +11,7 @@ use borsh::BorshDeserialize; use chiefstaker::{ math::{wad_mul, U256, WAD}, - state::{StakingPool, UserStake, POOL_SEED, STAKE_SEED, TOKEN_VAULT_SEED}, + state::{StakingPool, UserStake, METADATA_SEED, POOL_SEED, STAKE_SEED, TOKEN_VAULT_SEED}, StakingInstruction, }; use solana_program::{ @@ -177,6 +177,7 @@ async fn legacy_complete_full_removes_and_closes() { let user_token = Keypair::new(); let (stake_pda, stake_bump) = Pubkey::find_program_address(&[STAKE_SEED, pool_pda.as_ref(), user.pubkey().as_ref()], &program_id); + let (metadata_pda, _) = Pubkey::find_program_address(&[METADATA_SEED, pool_pda.as_ref()], &program_id); let staked: u64 = 1_000_000_000; @@ -262,6 +263,8 @@ async fn legacy_complete_full_removes_and_closes() { AccountMeta::new(user.pubkey(), true), AccountMeta::new_readonly(spl_token_2022::id(), false), AccountMeta::new_readonly(system_program::id(), false), + // Metadata PDA (required; uninitialized here — this pool has no metadata) + AccountMeta::new(metadata_pda, false), ], data: borsh::to_vec(&StakingInstruction::CompleteUnstake).unwrap(), }; diff --git a/tests/typescript/test_staking.ts b/tests/typescript/test_staking.ts index 4ac59f3..6d1747d 100644 --- a/tests/typescript/test_staking.ts +++ b/tests/typescript/test_staking.ts @@ -162,6 +162,7 @@ function createStakeInstruction( mint: PublicKey, user: PublicKey, amount: bigint, + metadata: PublicKey, tokenProgramId: PublicKey = TOKEN_2022_PROGRAM_ID, ): TransactionInstruction { const data = Buffer.alloc(1 + 8); @@ -178,6 +179,8 @@ function createStakeInstruction( { pubkey: user, isSigner: true, isWritable: true }, { pubkey: SystemProgram.programId, isSigner: false, isWritable: false }, { pubkey: tokenProgramId, isSigner: false, isWritable: false }, + // Metadata PDA is required (uninitialized account tolerated for pools without metadata) + { pubkey: metadata, isSigner: false, isWritable: true }, ], programId: PROGRAM_ID, data, @@ -715,6 +718,7 @@ class TestContext { async stake(user: Keypair, userToken: PublicKey, amount: bigint): Promise { const [userStakePDA] = deriveUserStakePDA(this.poolPDA, user.publicKey); + const [metadataPDA] = deriveMetadataPDA(this.poolPDA); const ix = createStakeInstruction( this.poolPDA, @@ -724,6 +728,7 @@ class TestContext { this.mint, user.publicKey, amount, + metadataPDA, this.tokenProgramId, ); @@ -733,6 +738,7 @@ class TestContext { async stakeOnBehalf(staker: Keypair, stakerToken: PublicKey, beneficiary: PublicKey, amount: bigint): Promise { const [beneficiaryStakePDA] = deriveUserStakePDA(this.poolPDA, beneficiary); + const [metadataPDA] = deriveMetadataPDA(this.poolPDA); const ix = createStakeOnBehalfInstruction( this.poolPDA, @@ -743,7 +749,7 @@ class TestContext { staker.publicKey, beneficiary, amount, - undefined, + metadataPDA, this.tokenProgramId, ); @@ -764,7 +770,7 @@ class TestContext { user.publicKey, amount, this.tokenProgramId, - withMetadata ? metadataPDA : undefined, + metadataPDA, // metadata PDA is now required on this instruction ); const tx = new Transaction().add(ix); @@ -890,7 +896,7 @@ class TestContext { this.mint, user.publicKey, this.tokenProgramId, - withMetadata ? metadataPDA : undefined, + metadataPDA, // metadata PDA is now required on this instruction ); const tx = new Transaction().add(ix); @@ -1029,7 +1035,7 @@ class TestContext { this.poolPDA, userStakePDA, user.publicKey, - withMetadata ? metadataPDA : undefined, + metadataPDA, // metadata PDA is now required on this instruction ); const tx = new Transaction().add(ix); From 050975568ddbdf11ad6f51ac401dab383123437b Mon Sep 17 00:00:00 2001 From: Mark Karpeles Date: Sun, 24 May 2026 04:45:42 +0900 Subject: [PATCH 3/3] test: update obsolete no-metadata stake test for required-metadata design MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ctx.stake now always passes the metadata PDA, so staking on a metadata pool increments member_count. Repurpose the old "stake without metadata is backwards compatible" test (which asserted member_count stayed 0) to assert it increments to 1 — matching the now-required metadata account. Co-Authored-By: Claude Opus 4.7 (1M context) --- tests/typescript/test_staking.ts | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/tests/typescript/test_staking.ts b/tests/typescript/test_staking.ts index 6d1747d..f27c65d 100644 --- a/tests/typescript/test_staking.ts +++ b/tests/typescript/test_staking.ts @@ -3831,8 +3831,9 @@ async function runTests() { if (meta2.memberCount !== 1n) throw new Error(`Expected 1 after re-set, got ${meta2.memberCount}`); }); - // Test: Stake without metadata account still works (backwards compatible) - await test(`[${tokenProgramLabel}] Stake without metadata account is backwards compatible`, async () => { + // Test: Stake on a metadata pool increments member_count (metadata PDA is now + // a required account on Stake, so the count is maintained exactly). + await test(`[${tokenProgramLabel}] Stake increments member_count on a metadata pool`, async () => { const ctx = new TestContext(connection, Keypair.generate(), programAuthority, tokenProgramId); await ctx.setup(); await ctx.createMintWithMetadata(9, 'BackCompat', 'BCK'); @@ -3844,12 +3845,11 @@ async function runTests() { const userToken = await ctx.createUserTokenAccount(user.publicKey); await ctx.mintTokens(userToken, BigInt(1_000_000_000)); - // Stake WITHOUT passing metadata account (old-style 8-account call) + // ctx.stake always passes the metadata PDA now; a new stake increments the count. await ctx.stake(user, userToken, BigInt(1_000_000_000)); - // member_count should remain 0 since we didn't pass metadata const meta = await ctx.readMetadata(); - if (meta.memberCount !== 0n) throw new Error(`Expected 0 (no metadata passed), got ${meta.memberCount}`); + if (meta.memberCount !== 1n) throw new Error(`Expected 1 after stake, got ${meta.memberCount}`); }); } // end Token 2022-only metadata tests