From e1ee9b241d4db440884005f4c3e0ad88db5ec9e9 Mon Sep 17 00:00:00 2001 From: Mark Karpeles Date: Sat, 23 May 2026 23:04:36 +0900 Subject: [PATCH 1/6] Freeze unstaking coins during cooldown; close account on full unstake RequestUnstake now settles and removes the requested coins from the pool's reward accounting immediately, so coins awaiting unstake stop earning new rewards during the cooldown. Their already-earned rewards are paid out at request time; only the token transfer is deferred to CompleteUnstake. The proportional immature-reward forfeiture and weight removal (linear in amount) already applied on direct Unstake now also apply at request time, so a partial request loses only the proportional share of immature rewards and weight. A full unstake (direct or via CompleteUnstake) now fully resets the position and closes the stake account to reclaim rent, except when the pool still owes residual SOL (kept open so the user can claim it, then close). CancelUnstakeRequest restores the frozen coins to the active position: stake and weight are added back at the original maturity, and the re-added tokens get a fresh reward snapshot so they earn no cooldown-period rewards. - unstake.rs: extract settle_unstake_accounting + close_user_stake_account helpers; execute_unstake closes on full unstake; optional metadata account - request_unstake.rs: settle + auto-pay + freeze; owner now writable - complete_unstake.rs: deliver frozen tokens + close on full unstake - cancel_unstake.rs: restore frozen stake; pool now writable; residual guard - error.rs: add ResidualRewardsPending - lib.rs/idl.json/README.md: updated account tables and docs - test_staking.ts: writable flags, optional metadata, new cooldown tests Co-Authored-By: Claude Opus 4.7 (1M context) --- README.md | 8 +- idl.json | 61 ++++- programs/chiefstaker/src/error.rs | 3 + .../src/instructions/cancel_unstake.rs | 95 +++++++- .../src/instructions/complete_unstake.rs | 97 +++++--- .../src/instructions/request_unstake.rs | 54 ++++- .../chiefstaker/src/instructions/unstake.rs | 132 +++++++++-- programs/chiefstaker/src/lib.rs | 21 +- tests/typescript/test_staking.ts | 222 +++++++++++++++--- 9 files changed, 587 insertions(+), 106 deletions(-) diff --git a/README.md b/README.md index a85b177..2fd1b20 100644 --- a/README.md +++ b/README.md @@ -63,16 +63,16 @@ Rewards can be deposited directly via instruction or sent to the pool PDA (e.g., |---|-------------|-------------| | 0 | `InitializePool` | Create a new staking pool for a Token 2022 mint | | 1 | `Stake` | Stake tokens into the pool | -| 2 | `Unstake` | Unstake tokens (direct, when no cooldown) | +| 2 | `Unstake` | Unstake tokens (direct, when no cooldown); full unstake closes the account | | 3 | `ClaimRewards` | Claim accumulated SOL rewards | | 4 | `DepositRewards` | Deposit SOL rewards into the pool | | 5 | `SyncPool` | Rebase pool math to prevent overflow | | 6 | `SyncRewards` | Sync SOL sent directly to the pool PDA | | 7 | `UpdatePoolSettings` | Set min stake, lock duration, cooldown (authority only) | | 8 | `TransferAuthority` | Transfer or renounce pool authority | -| 9 | `RequestUnstake` | Start unstake cooldown (tokens keep earning) | -| 10 | `CompleteUnstake` | Finish unstake after cooldown elapsed | -| 11 | `CancelUnstakeRequest` | Cancel a pending unstake request | +| 9 | `RequestUnstake` | Start unstake cooldown (requested coins stop earning; earned rewards settled) | +| 10 | `CompleteUnstake` | Finish unstake after cooldown; full unstake closes the account | +| 11 | `CancelUnstakeRequest` | Cancel a pending request, restoring coins to the active position | | 12 | `CloseStakeAccount` | Close zero-balance stake account to reclaim rent | | 13 | ~~`FixTotalRewardDebt`~~ | Deprecated (no-op, returns error) | | 14 | `SetPoolMetadata` | Set pool name, tags, and URL (permissionless) | diff --git a/idl.json b/idl.json index ee12c25..2f2e565 100644 --- a/idl.json +++ b/idl.json @@ -123,7 +123,8 @@ "docs": [ "Unstake tokens from the pool (direct unstake).", "Only available when pool has no cooldown configured.", - "Claims pending rewards and redistributes stranded rewards." + "Claims pending rewards and redistributes stranded rewards.", + "A full unstake closes the user stake account." ], "accounts": [ { @@ -159,6 +160,18 @@ { "name": "tokenProgram", "docs": ["Token program (SPL Token or Token 2022)"] + }, + { + "name": "systemProgram", + "optional": true, + "address": "11111111111111111111111111111111", + "docs": ["System program (optional, for legacy account reallocation)"] + }, + { + "name": "metadata", + "writable": true, + "optional": true, + "docs": ["Metadata PDA (optional, to decrement member_count on full-unstake close)"] } ], "args": [ @@ -332,7 +345,10 @@ "discriminator": [9], "docs": [ "Request unstake - starts cooldown period.", - "Tokens remain staked during cooldown. Only one pending request at a time." + "The requested coins stop earning rewards immediately (removed from pool", + "weight) and their already-earned rewards are settled and paid out now.", + "Only the token transfer is deferred to completeUnstake.", + "Only one pending request at a time." ], "accounts": [ { @@ -347,8 +363,15 @@ }, { "name": "owner", + "writable": true, "signer": true, - "docs": ["User / owner"] + "docs": ["User / owner (receives settled reward SOL)"] + }, + { + "name": "systemProgram", + "optional": true, + "address": "11111111111111111111111111111111", + "docs": ["System program (optional, for legacy account reallocation)"] } ], "args": [ @@ -363,7 +386,8 @@ "discriminator": [10], "docs": [ "Complete unstake after cooldown has elapsed.", - "Claims pending rewards and redistributes stranded rewards." + "Delivers the frozen tokens from the vault to the user (reward settlement", + "already happened at requestUnstake). A full unstake closes the account." ], "accounts": [ { @@ -399,6 +423,18 @@ { "name": "tokenProgram", "docs": ["Token program (SPL Token or Token 2022)"] + }, + { + "name": "systemProgram", + "optional": true, + "address": "11111111111111111111111111111111", + "docs": ["System program (optional)"] + }, + { + "name": "metadata", + "writable": true, + "optional": true, + "docs": ["Metadata PDA (optional, to decrement member_count on close)"] } ], "args": [] @@ -407,12 +443,14 @@ "name": "cancelUnstakeRequest", "discriminator": [11], "docs": [ - "Cancel a pending unstake request.", - "Tokens remain staked." + "Cancel a pending unstake request, restoring the frozen coins to the", + "active staking position (re-grants weight; re-added tokens get a fresh", + "reward snapshot so they earn no cooldown-period rewards)." ], "accounts": [ { "name": "pool", + "writable": true, "docs": ["Pool account"] }, { @@ -424,6 +462,12 @@ "name": "owner", "signer": true, "docs": ["User / owner"] + }, + { + "name": "systemProgram", + "optional": true, + "address": "11111111111111111111111111111111", + "docs": ["System program (optional, for legacy account reallocation)"] } ], "args": [] @@ -1169,6 +1213,11 @@ "code": 6034, "name": "RewardDebtExceedsBound", "msg": "New total_reward_debt exceeds maximum accumulated rewards" + }, + { + "code": 6035, + "name": "ResidualRewardsPending", + "msg": "Claim residual rewards before cancelling a full unstake request" } ] } diff --git a/programs/chiefstaker/src/error.rs b/programs/chiefstaker/src/error.rs index e95e2bf..7d6ba5e 100644 --- a/programs/chiefstaker/src/error.rs +++ b/programs/chiefstaker/src/error.rs @@ -108,6 +108,9 @@ pub enum StakingError { #[error("New total_reward_debt exceeds maximum accumulated rewards")] RewardDebtExceedsBound, + + #[error("Claim residual rewards before cancelling a full unstake request")] + ResidualRewardsPending, } impl From for ProgramError { diff --git a/programs/chiefstaker/src/instructions/cancel_unstake.rs b/programs/chiefstaker/src/instructions/cancel_unstake.rs index ea4fef2..f202718 100644 --- a/programs/chiefstaker/src/instructions/cancel_unstake.rs +++ b/programs/chiefstaker/src/instructions/cancel_unstake.rs @@ -10,15 +10,24 @@ use solana_program::{ use crate::{ error::StakingError, + math::{wad_mul, U256, WAD}, state::{StakingPool, UserStake}, }; -/// Cancel a pending unstake request +/// Cancel a pending unstake request, restoring the frozen coins to the active +/// staking position so they earn rewards again. +/// +/// RequestUnstake removed the frozen coins from the pool's reward accounting. This +/// adds them back: total_staked and sum_stake_exp are restored (re-granting the +/// weight at the original maturity, since exp_start_factor is unchanged), and the +/// re-added tokens receive a fresh reward snapshot so they do not retroactively +/// claim rewards distributed during the cooldown (mirrors add-stake). /// /// Accounts: -/// 0. `[]` Pool account +/// 0. `[writable]` Pool account /// 1. `[writable]` User stake account /// 2. `[signer]` User/owner +/// 3. `[]` System program (optional, for legacy account reallocation) pub fn process_cancel_unstake_request( program_id: &Pubkey, accounts: &[AccountInfo], @@ -38,7 +47,7 @@ pub fn process_cancel_unstake_request( if pool_info.owner != program_id { return Err(StakingError::InvalidAccountOwner.into()); } - let pool = StakingPool::try_from_slice(&pool_info.try_borrow_data()?)?; + let mut pool = StakingPool::try_from_slice(&pool_info.try_borrow_data()?)?; if !pool.is_initialized() { return Err(StakingError::NotInitialized.into()); } @@ -49,6 +58,11 @@ pub fn process_cancel_unstake_request( return Err(StakingError::InvalidPDA.into()); } + // Check if pool needs rebasing (we are about to grow sum_stake_exp) + if pool.get_sum_stake_exp().needs_rebase() { + return Err(StakingError::PoolRequiresSync.into()); + } + // Realloc legacy accounts to current size (payer = user) // System program is optional trailing account, only needed for legacy accounts let system_program_info = account_info_iter.next(); @@ -78,25 +92,84 @@ pub fn process_cancel_unstake_request( return Err(StakingError::InvalidPDA.into()); } - // Lazily adjust exp_start_factor if pool has been rebased - user_stake.sync_to_pool(&pool)?; - // Check there is a pending request if !user_stake.has_pending_unstake_request() { return Err(StakingError::NoPendingUnstakeRequest.into()); } - let cancelled_amount = user_stake.unstake_request_amount; + // Lazily adjust exp_start_factor if pool has been rebased + user_stake.sync_to_pool(&pool)?; + + let frozen = user_stake.unstake_request_amount; + + // Edge case: a full unstake request (amount == 0) that still has unpaid + // residual rewards stores those in reward_debt. Reactivating the position + // would overwrite that storage with a snapshot, losing the residual. Require + // the user to claim it first (ClaimRewards works on the amount==0 path). + if user_stake.amount == 0 && user_stake.reward_debt / WAD > 0 { + return Err(StakingError::ResidualRewardsPending.into()); + } + + // Restore the frozen coins to the active position. + let frozen_wad = (frozen as u128) + .checked_mul(WAD) + .ok_or(StakingError::MathOverflow)?; + + // sum_stake_exp += frozen * exp_start_factor (restores weight / maturity) + let contribution = wad_mul(frozen_wad, user_stake.exp_start_factor)?; + let new_sum = pool + .get_sum_stake_exp() + .checked_add(U256::from_u128(contribution)) + .ok_or(StakingError::MathOverflow)?; + pool.set_sum_stake_exp(new_sum); + + // total_staked += frozen + pool.total_staked = pool + .total_staked + .checked_add(frozen as u128) + .ok_or(StakingError::MathOverflow)?; + + // Fresh reward snapshot for the re-added tokens (no retroactive rewards) + let new_token_debt = wad_mul(frozen_wad, pool.acc_reward_per_weighted_share)?; + + if user_stake.amount == 0 { + // Was a full request (no residual remains): start clean. + user_stake.reward_debt = new_token_debt; + user_stake.claimed_rewards_wad = 0; + } else { + // Partial request: keep the active remainder's pending intact and add a + // fresh debt snapshot for the re-added tokens (mirrors add-stake). + user_stake.reward_debt = user_stake + .reward_debt + .checked_add(new_token_debt) + .ok_or(StakingError::MathOverflow)?; + } + + user_stake.amount = user_stake + .amount + .checked_add(frozen) + .ok_or(StakingError::MathOverflow)?; + + pool.total_reward_debt = pool + .total_reward_debt + .checked_add(new_token_debt) + .ok_or(StakingError::MathOverflow)?; // Clear the request fields user_stake.unstake_request_amount = 0; user_stake.unstake_request_time = 0; - // Save user stake - let mut stake_data = user_stake_info.try_borrow_mut_data()?; - user_stake.serialize(&mut &mut stake_data[..])?; + // Save pool and user stake + { + let mut pool_data = pool_info.try_borrow_mut_data()?; + pool.serialize(&mut &mut pool_data[..])?; + } + { + let mut stake_data = user_stake_info.try_borrow_mut_data()?; + user_stake.serialize(&mut &mut stake_data[..])?; + } - msg!("Cancelled unstake request for {} tokens", cancelled_amount); + msg!("Cancelled unstake request for {} tokens", frozen); Ok(()) } diff --git a/programs/chiefstaker/src/instructions/complete_unstake.rs b/programs/chiefstaker/src/instructions/complete_unstake.rs index 954b6e8..9d6c12e 100644 --- a/programs/chiefstaker/src/instructions/complete_unstake.rs +++ b/programs/chiefstaker/src/instructions/complete_unstake.rs @@ -1,22 +1,31 @@ //! Complete unstake instruction (after cooldown period has elapsed) -use borsh::BorshDeserialize; +use borsh::{BorshDeserialize, BorshSerialize}; use solana_program::{ account_info::{next_account_info, AccountInfo}, clock::Clock, entrypoint::ProgramResult, + msg, + program::invoke_signed, pubkey::Pubkey, sysvar::Sysvar, }; +use spl_token_2022::extension::StateWithExtensions; use crate::{ error::StakingError, - state::{is_valid_token_program, StakingPool, UserStake}, + math::WAD, + state::{is_valid_token_program, StakingPool, UserStake, POOL_SEED}, }; -use super::unstake::execute_unstake; +use super::unstake::close_user_stake_account; -/// Complete unstake after cooldown has elapsed +/// Complete unstake after the cooldown has elapsed. +/// +/// The reward settlement and pool accounting already happened at RequestUnstake; +/// this instruction only delivers the frozen tokens from the vault to the user +/// and, when the position is fully unstaked with nothing owed, closes the stake +/// account to reclaim its rent. /// /// Accounts (same as Unstake): /// 0. `[writable]` Pool account @@ -26,6 +35,8 @@ use super::unstake::execute_unstake; /// 4. `[]` Token mint /// 5. `[writable, signer]` User/owner /// 6. `[]` Token 2022 program +/// 7. `[]` System program (optional, accepted for call-site symmetry) +/// 8. `[writable]` Metadata PDA (optional, to decrement member_count on close) pub fn process_complete_unstake( program_id: &Pubkey, accounts: &[AccountInfo], @@ -54,7 +65,7 @@ pub fn process_complete_unstake( if pool_info.owner != program_id { return Err(StakingError::InvalidAccountOwner.into()); } - let mut pool = StakingPool::try_from_slice(&pool_info.try_borrow_data()?)?; + let pool = StakingPool::try_from_slice(&pool_info.try_borrow_data()?)?; if !pool.is_initialized() { return Err(StakingError::NotInitialized.into()); } @@ -117,31 +128,61 @@ pub fn process_complete_unstake( return Err(StakingError::CooldownNotElapsed.into()); } - // Lazily adjust exp_start_factor if pool has been rebased - user_stake.sync_to_pool(&pool)?; - - let amount = user_stake.unstake_request_amount; + // The frozen tokens awaiting withdrawal (already settled at request time) + let withdraw_amount = user_stake.unstake_request_amount; - // Clear the request fields before execute_unstake (which serializes) + // Clear the request fields user_stake.unstake_request_amount = 0; user_stake.unstake_request_time = 0; - // Optional trailing system program for legacy account reallocation - let system_program_info = account_info_iter.next(); - - // Execute the shared unstake logic - execute_unstake( - program_id, - &mut pool, - &mut user_stake, - pool_info, - user_stake_info, - token_vault_info, - user_token_info, - mint_info, - user_info, - amount, - current_time, - system_program_info, - ) + // Close iff fully unstaked and nothing owed; otherwise keep the account so the + // remaining position keeps earning / residual rewards stay claimable. + let should_close = user_stake.amount == 0 && user_stake.reward_debt / WAD == 0; + + // Persist the cleared request fields before the CPI + { + let mut stake_data = user_stake_info.try_borrow_mut_data()?; + user_stake.serialize(&mut &mut stake_data[..])?; + } + + // Transfer the frozen tokens from vault to user (CPI, pool PDA signs) + let mint_data = mint_info.try_borrow_data()?; + let mint = StateWithExtensions::::unpack(&mint_data)?; + let decimals = mint.base.decimals; + drop(mint_data); + + let pool_seeds = &[POOL_SEED, pool.mint.as_ref(), &[pool.bump]]; + + invoke_signed( + &spl_token_2022::instruction::transfer_checked( + mint_info.owner, + token_vault_info.key, + mint_info.key, + user_token_info.key, + pool_info.key, + &[], + withdraw_amount, + decimals, + )?, + &[ + token_vault_info.clone(), + mint_info.clone(), + user_token_info.clone(), + pool_info.clone(), + ], + &[pool_seeds], + )?; + + msg!("Completed unstake of {} tokens", withdraw_amount); + + // Optional trailing accounts: system program (unused here) then metadata (close) + let _system_program_info = account_info_iter.next(); + let metadata_info = account_info_iter.next(); + + // 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)?; + } + + Ok(()) } diff --git a/programs/chiefstaker/src/instructions/request_unstake.rs b/programs/chiefstaker/src/instructions/request_unstake.rs index e04ff2e..b16f89c 100644 --- a/programs/chiefstaker/src/instructions/request_unstake.rs +++ b/programs/chiefstaker/src/instructions/request_unstake.rs @@ -12,15 +12,26 @@ use solana_program::{ use crate::{ error::StakingError, + events::{emit_reward_payout, RewardPayoutType}, state::{StakingPool, UserStake}, }; -/// Request unstake - starts cooldown period. Tokens remain staked and earn rewards. +use super::unstake::settle_unstake_accounting; + +/// Request unstake - starts the cooldown period. +/// +/// The requested coins stop earning rewards immediately: their stake and weight +/// are removed from the pool (so they no longer share in new rewards), and the +/// rewards they have already earned are settled and paid out now. Only the token +/// transfer is deferred to CompleteUnstake — the tokens stay in the vault until +/// the cooldown elapses. A pending request can be reversed with CancelUnstake, +/// which restores the frozen coins to the active position. /// /// Accounts: /// 0. `[writable]` Pool account /// 1. `[writable]` User stake account -/// 2. `[signer]` User/owner +/// 2. `[writable, signer]` User/owner (receives settled reward SOL) +/// 3. `[]` System program (optional, for legacy account reallocation) pub fn process_request_unstake( program_id: &Pubkey, accounts: &[AccountInfo], @@ -45,7 +56,7 @@ pub fn process_request_unstake( if pool_info.owner != program_id { return Err(StakingError::InvalidAccountOwner.into()); } - let pool = StakingPool::try_from_slice(&pool_info.try_borrow_data()?)?; + let mut pool = StakingPool::try_from_slice(&pool_info.try_borrow_data()?)?; if !pool.is_initialized() { return Err(StakingError::NotInitialized.into()); } @@ -120,13 +131,42 @@ pub fn process_request_unstake( } } - // Set unstake request fields + // Settle and remove the requested coins from the pool's reward accounting so + // they stop earning during the cooldown. This pays out their already-earned + // rewards and reduces user_stake.amount to the still-active remainder; the + // requested tokens remain in the vault until CompleteUnstake. + let reward_transfer_amount = settle_unstake_accounting( + &mut pool, + &mut user_stake, + pool_info, + user_stake_info, + user_info, + amount, + current_time, + system_program_info, + )?; + + // Record the pending request (the frozen amount awaiting withdrawal) user_stake.unstake_request_amount = amount; user_stake.unstake_request_time = current_time; - // Save user stake - let mut stake_data = user_stake_info.try_borrow_mut_data()?; - user_stake.serialize(&mut &mut stake_data[..])?; + // Save pool and user stake + { + let mut pool_data = pool_info.try_borrow_mut_data()?; + pool.serialize(&mut &mut pool_data[..])?; + } + { + let mut stake_data = user_stake_info.try_borrow_mut_data()?; + user_stake.serialize(&mut &mut stake_data[..])?; + } + + // Pay out settled rewards (no token CPI here, so SOL can move directly) + if reward_transfer_amount > 0 { + **pool_info.try_borrow_mut_lamports()? -= reward_transfer_amount; + **user_info.try_borrow_mut_lamports()? += reward_transfer_amount; + msg!("Claimed {} lamports in rewards", reward_transfer_amount); + emit_reward_payout(pool_info.key, user_info.key, reward_transfer_amount, RewardPayoutType::Unstake); + } msg!( "Unstake request created for {} tokens, cooldown {} seconds", diff --git a/programs/chiefstaker/src/instructions/unstake.rs b/programs/chiefstaker/src/instructions/unstake.rs index 19e5ca7..69545d1 100644 --- a/programs/chiefstaker/src/instructions/unstake.rs +++ b/programs/chiefstaker/src/instructions/unstake.rs @@ -7,6 +7,7 @@ use solana_program::{ entrypoint::ProgramResult, msg, program::invoke_signed, + program_error::ProgramError, pubkey::Pubkey, sysvar::Sysvar, }; @@ -16,35 +17,38 @@ 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, StakingPool, UserStake, POOL_SEED}, + state::{is_valid_token_program, PoolMetadata, StakingPool, UserStake, POOL_SEED}, }; -/// Shared unstake logic used by both process_unstake and process_complete_unstake. -/// Handles: reward claiming, pool math updates (sum_stake_exp, total_staked), -/// reward_debt recalculation, and token transfer. +/// Accounting-only portion of an unstake, shared by direct `Unstake` +/// (`execute_unstake`) and the cooldown `RequestUnstake` flow. +/// +/// Settles pending rewards, redistributes the unstaked portion's forfeited +/// immature SOL, removes the unstaked weight/stake from the pool, recomputes the +/// remaining position's reward debt, and pre-updates `last_synced_lamports`. +/// +/// Mutates `pool` and `user_stake` in memory and reallocs the stake account if it +/// is a legacy size, but does NOT serialize accounts, transfer tokens, or move +/// reward lamports — the caller does those. Returns the lamports the caller must +/// transfer pool -> user as the reward payout. /// /// Assumes all account validation has been done by the caller. -pub fn execute_unstake<'a>( - _program_id: &Pubkey, +pub fn settle_unstake_accounting<'a>( pool: &mut StakingPool, user_stake: &mut UserStake, pool_info: &AccountInfo<'a>, user_stake_info: &AccountInfo<'a>, - token_vault_info: &AccountInfo<'a>, - user_token_info: &AccountInfo<'a>, - mint_info: &AccountInfo<'a>, user_info: &AccountInfo<'a>, amount: u64, current_time: i64, system_program_info: Option<&AccountInfo<'a>>, -) -> ProgramResult { - +) -> Result { // Capture old reward_debt for total_reward_debt bookkeeping let old_reward_debt = user_stake.reward_debt; - // Calculate pending rewards (but defer SOL transfer until after token CPI, - // because the Solana runtime verifies CPI account balances and user_info - // is not a CPI account) + // Calculate pending rewards (the caller defers the SOL transfer until after + // any token CPI, because the Solana runtime verifies CPI account balances and + // user_info is not a CPI account). let mut reward_transfer_amount: u64 = 0; let user_weighted = calculate_user_weighted_stake( @@ -92,7 +96,7 @@ pub fn execute_unstake<'a>( .ok_or(StakingError::MathOverflow)?; unpaid_rewards_wad = pending.saturating_sub(paid_wad); - // Pre-update last_synced_lamports (actual SOL transfer deferred to after CPI) + // Pre-update last_synced_lamports (actual SOL transfer deferred to caller) if reward_transfer_amount > 0 { pool.last_synced_lamports = pool.last_synced_lamports.saturating_sub(reward_transfer_amount); } @@ -205,6 +209,93 @@ pub fn execute_unstake<'a>( // Realloc legacy accounts to current size (payer = user) UserStake::maybe_realloc(user_stake_info, user_info, system_program_info)?; + Ok(reward_transfer_amount) +} + +/// Close a fully-unstaked user stake account: drain its rent lamports to the +/// user, zero its data so it can't be re-read as a valid stake, and decrement +/// the pool member_count if a metadata account is supplied. +/// +/// Mirrors `process_close_stake_account`. The caller must ensure the position is +/// fully unstaked with no residual rewards owed before calling. +pub fn close_user_stake_account<'a>( + program_id: &Pubkey, + pool_info: &AccountInfo<'a>, + user_stake_info: &AccountInfo<'a>, + user_info: &AccountInfo<'a>, + metadata_info: Option<&AccountInfo<'a>>, +) -> ProgramResult { + // Transfer all lamports from stake account to user (closes the account) + let stake_lamports = user_stake_info.lamports(); + **user_stake_info.try_borrow_mut_lamports()? = 0; + **user_info.try_borrow_mut_lamports()? += stake_lamports; + + // Zero out the account data so it can't be re-read as a valid stake + { + 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) = 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[..])?; + } + } + } + } + + msg!("Closed user stake account, returned {} lamports", stake_lamports); + Ok(()) +} + +/// Shared unstake logic for the direct `Unstake` path (cooldown == 0). +/// +/// Settles rewards and pool math via `settle_unstake_accounting`, transfers the +/// staked tokens back to the user, pays out the reward SOL, and — when the +/// position is fully unstaked with nothing owed — closes the stake account. +/// +/// Assumes all account validation has been done by the caller. +#[allow(clippy::too_many_arguments)] +pub fn execute_unstake<'a>( + program_id: &Pubkey, + pool: &mut StakingPool, + user_stake: &mut UserStake, + pool_info: &AccountInfo<'a>, + user_stake_info: &AccountInfo<'a>, + token_vault_info: &AccountInfo<'a>, + user_token_info: &AccountInfo<'a>, + mint_info: &AccountInfo<'a>, + user_info: &AccountInfo<'a>, + amount: u64, + current_time: i64, + system_program_info: Option<&AccountInfo<'a>>, + metadata_info: Option<&AccountInfo<'a>>, +) -> ProgramResult { + let reward_transfer_amount = settle_unstake_accounting( + pool, + user_stake, + pool_info, + user_stake_info, + user_info, + amount, + current_time, + system_program_info, + )?; + + // A full unstake fully resets the position. Close the account afterward when + // nothing is owed; if residual rewards remain (pool underfunded), keep it open + // so the user can claim them, then close. + let fully_unstaked = user_stake.amount == 0; + let has_residual = user_stake.reward_debt / WAD > 0; + // Save states (before CPI — pool data includes pre-updated last_synced_lamports) { let mut pool_data = pool_info.try_borrow_mut_data()?; @@ -254,6 +345,11 @@ pub fn execute_unstake<'a>( msg!("Unstaked {} tokens", amount); + // Close the account when fully unstaked and nothing is owed. + if fully_unstaked && !has_residual { + close_user_stake_account(program_id, pool_info, user_stake_info, user_info, metadata_info)?; + } + Ok(()) } @@ -267,6 +363,8 @@ pub fn execute_unstake<'a>( /// 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) pub fn process_unstake( program_id: &Pubkey, accounts: &[AccountInfo], @@ -380,8 +478,9 @@ pub fn process_unstake( } } - // Optional trailing system program for legacy account reallocation + // 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(); // Execute the shared unstake logic execute_unstake( @@ -397,5 +496,6 @@ pub fn process_unstake( amount, current_time, system_program_info, + metadata_info, ) } diff --git a/programs/chiefstaker/src/lib.rs b/programs/chiefstaker/src/lib.rs index 0601d51..45b823f 100644 --- a/programs/chiefstaker/src/lib.rs +++ b/programs/chiefstaker/src/lib.rs @@ -70,6 +70,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) Unstake { /// Amount of tokens to unstake amount: u64, @@ -128,17 +130,22 @@ pub enum StakingInstruction { new_authority: Pubkey, }, - /// Request unstake - starts cooldown period (tokens remain staked) + /// Request unstake - starts cooldown period. + /// The requested coins stop earning immediately (removed from pool weight) and + /// their already-earned rewards are settled/paid out. Only the token transfer + /// is deferred to CompleteUnstake. /// /// Accounts: /// 0. `[writable]` Pool account /// 1. `[writable]` User stake account - /// 2. `[signer]` User/owner + /// 2. `[writable, signer]` User/owner (receives settled reward SOL) + /// 3. `[]` System program (optional, for legacy account reallocation) RequestUnstake { amount: u64, }, - /// Complete unstake after cooldown elapsed + /// Complete unstake after cooldown elapsed. + /// Delivers the frozen tokens and closes the account on a full unstake. /// /// Accounts (same as Unstake): /// 0. `[writable]` Pool account @@ -148,14 +155,18 @@ 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) CompleteUnstake, - /// Cancel a pending unstake request + /// Cancel a pending unstake request, restoring the frozen coins to the + /// active staking position. /// /// Accounts: - /// 0. `[]` Pool account + /// 0. `[writable]` Pool account /// 1. `[writable]` User stake account /// 2. `[signer]` User/owner + /// 3. `[]` System program (optional, for legacy account reallocation) CancelUnstakeRequest, /// Close a zero-balance user stake account to reclaim rent diff --git a/tests/typescript/test_staking.ts b/tests/typescript/test_staking.ts index e0f149f..197d98c 100644 --- a/tests/typescript/test_staking.ts +++ b/tests/typescript/test_staking.ts @@ -172,21 +172,29 @@ function createUnstakeInstruction( user: PublicKey, amount: bigint, tokenProgramId: PublicKey = TOKEN_2022_PROGRAM_ID, + metadata?: PublicKey, ): TransactionInstruction { const data = Buffer.alloc(1 + 8); data.writeUInt8(InstructionType.Unstake, 0); data.writeBigUInt64LE(amount, 1); + const keys = [ + { pubkey: pool, isSigner: false, isWritable: true }, + { pubkey: userStake, isSigner: false, isWritable: true }, + { pubkey: tokenVault, isSigner: false, isWritable: true }, + { pubkey: userToken, isSigner: false, isWritable: true }, + { pubkey: mint, isSigner: false, isWritable: false }, + { pubkey: user, isSigner: true, isWritable: true }, + { pubkey: tokenProgramId, isSigner: false, isWritable: false }, + ]; + if (metadata) { + // Optional trailing accounts are positional: system program then metadata. + keys.push({ pubkey: SystemProgram.programId, isSigner: false, isWritable: false }); + keys.push({ pubkey: metadata, isSigner: false, isWritable: true }); + } + return new TransactionInstruction({ - keys: [ - { pubkey: pool, isSigner: false, isWritable: true }, - { pubkey: userStake, isSigner: false, isWritable: true }, - { pubkey: tokenVault, isSigner: false, isWritable: true }, - { pubkey: userToken, isSigner: false, isWritable: true }, - { pubkey: mint, isSigner: false, isWritable: false }, - { pubkey: user, isSigner: true, isWritable: true }, - { pubkey: tokenProgramId, isSigner: false, isWritable: false }, - ], + keys, programId: PROGRAM_ID, data, }); @@ -365,7 +373,8 @@ function createRequestUnstakeInstruction( keys: [ { pubkey: pool, isSigner: false, isWritable: true }, { pubkey: userStake, isSigner: false, isWritable: true }, - { pubkey: user, isSigner: true, isWritable: false }, + // owner is writable: RequestUnstake settles & pays out earned rewards now + { pubkey: user, isSigner: true, isWritable: true }, ], programId: PROGRAM_ID, data, @@ -380,20 +389,28 @@ function createCompleteUnstakeInstruction( mint: PublicKey, user: PublicKey, tokenProgramId: PublicKey = TOKEN_2022_PROGRAM_ID, + metadata?: PublicKey, ): TransactionInstruction { const data = Buffer.alloc(1); data.writeUInt8(InstructionType.CompleteUnstake, 0); + const keys = [ + { pubkey: pool, isSigner: false, isWritable: true }, + { pubkey: userStake, isSigner: false, isWritable: true }, + { pubkey: tokenVault, isSigner: false, isWritable: true }, + { pubkey: userToken, isSigner: false, isWritable: true }, + { pubkey: mint, isSigner: false, isWritable: false }, + { pubkey: user, isSigner: true, isWritable: true }, + { pubkey: tokenProgramId, isSigner: false, isWritable: false }, + ]; + if (metadata) { + // Optional trailing accounts are positional: system program then metadata. + keys.push({ pubkey: SystemProgram.programId, isSigner: false, isWritable: false }); + keys.push({ pubkey: metadata, isSigner: false, isWritable: true }); + } + return new TransactionInstruction({ - keys: [ - { pubkey: pool, isSigner: false, isWritable: true }, - { pubkey: userStake, isSigner: false, isWritable: true }, - { pubkey: tokenVault, isSigner: false, isWritable: true }, - { pubkey: userToken, isSigner: false, isWritable: true }, - { pubkey: mint, isSigner: false, isWritable: false }, - { pubkey: user, isSigner: true, isWritable: true }, - { pubkey: tokenProgramId, isSigner: false, isWritable: false }, - ], + keys, programId: PROGRAM_ID, data, }); @@ -409,7 +426,8 @@ function createCancelUnstakeRequestInstruction( return new TransactionInstruction({ keys: [ - { pubkey: pool, isSigner: false, isWritable: false }, + // pool is writable: CancelUnstakeRequest restores stake to pool accounting + { pubkey: pool, isSigner: false, isWritable: true }, { pubkey: userStake, isSigner: false, isWritable: true }, { pubkey: user, isSigner: true, isWritable: false }, ], @@ -712,8 +730,9 @@ class TestContext { return await sendAndConfirmTransaction(this.connection, tx, [this.payer, staker]); } - async unstake(user: Keypair, userToken: PublicKey, amount: bigint): Promise { + async unstake(user: Keypair, userToken: PublicKey, amount: bigint, withMetadata: boolean = false): Promise { const [userStakePDA] = deriveUserStakePDA(this.poolPDA, user.publicKey); + const [metadataPDA] = deriveMetadataPDA(this.poolPDA); const ix = createUnstakeInstruction( this.poolPDA, @@ -724,6 +743,7 @@ class TestContext { user.publicKey, amount, this.tokenProgramId, + withMetadata ? metadataPDA : undefined, ); const tx = new Transaction().add(ix); @@ -837,8 +857,9 @@ class TestContext { return await sendAndConfirmTransaction(this.connection, tx, [this.payer, user]); } - async completeUnstake(user: Keypair, userToken: PublicKey): Promise { + async completeUnstake(user: Keypair, userToken: PublicKey, withMetadata: boolean = false): Promise { const [userStakePDA] = deriveUserStakePDA(this.poolPDA, user.publicKey); + const [metadataPDA] = deriveMetadataPDA(this.poolPDA); const ix = createCompleteUnstakeInstruction( this.poolPDA, @@ -848,6 +869,7 @@ class TestContext { this.mint, user.publicKey, this.tokenProgramId, + withMetadata ? metadataPDA : undefined, ); const tx = new Transaction().add(ix); @@ -1704,6 +1726,142 @@ async function runTests() { if (!doubleRequestFailed) throw new Error('Double request should fail'); }); + // Test: requested coins stop earning rewards during cooldown (req #4) + await test(`[${tokenProgramLabel}] Cooldown: requested coins stop earning rewards`, async () => { + const ctx = new TestContext(connection, Keypair.generate(), programAuthority, tokenProgramId); + await ctx.setup(); + await ctx.createMint(9); + await ctx.initializePool(BigInt(60)); + await ctx.updatePoolSettings(ctx.payer, null, null, BigInt(60)); + + // Two equal stakers + const alice = Keypair.generate(); + const bob = Keypair.generate(); + await airdropAndConfirm(connection, alice.publicKey, LAMPORTS_PER_SOL); + await airdropAndConfirm(connection, bob.publicKey, LAMPORTS_PER_SOL); + const aliceToken = await ctx.createUserTokenAccount(alice.publicKey); + const bobToken = await ctx.createUserTokenAccount(bob.publicKey); + await ctx.mintTokens(aliceToken, BigInt(1_000_000_000)); + await ctx.mintTokens(bobToken, BigInt(1_000_000_000)); + await ctx.stake(alice, aliceToken, BigInt(1_000_000_000)); + await ctx.stake(bob, bobToken, BigInt(1_000_000_000)); + + // Let weight accrue a little + console.log(' Waiting 4s for weight...'); + await new Promise(r => setTimeout(r, 4000)); + + // Alice requests a full unstake -> her coins leave the pool's reward accounting + await ctx.requestUnstake(alice, BigInt(1_000_000_000)); + + const poolAfterReq = await ctx.readPoolState(); + if (poolAfterReq.totalStaked !== BigInt(1_000_000_000)) { + throw new Error(`Expected total_staked 1000000000 after Alice's request, got ${poolAfterReq.totalStaked}`); + } + + // New rewards arrive AFTER the request + await ctx.depositRewards(BigInt(LAMPORTS_PER_SOL / 2)); + + // Bob (still active) earns; Alice (frozen) earns nothing new + const bobBefore = await ctx.getBalance(bob.publicKey); + await ctx.claimRewards(bob); + const bobReward = await ctx.getBalance(bob.publicKey) - bobBefore; + + const aliceBefore = await ctx.getBalance(alice.publicKey); + await ctx.claimRewards(alice); + const aliceReward = await ctx.getBalance(alice.publicKey) - aliceBefore; + + console.log(` Bob reward: ${bobReward}, Alice (frozen) reward: ${aliceReward}`); + if (bobReward <= 0) throw new Error(`Active staker Bob should earn from the new deposit, got ${bobReward}`); + if (aliceReward > 0) throw new Error(`Frozen staker Alice should earn nothing new, got ${aliceReward}`); + }); + + // Test: partial request removes proportional stake/weight (req #1, #3) + await test(`[${tokenProgramLabel}] Cooldown: partial request removes proportional stake`, async () => { + const ctx = new TestContext(connection, Keypair.generate(), programAuthority, tokenProgramId); + await ctx.setup(); + await ctx.createMint(9); + await ctx.initializePool(BigInt(2592000)); + await ctx.updatePoolSettings(ctx.payer, null, null, BigInt(60)); + + const user = Keypair.generate(); + await airdropAndConfirm(connection, user.publicKey, LAMPORTS_PER_SOL); + const userToken = await ctx.createUserTokenAccount(user.publicKey); + await ctx.mintTokens(userToken, BigInt(1_000_000_000)); + await ctx.stake(user, userToken, BigInt(1_000_000_000)); + + // Request 40% -> active stake (and thus weight) drops to 60% + await ctx.requestUnstake(user, BigInt(400_000_000)); + + const pool = await ctx.readPoolState(); + if (pool.totalStaked !== BigInt(600_000_000)) { + throw new Error(`Expected total_staked 600000000, got ${pool.totalStaked}`); + } + const state = await ctx.readUserStakeState(user.publicKey); + if (state.amount !== BigInt(600_000_000)) { + throw new Error(`Expected active amount 600000000, got ${state.amount}`); + } + }); + + // Test: cooldown full unstake closes the account (req #2) + await test(`[${tokenProgramLabel}] Cooldown: full unstake closes the account`, async () => { + const ctx = new TestContext(connection, Keypair.generate(), programAuthority, tokenProgramId); + await ctx.setup(); + await ctx.createMint(9); + await ctx.initializePool(BigInt(2592000)); + await ctx.updatePoolSettings(ctx.payer, null, null, BigInt(5)); + + const user = Keypair.generate(); + await airdropAndConfirm(connection, user.publicKey, LAMPORTS_PER_SOL); + const userToken = await ctx.createUserTokenAccount(user.publicKey); + await ctx.mintTokens(userToken, BigInt(1_000_000_000)); + await ctx.stake(user, userToken, BigInt(1_000_000_000)); + + await ctx.requestUnstake(user, BigInt(1_000_000_000)); + console.log(' Waiting 6s for cooldown...'); + await new Promise(r => setTimeout(r, 6000)); + await ctx.completeUnstake(user, userToken); + + const [userStakePDA] = deriveUserStakePDA(ctx.poolPDA, user.publicKey); + const closed = await connection.getAccountInfo(userStakePDA); + if (closed !== null) throw new Error('Stake account should be closed after full cooldown unstake'); + + const balance = await ctx.getTokenBalance(userToken); + if (balance !== BigInt(1_000_000_000)) throw new Error(`Expected 1000000000 tokens back, got ${balance}`); + }); + + // Test: cancel restores the frozen coins to the active position + await test(`[${tokenProgramLabel}] Cooldown: cancel restores stake to the pool`, async () => { + const ctx = new TestContext(connection, Keypair.generate(), programAuthority, tokenProgramId); + await ctx.setup(); + await ctx.createMint(9); + await ctx.initializePool(BigInt(2592000)); + await ctx.updatePoolSettings(ctx.payer, null, null, BigInt(60)); + + const user = Keypair.generate(); + await airdropAndConfirm(connection, user.publicKey, LAMPORTS_PER_SOL); + const userToken = await ctx.createUserTokenAccount(user.publicKey); + await ctx.mintTokens(userToken, BigInt(1_000_000_000)); + await ctx.stake(user, userToken, BigInt(1_000_000_000)); + + // Full request removes the stake from pool accounting + await ctx.requestUnstake(user, BigInt(1_000_000_000)); + const poolAfterReq = await ctx.readPoolState(); + if (poolAfterReq.totalStaked !== 0n) { + throw new Error(`Expected total_staked 0 after full request, got ${poolAfterReq.totalStaked}`); + } + + // Cancel restores it + await ctx.cancelUnstakeRequest(user); + const poolAfterCancel = await ctx.readPoolState(); + if (poolAfterCancel.totalStaked !== BigInt(1_000_000_000)) { + throw new Error(`Expected total_staked 1000000000 after cancel, got ${poolAfterCancel.totalStaked}`); + } + const state = await ctx.readUserStakeState(user.publicKey); + if (state.amount !== BigInt(1_000_000_000)) { + throw new Error(`Expected active amount 1000000000 after cancel, got ${state.amount}`); + } + }); + // Test: Existing pools (zero reserved fields) work unchanged await test(`[${tokenProgramLabel}] Backward compat: existing pool works with zero settings`, async () => { const ctx = new TestContext(connection, Keypair.generate(), programAuthority, tokenProgramId); @@ -3590,8 +3748,8 @@ async function runTests() { if (meta2.memberCount !== 1n) throw new Error(`Expected still 1, got ${meta2.memberCount}`); }); - // Test: CloseStakeAccount with metadata decrements member_count - await test(`[${tokenProgramLabel}] CloseStakeAccount with metadata decrements member_count`, async () => { + // Test: full Unstake auto-closes the account and decrements member_count + await test(`[${tokenProgramLabel}] Full unstake auto-closes account and decrements member_count`, async () => { const ctx = new TestContext(connection, Keypair.generate(), programAuthority, tokenProgramId); await ctx.setup(); await ctx.createMintWithMetadata(9, 'CloseTest', 'CLZ'); @@ -3608,13 +3766,19 @@ async function runTests() { const meta1 = await ctx.readMetadata(); if (meta1.memberCount !== 1n) throw new Error(`Expected 1, got ${meta1.memberCount}`); - // Unstake everything - await ctx.unstake(user, userToken, BigInt(1_000_000_000)); + // Full unstake (with metadata) auto-closes the account and decrements member_count + await ctx.unstake(user, userToken, BigInt(1_000_000_000), true); + + const [userStakePDA] = deriveUserStakePDA(ctx.poolPDA, user.publicKey); + const closed = await connection.getAccountInfo(userStakePDA); + if (closed !== null) throw new Error('Stake account should be closed after full unstake'); - // Close stake account with metadata - await ctx.closeStakeAccount(user, true); const meta2 = await ctx.readMetadata(); - if (meta2.memberCount !== 0n) throw new Error(`Expected 0 after close, got ${meta2.memberCount}`); + if (meta2.memberCount !== 0n) throw new Error(`Expected 0 after full unstake, got ${meta2.memberCount}`); + + // Tokens fully returned + const balance = await ctx.getTokenBalance(userToken); + if (balance !== BigInt(1_000_000_000)) throw new Error(`Expected 1000000000 tokens back, got ${balance}`); }); // Test: SetPoolMetadata preserves member_count on update From 5bd8c9d9eee9620c8802900bf86f9cb5cdf07d21 Mon Sep 17 00:00:00 2001 From: Mark Karpeles Date: Sat, 23 May 2026 23:47:14 +0900 Subject: [PATCH 2/6] Always close account on full unstake; pay all owed SOL, no residual MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per review: on a full unstake we pay out the user's rewards in full (the pool holds every staker's unclaimed rewards, so a shortfall should not occur) and close the account 100% of the time. Any theoretical unpayable remainder is redistributed to remaining stakers via last_synced_lamports rather than stranded on a closed account — so a fully-unstaked position never carries residual. - settle_unstake_accounting: full-unstake branch zeroes reward_debt and redistributes any remainder instead of storing residual / total_residual_unpaid - execute_unstake / complete_unstake: close unconditionally when amount == 0 Fix E2E tests for the auto-close behavior (close refunds the user's own stake-account rent, which is not reward SOL): - add rewardExcludingRent() helper that subtracts refunded rent on close - conservation + multi-phase reconciliation tests use it for full unstakes - "Cannot claim after full unstake": assert the account is closed and a follow-up claim fails, instead of expecting a zero-value claim Co-Authored-By: Claude Opus 4.7 (1M context) --- .../src/instructions/complete_unstake.rs | 11 ++- .../chiefstaker/src/instructions/unstake.rs | 39 ++++----- tests/typescript/test_staking.ts | 82 ++++++++++++------- 3 files changed, 74 insertions(+), 58 deletions(-) diff --git a/programs/chiefstaker/src/instructions/complete_unstake.rs b/programs/chiefstaker/src/instructions/complete_unstake.rs index 9d6c12e..2ff648f 100644 --- a/programs/chiefstaker/src/instructions/complete_unstake.rs +++ b/programs/chiefstaker/src/instructions/complete_unstake.rs @@ -14,7 +14,6 @@ use spl_token_2022::extension::StateWithExtensions; use crate::{ error::StakingError, - math::WAD, state::{is_valid_token_program, StakingPool, UserStake, POOL_SEED}, }; @@ -24,8 +23,8 @@ use super::unstake::close_user_stake_account; /// /// The reward settlement and pool accounting already happened at RequestUnstake; /// this instruction only delivers the frozen tokens from the vault to the user -/// and, when the position is fully unstaked with nothing owed, closes the stake -/// account to reclaim its rent. +/// and, when the position is fully unstaked, closes the stake account to reclaim +/// its rent. /// /// Accounts (same as Unstake): /// 0. `[writable]` Pool account @@ -135,9 +134,9 @@ pub fn process_complete_unstake( user_stake.unstake_request_amount = 0; user_stake.unstake_request_time = 0; - // Close iff fully unstaked and nothing owed; otherwise keep the account so the - // remaining position keeps earning / residual rewards stay claimable. - let should_close = user_stake.amount == 0 && user_stake.reward_debt / WAD == 0; + // A full unstake (no active stake remaining) closes the account; a partial + // completion leaves the still-active position open and earning. + let should_close = user_stake.amount == 0; // Persist the cleared request fields before the CPI { diff --git a/programs/chiefstaker/src/instructions/unstake.rs b/programs/chiefstaker/src/instructions/unstake.rs index 69545d1..dbd1d44 100644 --- a/programs/chiefstaker/src/instructions/unstake.rs +++ b/programs/chiefstaker/src/instructions/unstake.rs @@ -180,25 +180,24 @@ pub fn settle_unstake_accounting<'a>( .checked_add(user_stake.reward_debt) .ok_or(StakingError::MathOverflow)?; } else { - // Full unstake: preserve any unpaid rewards in reward_debt so the user - // can claim them later via the amount==0 claim path. When amount==0, - // reward_debt is reinterpreted as "unclaimed WAD-scaled rewards". - user_stake.reward_debt = unpaid_rewards_wad; + // Full unstake: the position is fully reset and the account will be closed. + // The user's pending rewards were paid above (capped by the pool balance). + // Any unpayable remainder should not occur — the pool holds every staker's + // unclaimed rewards — but if it does, we redistribute it to the remaining + // stakers via last_synced_lamports rather than stranding it on a closed + // account or owing it to a user who is gone. + let remainder_lamports = (unpaid_rewards_wad / WAD) as u64; + if remainder_lamports > 0 { + pool.last_synced_lamports = pool.last_synced_lamports.saturating_sub(remainder_lamports); + } + + user_stake.reward_debt = 0; user_stake.claimed_rewards_wad = 0; - // Remove old debt from total_reward_debt but do NOT add the residual. - // Residual debts are tracked separately in total_residual_unpaid because - // the user's amount is 0 (no allocation in total_staked * acc_rps), and - // including them in total_reward_debt would break FixTotalRewardDebt. + // Remove the user's old debt from the pool aggregate; nothing is owed. pool.total_reward_debt = pool .total_reward_debt .saturating_sub(old_reward_debt); - - let residual_lamports = (unpaid_rewards_wad / WAD) as u64; - pool.total_residual_unpaid = pool - .total_residual_unpaid - .checked_add(residual_lamports) - .ok_or(StakingError::MathOverflow)?; } // Increment cumulative rewards counter @@ -217,7 +216,7 @@ pub fn settle_unstake_accounting<'a>( /// the pool member_count if a metadata account is supplied. /// /// Mirrors `process_close_stake_account`. The caller must ensure the position is -/// fully unstaked with no residual rewards owed before calling. +/// fully unstaked before calling. pub fn close_user_stake_account<'a>( program_id: &Pubkey, pool_info: &AccountInfo<'a>, @@ -290,11 +289,9 @@ pub fn execute_unstake<'a>( system_program_info, )?; - // A full unstake fully resets the position. Close the account afterward when - // nothing is owed; if residual rewards remain (pool underfunded), keep it open - // so the user can claim them, then close. + // A full unstake fully resets the position; the account is always closed + // afterward (the user's rewards were settled and paid in full above). let fully_unstaked = user_stake.amount == 0; - let has_residual = user_stake.reward_debt / WAD > 0; // Save states (before CPI — pool data includes pre-updated last_synced_lamports) { @@ -345,8 +342,8 @@ pub fn execute_unstake<'a>( msg!("Unstaked {} tokens", amount); - // Close the account when fully unstaked and nothing is owed. - if fully_unstaked && !has_residual { + // Full unstake closes the account to reclaim its rent. + if fully_unstaked { close_user_stake_account(program_id, pool_info, user_stake_info, user_info, metadata_info)?; } diff --git a/tests/typescript/test_staking.ts b/tests/typescript/test_staking.ts index 197d98c..4ac59f3 100644 --- a/tests/typescript/test_staking.ts +++ b/tests/typescript/test_staking.ts @@ -99,6 +99,27 @@ function deriveMetadataPDA(pool: PublicKey): [PublicKey, number] { ); } +// Reward SOL from an operation, excluding any stake-account rent refunded when a +// full unstake auto-closes the account. That rent is the user's own deposit +// coming back (paid at stake time), not reward SOL, so it must not be counted in +// conservation accounting. +async function rewardExcludingRent( + connection: Connection, + poolPDA: PublicKey, + user: Keypair, + op: () => Promise, +): Promise { + const [stakePDA] = deriveUserStakePDA(poolPDA, user.publicKey); + const infoBefore = await connection.getAccountInfo(stakePDA); + const rentRefundable = infoBefore ? BigInt(infoBefore.lamports) : 0n; + const before = BigInt(await connection.getBalance(user.publicKey)); + await op(); + const after = BigInt(await connection.getBalance(user.publicKey)); + const infoAfter = await connection.getAccountInfo(stakePDA); + const closed = infoBefore !== null && infoAfter === null; + return (after - before) - (closed ? rentRefundable : 0n); +} + async function airdropAndConfirm(connection: Connection, publicKey: PublicKey, lamports: number): Promise { const sig = await connection.requestAirdrop(publicKey, lamports); await connection.confirmTransaction(sig); @@ -2389,22 +2410,26 @@ async function runTests() { // Claim rewards first await ctx.claimRewards(user); - // Fully unstake + // Fully unstake — this auto-closes the stake account await ctx.unstake(user, userToken, BigInt(1_000_000_000)); - // Deposit more rewards (user has amount=0, should not benefit) - await ctx.depositRewards(BigInt(LAMPORTS_PER_SOL)); + const [userStakePDA] = deriveUserStakePDA(ctx.poolPDA, user.publicKey); + if ((await connection.getAccountInfo(userStakePDA)) !== null) { + throw new Error('Stake account should be closed after full unstake'); + } - // Claim again — should succeed silently (amount==0 path returns Ok) - // but should give the user 0 SOL from the new deposit - const balBefore = BigInt(await ctx.getBalance(user.publicKey)); - await ctx.claimRewards(user); - const balAfter = BigInt(await ctx.getBalance(user.publicKey)); + // Deposit more rewards — the exited user must not benefit from them + await ctx.depositRewards(BigInt(LAMPORTS_PER_SOL)); - // User should gain nothing (or lose tx fee) - const gained = balAfter - balBefore; - if (gained > BigInt(0)) { - throw new Error(`User gained ${gained} lamports after full unstake — should be 0`); + // Claiming now must fail: there is no stake account left to claim from + let claimFailed = false; + try { + await ctx.claimRewards(user); + } catch (e) { + claimFailed = true; + } + if (!claimFailed) { + throw new Error('Claim after full unstake should fail (account closed)'); } }); @@ -3140,13 +3165,11 @@ async function runTests() { let totalDeposited = BigInt(0); const rewards: Record = { alice: BigInt(0), bob: BigInt(0), carol: BigInt(0), dave: BigInt(0) }; - // Helper: measure SOL reward from an operation + // Helper: measure SOL reward from an operation. Excludes any stake-account + // rent refunded when a full unstake auto-closes the account (that is the + // user's own rent returning, not a reward). Payer covers tx fees. async function measureReward(name: string, user: Keypair, op: () => Promise): Promise { - const before = BigInt(await ctx.getBalance(user.publicKey)); - await op(); - const after = BigInt(await ctx.getBalance(user.publicKey)); - // Payer covers tx fee, so balance diff = pure reward - const reward = after - before; + const reward = await rewardExcludingRent(connection, ctx.poolPDA, user, op); if (reward > BigInt(0)) { rewards[name] += reward; } @@ -3366,10 +3389,9 @@ async function runTests() { let balAfter = BigInt(await ctx.getBalance(staker1.publicKey)); totalClaimed += (balAfter - bal); - bal = BigInt(await ctx.getBalance(staker1.publicKey)); - await ctx.unstake(staker1, token1, BigInt(1_500_000_000)); - balAfter = BigInt(await ctx.getBalance(staker1.publicKey)); - totalClaimed += (balAfter - bal); // reward portion from unstake + // Full unstake auto-closes the account; exclude refunded rent from rewards + totalClaimed += await rewardExcludingRent(connection, ctx.poolPDA, staker1, () => + ctx.unstake(staker1, token1, BigInt(1_500_000_000))); // Staker2: claim + unstake bal = BigInt(await ctx.getBalance(staker2.publicKey)); @@ -3377,10 +3399,9 @@ async function runTests() { balAfter = BigInt(await ctx.getBalance(staker2.publicKey)); totalClaimed += (balAfter - bal); - bal = BigInt(await ctx.getBalance(staker2.publicKey)); - await ctx.unstake(staker2, token2, BigInt(1_000_000_000)); - balAfter = BigInt(await ctx.getBalance(staker2.publicKey)); - totalClaimed += (balAfter - bal); + // Full unstake auto-closes the account; exclude refunded rent from rewards + totalClaimed += await rewardExcludingRent(connection, ctx.poolPDA, staker2, () => + ctx.unstake(staker2, token2, BigInt(1_000_000_000))); // Check pool remaining const poolBalance = BigInt(await ctx.getBalance(ctx.poolPDA)); @@ -3556,11 +3577,10 @@ async function runTests() { const stakeAmount = i === 0 ? BigInt(1_500_000_000) // 1B + 500M : BigInt((i + 1) * 1_000_000_000); - const balBeforeUnstake = BigInt(await ctx.getBalance(stakers[i].publicKey)); - await ctx.unstake(stakers[i], tokens[i], stakeAmount); - const balAfterUnstake = BigInt(await ctx.getBalance(stakers[i].publicKey)); - // Count only the reward portion of unstake (unstake also claims rewards) - totalClaimed += (balAfterUnstake - balBeforeUnstake); + // Full unstake auto-closes the account; count only the reward portion + // (exclude the refunded stake-account rent). + totalClaimed += await rewardExcludingRent(connection, ctx.poolPDA, stakers[i], () => + ctx.unstake(stakers[i], tokens[i], stakeAmount)); } const poolBalance = BigInt(await ctx.getBalance(ctx.poolPDA)); From 1ee512c881cf84cb77b0751861d76e7b479c2d87 Mon Sep 17 00:00:00 2001 From: Mark Karpeles Date: Sat, 23 May 2026 23:54:08 +0900 Subject: [PATCH 3/6] Don't restore/double-remove legacy in-flight unstake requests on upgrade MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A request created by the pre-upgrade code never removed the coins from the pool (old RequestUnstake only set the request fields). The new Cancel/Complete would otherwise add those coins back (Cancel) or under-remove them (Complete), crediting the user stake they never lost. Add a 1-byte UserStake.unstake_request_settled marker, set to 1 only by the new RequestUnstake. Legacy accounts deserialize it as 0 (relying on realloc zero-fill, same as claimed_rewards_wad). Branch on it: - Cancel: settled==1 restores coins to the pool; settled==0 (legacy) just clears the request — nothing to restore. - Complete: settled==1 delivers tokens + closes; settled==0 (legacy) runs the full execute_unstake (settle + remove from pool + transfer + close), matching the pre-upgrade flow. UserStake::LEN 177 -> 178; LEGACY_LEN stays 161. Updated new(), BorshDeserialize, size/roundtrip tests (incl. legacy classifies as settled==0), and idl.json. Co-Authored-By: Claude Opus 4.7 (1M context) --- idl.json | 7 +- .../src/instructions/cancel_unstake.rs | 106 ++++++++++-------- .../src/instructions/complete_unstake.rs | 61 +++++++--- .../src/instructions/request_unstake.rs | 6 +- programs/chiefstaker/src/state.rs | 38 ++++++- 5 files changed, 149 insertions(+), 69 deletions(-) diff --git a/idl.json b/idl.json index 2f2e565..ee14d7b 100644 --- a/idl.json +++ b/idl.json @@ -879,7 +879,7 @@ "docs": [ "User stake account.", "PDA seeds: [\"stake\", pool, owner].", - "Size: 177 bytes (legacy accounts at 153 or 161 bytes are lazily reallocated)." + "Size: 178 bytes (legacy accounts at 153, 161, or 177 bytes are lazily reallocated)." ], "type": { "kind": "struct", @@ -955,6 +955,11 @@ "name": "claimedRewardsWad", "type": "u128", "docs": ["Cumulative WAD-scaled rewards already paid for the current position. Used for frequency-independent claims: pending = full_entitlement - claimedRewardsWad. Reset to 0 on stake/unstake. Defaults to 0 for legacy accounts."] + }, + { + "name": "unstakeRequestSettled", + "type": "u8", + "docs": ["1 when the pending unstake request was created under the settle-at-request flow (frozen coins already removed from the pool). 0 = no request, or a legacy in-flight request from before this upgrade. Cancel/Complete branch on this to avoid double-restoring or double-removing legacy requests. Defaults to 0 for legacy accounts."] } ] } diff --git a/programs/chiefstaker/src/instructions/cancel_unstake.rs b/programs/chiefstaker/src/instructions/cancel_unstake.rs index f202718..05e5384 100644 --- a/programs/chiefstaker/src/instructions/cancel_unstake.rs +++ b/programs/chiefstaker/src/instructions/cancel_unstake.rs @@ -102,62 +102,72 @@ pub fn process_cancel_unstake_request( let frozen = user_stake.unstake_request_amount; - // Edge case: a full unstake request (amount == 0) that still has unpaid - // residual rewards stores those in reward_debt. Reactivating the position - // would overwrite that storage with a snapshot, losing the residual. Require - // the user to claim it first (ClaimRewards works on the amount==0 path). - if user_stake.amount == 0 && user_stake.reward_debt / WAD > 0 { - return Err(StakingError::ResidualRewardsPending.into()); - } + if user_stake.unstake_request_settled == 1 { + // New-style request: RequestUnstake removed the frozen coins from the pool + // and settled their rewards. Restore them to the active position. + + // Edge case: a full unstake request (amount == 0) that still has unpaid + // residual rewards stores those in reward_debt. Reactivating the position + // would overwrite that storage with a snapshot, losing the residual. + // Require the user to claim it first (ClaimRewards works on amount==0). + if user_stake.amount == 0 && user_stake.reward_debt / WAD > 0 { + return Err(StakingError::ResidualRewardsPending.into()); + } + + let frozen_wad = (frozen as u128) + .checked_mul(WAD) + .ok_or(StakingError::MathOverflow)?; - // Restore the frozen coins to the active position. - let frozen_wad = (frozen as u128) - .checked_mul(WAD) - .ok_or(StakingError::MathOverflow)?; - - // sum_stake_exp += frozen * exp_start_factor (restores weight / maturity) - let contribution = wad_mul(frozen_wad, user_stake.exp_start_factor)?; - let new_sum = pool - .get_sum_stake_exp() - .checked_add(U256::from_u128(contribution)) - .ok_or(StakingError::MathOverflow)?; - pool.set_sum_stake_exp(new_sum); - - // total_staked += frozen - pool.total_staked = pool - .total_staked - .checked_add(frozen as u128) - .ok_or(StakingError::MathOverflow)?; - - // Fresh reward snapshot for the re-added tokens (no retroactive rewards) - let new_token_debt = wad_mul(frozen_wad, pool.acc_reward_per_weighted_share)?; - - if user_stake.amount == 0 { - // Was a full request (no residual remains): start clean. - user_stake.reward_debt = new_token_debt; - user_stake.claimed_rewards_wad = 0; - } else { - // Partial request: keep the active remainder's pending intact and add a - // fresh debt snapshot for the re-added tokens (mirrors add-stake). - user_stake.reward_debt = user_stake - .reward_debt - .checked_add(new_token_debt) + // sum_stake_exp += frozen * exp_start_factor (restores weight / maturity) + let contribution = wad_mul(frozen_wad, user_stake.exp_start_factor)?; + let new_sum = pool + .get_sum_stake_exp() + .checked_add(U256::from_u128(contribution)) + .ok_or(StakingError::MathOverflow)?; + pool.set_sum_stake_exp(new_sum); + + // total_staked += frozen + pool.total_staked = pool + .total_staked + .checked_add(frozen as u128) .ok_or(StakingError::MathOverflow)?; - } - user_stake.amount = user_stake - .amount - .checked_add(frozen) - .ok_or(StakingError::MathOverflow)?; + // Fresh reward snapshot for the re-added tokens (no retroactive rewards) + let new_token_debt = wad_mul(frozen_wad, pool.acc_reward_per_weighted_share)?; + + if user_stake.amount == 0 { + // Was a full request (no residual remains): start clean. + user_stake.reward_debt = new_token_debt; + user_stake.claimed_rewards_wad = 0; + } else { + // Partial request: keep the active remainder's pending intact and add a + // fresh debt snapshot for the re-added tokens (mirrors add-stake). + user_stake.reward_debt = user_stake + .reward_debt + .checked_add(new_token_debt) + .ok_or(StakingError::MathOverflow)?; + } + + user_stake.amount = user_stake + .amount + .checked_add(frozen) + .ok_or(StakingError::MathOverflow)?; - pool.total_reward_debt = pool - .total_reward_debt - .checked_add(new_token_debt) - .ok_or(StakingError::MathOverflow)?; + pool.total_reward_debt = pool + .total_reward_debt + .checked_add(new_token_debt) + .ok_or(StakingError::MathOverflow)?; + } else { + // Legacy in-flight request (created before this upgrade): the coins were + // never removed from the pool's accounting, so there is nothing to restore + // — adding them back here would credit the user stake they never lost. + // Just clear the request, matching the pre-upgrade cancel behavior. + } // Clear the request fields user_stake.unstake_request_amount = 0; user_stake.unstake_request_time = 0; + user_stake.unstake_request_settled = 0; // Save pool and user stake { diff --git a/programs/chiefstaker/src/instructions/complete_unstake.rs b/programs/chiefstaker/src/instructions/complete_unstake.rs index 2ff648f..70a9f07 100644 --- a/programs/chiefstaker/src/instructions/complete_unstake.rs +++ b/programs/chiefstaker/src/instructions/complete_unstake.rs @@ -17,14 +17,19 @@ use crate::{ state::{is_valid_token_program, StakingPool, UserStake, POOL_SEED}, }; -use super::unstake::close_user_stake_account; +use super::unstake::{close_user_stake_account, execute_unstake}; /// Complete unstake after the cooldown has elapsed. /// -/// The reward settlement and pool accounting already happened at RequestUnstake; -/// this instruction only delivers the frozen tokens from the vault to the user -/// and, when the position is fully unstaked, closes the stake account to reclaim -/// its rent. +/// For requests created under the current flow (`unstake_request_settled == 1`) +/// the reward settlement and pool accounting already happened at RequestUnstake, +/// so this only delivers the frozen tokens and closes the account on a full +/// unstake. +/// +/// For a legacy in-flight request created before this upgrade +/// (`unstake_request_settled == 0`), the coins are still counted in the pool, so +/// this runs the full unstake (settle + remove from pool + transfer + close), +/// exactly as the pre-upgrade complete path did. /// /// Accounts (same as Unstake): /// 0. `[writable]` Pool account @@ -34,7 +39,7 @@ use super::unstake::close_user_stake_account; /// 4. `[]` Token mint /// 5. `[writable, signer]` User/owner /// 6. `[]` Token 2022 program -/// 7. `[]` System program (optional, accepted for call-site symmetry) +/// 7. `[]` System program (optional) /// 8. `[writable]` Metadata PDA (optional, to decrement member_count on close) pub fn process_complete_unstake( program_id: &Pubkey, @@ -64,7 +69,7 @@ pub fn process_complete_unstake( if pool_info.owner != program_id { return Err(StakingError::InvalidAccountOwner.into()); } - let pool = StakingPool::try_from_slice(&pool_info.try_borrow_data()?)?; + let mut pool = StakingPool::try_from_slice(&pool_info.try_borrow_data()?)?; if !pool.is_initialized() { return Err(StakingError::NotInitialized.into()); } @@ -127,12 +132,46 @@ pub fn process_complete_unstake( return Err(StakingError::CooldownNotElapsed.into()); } - // The frozen tokens awaiting withdrawal (already settled at request time) + // Lazily adjust exp_start_factor if pool has been rebased (needed by the + // legacy path's accounting; harmless for the new path). + user_stake.sync_to_pool(&pool)?; + let withdraw_amount = user_stake.unstake_request_amount; - // Clear the request fields + // Optional trailing accounts: system program (legacy realloc) then metadata + let system_program_info = account_info_iter.next(); + let metadata_info = account_info_iter.next(); + + if user_stake.unstake_request_settled == 0 { + // ── Legacy in-flight request ────────────────────────────────────── + // The frozen coins were never removed from the pool. Run the full unstake + // accounting + transfer + close, exactly like the pre-upgrade flow. + user_stake.unstake_request_amount = 0; + user_stake.unstake_request_time = 0; + + return execute_unstake( + program_id, + &mut pool, + &mut user_stake, + pool_info, + user_stake_info, + token_vault_info, + user_token_info, + mint_info, + user_info, + withdraw_amount, + current_time, + system_program_info, + metadata_info, + ); + } + + // ── New-style request ───────────────────────────────────────────────── + // Accounting already happened at RequestUnstake. Just deliver the frozen + // tokens and close the account on a full unstake. user_stake.unstake_request_amount = 0; user_stake.unstake_request_time = 0; + user_stake.unstake_request_settled = 0; // A full unstake (no active stake remaining) closes the account; a partial // completion leaves the still-active position open and earning. @@ -174,10 +213,6 @@ pub fn process_complete_unstake( msg!("Completed unstake of {} tokens", withdraw_amount); - // Optional trailing accounts: system program (unused here) then metadata (close) - let _system_program_info = account_info_iter.next(); - let metadata_info = account_info_iter.next(); - // 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)?; diff --git a/programs/chiefstaker/src/instructions/request_unstake.rs b/programs/chiefstaker/src/instructions/request_unstake.rs index b16f89c..4de6672 100644 --- a/programs/chiefstaker/src/instructions/request_unstake.rs +++ b/programs/chiefstaker/src/instructions/request_unstake.rs @@ -146,9 +146,13 @@ pub fn process_request_unstake( system_program_info, )?; - // Record the pending request (the frozen amount awaiting withdrawal) + // Record the pending request (the frozen amount awaiting withdrawal). + // Mark it as settled-at-request so Cancel/Complete know the coins were already + // removed from the pool here (distinguishes new requests from legacy in-flight + // requests created before this upgrade). user_stake.unstake_request_amount = amount; user_stake.unstake_request_time = current_time; + user_stake.unstake_request_settled = 1; // Save pool and user stake { diff --git a/programs/chiefstaker/src/state.rs b/programs/chiefstaker/src/state.rs index 3c8637e..67d7fbe 100644 --- a/programs/chiefstaker/src/state.rs +++ b/programs/chiefstaker/src/state.rs @@ -297,6 +297,15 @@ pub struct UserStake { /// (partial/full) when the position is restructured and pending is settled. /// Defaults to 0 for existing accounts (correct: first claim gets full pending). pub claimed_rewards_wad: u128, + + /// 1 when the current pending unstake request was created under the + /// "settle at request" flow, meaning the requested coins were already removed + /// from the pool's reward accounting (total_staked / sum_stake_exp) at request + /// time. 0 means no pending request, OR a legacy in-flight request created + /// before this upgrade (whose coins are still counted in the pool). + /// CancelUnstake / CompleteUnstake use this to avoid double-restoring or + /// double-removing a legacy request. Defaults to 0 for existing accounts. + pub unstake_request_settled: u8, } impl UserStake { @@ -314,10 +323,12 @@ impl UserStake { 8 + // last_stake_time 8 + // base_time_snapshot 8 + // total_rewards_claimed - 16; // claimed_rewards_wad + 16 + // claimed_rewards_wad + 1; // unstake_request_settled - /// Legacy account size (before claimed_rewards_wad was added) - pub const LEGACY_LEN: usize = Self::LEN - 16; + /// Legacy account size (before claimed_rewards_wad and unstake_request_settled + /// were added). Equals 161 bytes. + pub const LEGACY_LEN: usize = Self::LEN - 17; /// Create a new user stake pub fn new( @@ -344,6 +355,7 @@ impl UserStake { base_time_snapshot, total_rewards_claimed: 0, claimed_rewards_wad: 0, + unstake_request_settled: 0, } } @@ -426,6 +438,7 @@ impl BorshDeserialize for UserStake { // New fields — may not be present in legacy accounts let total_rewards_claimed = u64::deserialize_reader(reader).unwrap_or(0); let claimed_rewards_wad = u128::deserialize_reader(reader).unwrap_or(0); + let unstake_request_settled = u8::deserialize_reader(reader).unwrap_or(0); Ok(Self { discriminator, @@ -442,6 +455,7 @@ impl BorshDeserialize for UserStake { base_time_snapshot, total_rewards_claimed, claimed_rewards_wad, + unstake_request_settled, }) } } @@ -599,7 +613,7 @@ mod tests { ); let serialized = borsh::to_vec(&stake).unwrap(); assert_eq!(serialized.len(), UserStake::LEN); - assert_eq!(UserStake::LEN, 177); + assert_eq!(UserStake::LEN, 178); assert_eq!(UserStake::LEGACY_LEN, 161); } @@ -620,11 +634,14 @@ mod tests { // Truncate to legacy 161 bytes (no claimed_rewards_wad) let legacy = &full[..UserStake::LEGACY_LEN]; - // Deserialize should succeed with claimed_rewards_wad defaulting to 0 + // Deserialize should succeed with new fields defaulting to 0. + // Critically, a legacy in-flight unstake request must classify as + // unstake_request_settled == 0 so cancel/complete use the legacy path. let deserialized = UserStake::try_from_slice(legacy).unwrap(); assert_eq!(deserialized.amount, 1000); assert_eq!(deserialized.total_rewards_claimed, 0); assert_eq!(deserialized.claimed_rewards_wad, 0); + assert_eq!(deserialized.unstake_request_settled, 0); assert_eq!(deserialized.bump, 255); // Very old 153-byte accounts (no total_rewards_claimed or claimed_rewards_wad) @@ -633,11 +650,18 @@ mod tests { assert_eq!(deserialized_old.amount, 1000); assert_eq!(deserialized_old.total_rewards_claimed, 0); assert_eq!(deserialized_old.claimed_rewards_wad, 0); + assert_eq!(deserialized_old.unstake_request_settled, 0); + + // 177-byte accounts (pre-marker, with claimed_rewards_wad) classify as legacy + let pre_marker = &full[..177]; + let deserialized_pre = UserStake::try_from_slice(pre_marker).unwrap(); + assert_eq!(deserialized_pre.unstake_request_settled, 0); - // Full 177-byte deserialization should also work + // Full 178-byte deserialization should also work let deserialized_full = UserStake::try_from_slice(&full).unwrap(); assert_eq!(deserialized_full.total_rewards_claimed, 0); assert_eq!(deserialized_full.claimed_rewards_wad, 0); + assert_eq!(deserialized_full.unstake_request_settled, 0); } #[test] @@ -653,10 +677,12 @@ mod tests { ); stake.total_rewards_claimed = 999_999; stake.claimed_rewards_wad = 42_000_000_000_000_000_000; + stake.unstake_request_settled = 1; let serialized = borsh::to_vec(&stake).unwrap(); let deserialized = UserStake::try_from_slice(&serialized).unwrap(); assert_eq!(deserialized.total_rewards_claimed, 999_999); assert_eq!(deserialized.claimed_rewards_wad, 42_000_000_000_000_000_000); + assert_eq!(deserialized.unstake_request_settled, 1); } #[test] From 8d3e27c3f417c9063d362a8e7c7dac83422b27af Mon Sep 17 00:00:00 2001 From: Mark Karpeles Date: Sun, 24 May 2026 00:15:17 +0900 Subject: [PATCH 4/6] Add integration tests for legacy in-flight unstake requests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ProgramTest/BanksClient tests that craft the exact pre-upgrade on-chain state (177-byte UserStake with a pending request, coins still counted in the pool) and verify the upgraded program handles it: - legacy_cancel_does_not_credit_stake: CancelUnstakeRequest on a legacy request must restore nothing — total_staked and the user's amount stay unchanged. - legacy_complete_full_removes_and_closes: CompleteUnstake routes legacy requests through the full execute_unstake (real Token-2022 transfer CPI), draining the vault, zeroing total_staked, and closing the account. These are deterministic (no validator/old binary needed). Wire `cargo test` into CI, which previously only ran build-sbf + the TS E2E suite (so the 33 unit tests now run in CI too). Adds spl-token-2022 as a dev-dependency for the token CPI. Co-Authored-By: Claude Opus 4.7 (1M context) --- .github/workflows/verifiable-build.yml | 3 + programs/chiefstaker/Cargo.toml | 1 + .../chiefstaker/tests/legacy_migration.rs | 288 ++++++++++++++++++ 3 files changed, 292 insertions(+) create mode 100644 programs/chiefstaker/tests/legacy_migration.rs diff --git a/.github/workflows/verifiable-build.yml b/.github/workflows/verifiable-build.yml index 3391661..0aa1f5f 100644 --- a/.github/workflows/verifiable-build.yml +++ b/.github/workflows/verifiable-build.yml @@ -25,6 +25,9 @@ jobs: - name: Build program run: cargo build-sbf --workspace + - name: Run Rust tests (unit + legacy-migration integration) + run: cargo test -p chiefstaker + - name: Start test validator with upgradeable program run: | solana config set --url http://localhost:8899 --keypair ~/.config/solana/id.json diff --git a/programs/chiefstaker/Cargo.toml b/programs/chiefstaker/Cargo.toml index c6ae6df..5abc4b5 100644 --- a/programs/chiefstaker/Cargo.toml +++ b/programs/chiefstaker/Cargo.toml @@ -28,3 +28,4 @@ solana-security-txt = "1.1" solana-program-test = "2.0" solana-sdk = "2.0" tokio = { version = "1", features = ["full"] } +spl-token-2022 = { version = "5.0", features = ["no-entrypoint"] } diff --git a/programs/chiefstaker/tests/legacy_migration.rs b/programs/chiefstaker/tests/legacy_migration.rs new file mode 100644 index 0000000..e029466 --- /dev/null +++ b/programs/chiefstaker/tests/legacy_migration.rs @@ -0,0 +1,288 @@ +//! Integration tests for legacy in-flight unstake requests created BEFORE the +//! "settle at request" upgrade. +//! +//! Under the old code, `RequestUnstake` only set the request fields — it did NOT +//! remove the coins from the pool's accounting (`total_staked` / `sum_stake_exp`) +//! and the account had no `unstake_request_settled` byte (177 bytes). These tests +//! craft exactly that on-chain state and verify the upgraded program handles it: +//! * Cancel must NOT credit the user any stake (the coins never left the pool). +//! * Complete must run the full unstake (remove from pool + transfer + close). + +use borsh::BorshDeserialize; +use chiefstaker::{ + math::{wad_mul, U256, WAD}, + state::{StakingPool, UserStake, POOL_SEED, STAKE_SEED, TOKEN_VAULT_SEED}, + StakingInstruction, +}; +use solana_program::{ + instruction::{AccountMeta, Instruction}, + program_option::COption, + program_pack::Pack, + pubkey::Pubkey, + system_program, +}; +use solana_program_test::*; +use solana_sdk::{ + account::Account, + rent::Rent, + signature::{Keypair, Signer}, + transaction::Transaction, +}; + +fn rent_for(len: usize) -> u64 { + Rent::default().minimum_balance(len) +} + +/// Build a pool whose accounting still counts `staked` tokens (as the old code +/// left it during a pending request), staked at base_time so each token has full +/// exp_start_factor (= WAD). +fn make_pool(mint: Pubkey, vault: Pubkey, bump: u8, staked: u64, cooldown: u64) -> Vec { + let mut pool = StakingPool::new(mint, vault, Pubkey::default(), Pubkey::new_unique(), 60, 0, bump); + pool.total_staked = staked as u128; + // exp_start_factor = WAD for every token (staked at base_time), so + // sum_stake_exp = staked * WAD. + let contribution = wad_mul((staked as u128) * WAD, WAD).unwrap(); + pool.set_sum_stake_exp(U256::from_u128(contribution)); + pool.unstake_cooldown_seconds = cooldown; + borsh::to_vec(&pool).unwrap() +} + +/// Build a LEGACY (177-byte, pre-marker) UserStake with a pending request. +/// The serialized form is truncated to drop the `unstake_request_settled` byte, +/// exactly as an account written by the old program would appear. +fn make_legacy_stake(owner: Pubkey, pool: Pubkey, bump: u8, amount: u64, request: u64) -> Vec { + let mut us = UserStake::new(owner, pool, amount, 0, WAD, bump, 0); + us.unstake_request_amount = request; + us.unstake_request_time = 0; + let mut data = borsh::to_vec(&us).unwrap(); // 178 bytes (current layout) + data.truncate(UserStake::LEN - 1); // 177 bytes: legacy, no settled byte + data +} + +fn pack_mint(decimals: u8) -> Vec { + let mut data = vec![0u8; spl_token_2022::state::Mint::LEN]; + spl_token_2022::state::Mint { + mint_authority: COption::Some(Pubkey::new_unique()), + supply: 1_000_000, + decimals, + is_initialized: true, + freeze_authority: COption::None, + } + .pack_into_slice(&mut data); + data +} + +fn pack_token_account(mint: Pubkey, owner: Pubkey, amount: u64) -> Vec { + let mut data = vec![0u8; spl_token_2022::state::Account::LEN]; + spl_token_2022::state::Account { + mint, + owner, + amount, + delegate: COption::None, + state: spl_token_2022::state::AccountState::Initialized, + is_native: COption::None, + delegated_amount: 0, + close_authority: COption::None, + } + .pack_into_slice(&mut data); + data +} + +/// CancelUnstakeRequest on a legacy in-flight request must restore NOTHING to the +/// pool — the coins were never removed — so total_staked and the user's amount are +/// unchanged. (Regression guard: the new restore path must not run for legacy.) +#[tokio::test] +async fn legacy_cancel_does_not_credit_stake() { + let program_id = chiefstaker::id(); + let mint = Pubkey::new_unique(); + let (pool_pda, pool_bump) = 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 user = Keypair::new(); + let (stake_pda, stake_bump) = + Pubkey::find_program_address(&[STAKE_SEED, pool_pda.as_ref(), user.pubkey().as_ref()], &program_id); + + let staked: u64 = 1_000_000_000; + let request: u64 = 400_000_000; + + let mut pt = ProgramTest::new("chiefstaker", program_id, processor!(chiefstaker::process_instruction)); + pt.add_account( + pool_pda, + Account { + lamports: rent_for(StakingPool::LEN), + data: make_pool(mint, vault_pda, pool_bump, staked, 60), + owner: program_id, + executable: false, + rent_epoch: 0, + }, + ); + pt.add_account( + stake_pda, + Account { + lamports: rent_for(UserStake::LEN), // enough to cover realloc to 178 + data: make_legacy_stake(user.pubkey(), pool_pda, stake_bump, staked, request), + owner: program_id, + executable: false, + rent_epoch: 0, + }, + ); + pt.add_account( + user.pubkey(), + Account { + lamports: 10_000_000_000, + 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(stake_pda, false), + AccountMeta::new(user.pubkey(), true), + AccountMeta::new_readonly(system_program::id(), false), + ], + data: borsh::to_vec(&StakingInstruction::CancelUnstakeRequest).unwrap(), + }; + let tx = Transaction::new_signed_with_payer(&[ix], Some(&payer.pubkey()), &[&payer, &user], blockhash); + banks.process_transaction(tx).await.expect("legacy cancel should succeed"); + + // Pool accounting unchanged: total_staked must NOT have grown. + let pool_acc = banks.get_account(pool_pda).await.unwrap().unwrap(); + let pool = StakingPool::try_from_slice(&pool_acc.data).unwrap(); + assert_eq!(pool.total_staked, staked as u128, "legacy cancel must not credit stake to the pool"); + + // User stake: request cleared, amount unchanged, still marked legacy (0). + let stake_acc = banks.get_account(stake_pda).await.unwrap().unwrap(); + let us = UserStake::try_from_slice(&stake_acc.data).unwrap(); + assert_eq!(us.unstake_request_amount, 0, "request must be cleared"); + assert_eq!(us.amount, staked, "active amount must be unchanged"); + assert_eq!(us.unstake_request_settled, 0); +} + +/// CompleteUnstake on a legacy in-flight FULL request must run the full unstake: +/// remove the coins from the pool, transfer them out of the vault, and close the +/// account. +#[tokio::test] +async fn legacy_complete_full_removes_and_closes() { + let program_id = chiefstaker::id(); + let mint = Pubkey::new_unique(); + let (pool_pda, pool_bump) = 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 user = Keypair::new(); + 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 staked: u64 = 1_000_000_000; + + 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), + ); + + // cooldown 0 so the elapsed check passes regardless of the genesis clock. + pt.add_account( + pool_pda, + Account { + lamports: rent_for(StakingPool::LEN), + data: make_pool(mint, vault_pda, pool_bump, staked, 0), + owner: program_id, + executable: false, + rent_epoch: 0, + }, + ); + pt.add_account( + stake_pda, + Account { + lamports: rent_for(UserStake::LEN), + data: make_legacy_stake(user.pubkey(), pool_pda, stake_bump, staked, staked), + owner: program_id, + executable: false, + rent_epoch: 0, + }, + ); + pt.add_account( + mint, + Account { + lamports: rent_for(spl_token_2022::state::Mint::LEN), + data: pack_mint(9), + owner: spl_token_2022::id(), + executable: false, + rent_epoch: 0, + }, + ); + pt.add_account( + vault_pda, + Account { + lamports: rent_for(spl_token_2022::state::Account::LEN), + data: pack_token_account(mint, pool_pda, staked), + owner: spl_token_2022::id(), + executable: false, + rent_epoch: 0, + }, + ); + pt.add_account( + user_token.pubkey(), + Account { + lamports: rent_for(spl_token_2022::state::Account::LEN), + data: pack_token_account(mint, user.pubkey(), 0), + owner: spl_token_2022::id(), + executable: false, + rent_epoch: 0, + }, + ); + pt.add_account( + user.pubkey(), + Account { + lamports: 10_000_000_000, + 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(stake_pda, false), + AccountMeta::new(vault_pda, false), + AccountMeta::new(user_token.pubkey(), false), + AccountMeta::new_readonly(mint, false), + AccountMeta::new(user.pubkey(), true), + AccountMeta::new_readonly(spl_token_2022::id(), false), + AccountMeta::new_readonly(system_program::id(), false), + ], + data: borsh::to_vec(&StakingInstruction::CompleteUnstake).unwrap(), + }; + let tx = Transaction::new_signed_with_payer(&[ix], Some(&payer.pubkey()), &[&payer, &user], blockhash); + banks.process_transaction(tx).await.expect("legacy complete should succeed"); + + // Tokens delivered to the user, vault drained. + let user_tok = banks.get_account(user_token.pubkey()).await.unwrap().unwrap(); + let user_tok = spl_token_2022::state::Account::unpack(&user_tok.data).unwrap(); + assert_eq!(user_tok.amount, staked, "user must receive the full unstaked amount"); + + let vault = banks.get_account(vault_pda).await.unwrap().unwrap(); + let vault = spl_token_2022::state::Account::unpack(&vault.data).unwrap(); + assert_eq!(vault.amount, 0, "vault must be drained"); + + // Pool accounting reduced; stake account closed. + let pool_acc = banks.get_account(pool_pda).await.unwrap().unwrap(); + let pool = StakingPool::try_from_slice(&pool_acc.data).unwrap(); + assert_eq!(pool.total_staked, 0, "pool total_staked must drop by the unstaked amount"); + + let stake_acc = banks.get_account(stake_pda).await.unwrap(); + let closed = stake_acc.map(|a| a.lamports == 0 || a.data.iter().all(|b| *b == 0)).unwrap_or(true); + assert!(closed, "stake account must be closed after a full unstake"); +} From 6ddd909b0b9f0cd63de01a5645283695fc44e9a8 Mon Sep 17 00:00:00 2001 From: Mark Karpeles Date: Sun, 24 May 2026 00:21:47 +0900 Subject: [PATCH 5/6] CI: install protoc so cargo test can build solana-program-test solana-program-test -> solana-svm -> prost-build requires the Protocol Buffers compiler at build time. build-sbf doesn't need it, but the new cargo test step does. Install protobuf-compiler on the e2e-tests runner. Co-Authored-By: Claude Opus 4.7 (1M context) --- .github/workflows/verifiable-build.yml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/.github/workflows/verifiable-build.yml b/.github/workflows/verifiable-build.yml index 0aa1f5f..4dde583 100644 --- a/.github/workflows/verifiable-build.yml +++ b/.github/workflows/verifiable-build.yml @@ -22,6 +22,9 @@ jobs: - name: Generate keypair run: solana-keygen new --no-bip39-passphrase + - name: Install protobuf compiler (solana-program-test build dep) + run: sudo apt-get update && sudo apt-get install -y protobuf-compiler + - name: Build program run: cargo build-sbf --workspace From aafdfd4aa1e69d5fb0e09b39900c5e149b835669 Mon Sep 17 00:00:00 2001 From: Mark Karpeles Date: Sun, 24 May 2026 01:47:58 +0900 Subject: [PATCH 6/6] CI: cancel in-progress runs for the same ref+event Avoid stale runs piling up and competing for runners when pushing repeatedly to an open PR (push + pull_request each trigger a run). Co-Authored-By: Claude Opus 4.7 (1M context) --- .github/workflows/verifiable-build.yml | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/.github/workflows/verifiable-build.yml b/.github/workflows/verifiable-build.yml index 4dde583..924a91e 100644 --- a/.github/workflows/verifiable-build.yml +++ b/.github/workflows/verifiable-build.yml @@ -5,6 +5,12 @@ on: pull_request: workflow_dispatch: +# Cancel an in-progress run when a newer one starts for the same ref+event, +# so pushes don't pile up stale runs competing for runners. +concurrency: + group: ${{ github.workflow }}-${{ github.ref }}-${{ github.event_name }} + cancel-in-progress: true + env: SOLANA_VERSION: v2.0.25