From 18e4cf842be553f40971c79c5677135808ba663c Mon Sep 17 00:00:00 2001 From: juliaaschmidt Date: Thu, 23 Oct 2025 13:01:27 +0100 Subject: [PATCH 01/14] feat: add test for decrease 1e8c69eeb34a5050cd86ebe50f067770c5175839 --- program/tests/decrease.rs | 52 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 52 insertions(+) diff --git a/program/tests/decrease.rs b/program/tests/decrease.rs index 243ecabc..32e5933c 100644 --- a/program/tests/decrease.rs +++ b/program/tests/decrease.rs @@ -658,3 +658,55 @@ async fn fail_additional_with_increasing() { ) if code == StakePoolError::WrongStakeStake as u32 ); } + +#[test_case(DecreaseInstruction::Additional; "additional")] +#[test_case(DecreaseInstruction::Reserve; "reserve")] +#[test_case(DecreaseInstruction::Deprecated; "deprecated")] +#[tokio::test] +async fn fail_validator_marked_for_removal_decrease_stake(instruction_type: DecreaseInstruction) { + let ( + mut context, + stake_pool_accounts, + validator_stake, + _deposit_info, + decrease_lamports, + _reserve_lamports, + ) = setup().await; + + // First, remove the validator from the pool to mark it for removal + let error = stake_pool_accounts + .remove_validator_from_pool( + &mut context.banks_client, + &context.payer, + &context.last_blockhash, + &validator_stake.stake_account, + &validator_stake.transient_stake_account, + ) + .await; + assert!(error.is_none(), "Failed to remove validator: {:?}", error); + + // Now attempt to decrease stake on the removed validator - this should fail + let error = stake_pool_accounts + .decrease_validator_stake_either( + &mut context.banks_client, + &context.payer, + &context.last_blockhash, + &validator_stake.stake_account, + &validator_stake.transient_stake_account, + decrease_lamports, + validator_stake.transient_stake_seed, + instruction_type, + ) + .await + .unwrap() + .unwrap(); + + // Should fail with ValidatorNotFound error + assert_matches!( + error, + TransactionError::InstructionError( + _, + InstructionError::Custom(code) + ) if code == StakePoolError::ValidatorNotFound as u32 + ); +} From 4a7808ed1b0ae40d7ab5cc18a9087d22cf75073e Mon Sep 17 00:00:00 2001 From: juliaaschmidt Date: Thu, 23 Oct 2025 13:46:49 +0100 Subject: [PATCH 02/14] feat: add test for cluster-restart hijacks d99681a8e32a51fd60b480069d2c57892549f9e3 --- .../tests/update_validator_list_balance.rs | 219 +++++++++++++++++- 1 file changed, 218 insertions(+), 1 deletion(-) diff --git a/program/tests/update_validator_list_balance.rs b/program/tests/update_validator_list_balance.rs index c05defef..82b464a0 100644 --- a/program/tests/update_validator_list_balance.rs +++ b/program/tests/update_validator_list_balance.rs @@ -2,13 +2,14 @@ mod helpers; use { + bincode, helpers::*, solana_program::{ borsh1::try_from_slice_unchecked, instruction::InstructionError, program_pack::Pack, }, solana_program_test::*, solana_sdk::{ - hash::Hash, signature::Signer, stake::state::StakeStateV2, transaction::TransactionError, + hash::Hash, signature::{Keypair, Signer}, stake::state::StakeStateV2, transaction::TransactionError, }, spl_stake_pool::{ error::StakePoolError, @@ -779,3 +780,219 @@ async fn fail_with_uninitialized_validator_list() {} // TODO #[tokio::test] async fn success_with_force_destaked_validator() {} + +#[tokio::test] +async fn updates_validator_status_after_cluster_restart_merge() { + let num_validators = 1; + let ( + mut context, + last_blockhash, + stake_pool_accounts, + stake_accounts, + _, + _validator_lamports, + reserve_lamports, + _slot, + ) = setup(num_validators).await; + + let validator_stake_account = &stake_accounts[0]; + + // Get initial validator list state - should be Active + let initial_validator_list = stake_pool_accounts.get_validator_list(&mut context.banks_client).await; + let initial_validator_info = &initial_validator_list.validators[0]; + assert_eq!(initial_validator_info.status, StakeStatus::Active.into()); + + // Simulate cluster restart scenario where stake account gets reset to Initialized + let stake_pool = stake_pool_accounts.get_stake_pool(&mut context.banks_client).await; + + // Create an Initialized stake account (as would happen during cluster restart) + let initialized_meta = solana_stake_interface::state::Meta { + rent_exempt_reserve: 2282880, // Standard rent exemption + authorized: solana_stake_interface::state::Authorized { + staker: stake_pool_accounts.withdraw_authority, // Correct authorities + withdrawer: stake_pool_accounts.withdraw_authority, + }, + lockup: stake_pool.lockup, // Correct lockup + }; + + let initialized_stake_state = StakeStateV2::Initialized(initialized_meta); + + // Set the stake account to Initialized state (simulating cluster restart) + let mut stake_account_data = context + .banks_client + .get_account(validator_stake_account.stake_account) + .await + .unwrap() + .unwrap(); + + stake_account_data.data = bincode::serialize(&initialized_stake_state).unwrap(); + context.set_account(&validator_stake_account.stake_account, &stake_account_data.into()); + + // Run update_validator_list_balance - this should merge the Initialized account into reserve + // and CRITICALLY update the validator status to ReadyForRemoval + let error = stake_pool_accounts + .update_validator_list_balance( + &mut context.banks_client, + &context.payer, + &last_blockhash, + 1, + false, + ) + .await; + assert!(error.is_none(), "Update should succeed: {:?}", error); + + // MAIN TEST: Verify the validator status was properly updated to ReadyForRemoval + // This is the critical fix - without it, the status would remain Active + let post_merge_validator_list = stake_pool_accounts.get_validator_list(&mut context.banks_client).await; + let post_merge_validator_info = &post_merge_validator_list.validators[0]; + assert_eq!( + post_merge_validator_info.status, + StakeStatus::ReadyForRemoval.into(), + "Validator status should be updated to ReadyForRemoval after merging Initialized stake" + ); + + // Verify the funds were properly merged into reserve + let post_merge_reserve = context + .banks_client + .get_account(stake_pool_accounts.reserve_stake.pubkey()) + .await + .unwrap() + .unwrap(); + assert!( + post_merge_reserve.lamports > reserve_lamports, + "Reserve should have absorbed the Initialized stake funds" + ); + + // Verify active stake lamports are 0 (since account was merged) + assert_eq!( + u64::from(post_merge_validator_info.active_stake_lamports), + 0, + "Active stake lamports should be 0 after merging Initialized account" + ); + + // This test proves that Fix #2 works: "update the validator status correctly after merging the inactive stake into the reserve" +} + +#[tokio::test] +async fn ignores_unusable_stake_accounts_preventing_exploit() { + let num_validators = 1; + let ( + mut context, + last_blockhash, + stake_pool_accounts, + stake_accounts, + _, + validator_lamports, + _reserve_lamports, + _slot, + ) = setup(num_validators).await; + + let validator_stake_account = &stake_accounts[0]; + + // Verify the validator starts as Active with proper stake + let initial_validator_list = stake_pool_accounts.get_validator_list(&mut context.banks_client).await; + let initial_validator_info = &initial_validator_list.validators[0]; + assert_eq!(initial_validator_info.status, StakeStatus::Active.into()); + assert!(u64::from(initial_validator_info.active_stake_lamports) > 0); + let initial_stake_pool = stake_pool_accounts.get_stake_pool(&mut context.banks_client).await; + + // Now simulate the attack - malicious actor takes control of the stake account + // This could happen after a cluster restart where the account gets reset to Initialized + // and then the attacker manages to gain control before the pool processes it + let malicious_authority = Keypair::new(); + let extra_malicious_lamports = 1_000_000_000; // 1 SOL extra + let stake_pool = stake_pool_accounts.get_stake_pool(&mut context.banks_client).await; + + let malicious_meta = solana_stake_interface::state::Meta { + rent_exempt_reserve: 2282880, + authorized: solana_stake_interface::state::Authorized { + staker: malicious_authority.pubkey(), // WRONG - should be pool's withdraw authority + withdrawer: malicious_authority.pubkey(), // WRONG - should be pool's withdraw authority + }, + lockup: solana_stake_interface::state::Lockup { + custodian: malicious_authority.pubkey(), // WRONG - different from pool's lockup + epoch: stake_pool.lockup.epoch + 100, // WRONG - different epoch + ..stake_pool.lockup + }, + }; + + // Create a malicious delegation with the original validator + extra stake + let malicious_delegation = solana_stake_interface::state::Delegation { + voter_pubkey: validator_stake_account.vote.pubkey(), + stake: validator_lamports + extra_malicious_lamports, // Original stake + malicious extra + activation_epoch: 0, + deactivation_epoch: u64::MAX, + ..Default::default() + }; + + let malicious_stake = solana_stake_interface::state::Stake { + delegation: malicious_delegation, + credits_observed: 0, + }; + + let malicious_stake_state = StakeStateV2::Stake( + malicious_meta, + malicious_stake, + solana_stake_interface::stake_flags::StakeFlags::empty() + ); + + // Get original stake account for comparison + let original_stake_account = context + .banks_client + .get_account(validator_stake_account.stake_account) + .await + .unwrap() + .unwrap(); + let original_lamports = original_stake_account.lamports; + + // Replace the legitimate stake account with the malicious one + let mut malicious_account = original_stake_account.clone(); + malicious_account.lamports = original_lamports + extra_malicious_lamports; + malicious_account.data = bincode::serialize(&malicious_stake_state).unwrap(); + context.set_account(&validator_stake_account.stake_account, &malicious_account.into()); + + // Run update_validator_list_balance - this should succeed butignore the malicious account + let error = stake_pool_accounts + .update_validator_list_balance( + &mut context.banks_client, + &context.payer, + &last_blockhash, + 1, + false, + ) + .await; + assert!(error.is_none(), "Update should succeed but ignore malicious account: {:?}", error); + + let final_validator_list = stake_pool_accounts.get_validator_list(&mut context.banks_client).await; + let final_validator_info = &final_validator_list.validators[0]; + + // The validator should still be marked as Active but with 0 active stake + // because the malicious account was ignored + assert_eq!(final_validator_info.status, StakeStatus::Active.into()); + assert_eq!( + u64::from(final_validator_info.active_stake_lamports), + 0, + "Active stake lamports should be 0 because malicious account is ignored" + ); + + // The malicious stake account should still exist with all its lamports + // (proving it was ignored, not merged or processed) + let final_malicious_account = context + .banks_client + .get_account(validator_stake_account.stake_account) + .await + .unwrap() + .unwrap(); + + assert_eq!( + final_malicious_account.lamports, + original_lamports + extra_malicious_lamports, + "Malicious account should retain all its lamports since it was ignored" + ); + + // Verify the pool's total assets remained unchanged (malicious account was ignored) + let final_stake_pool = stake_pool_accounts.get_stake_pool(&mut context.banks_client).await; + assert_eq!(final_stake_pool.total_lamports, initial_stake_pool.total_lamports); + + // This test proves fix for "only count active validator stakes if they're usable by the pool" +} From c6051e287985dd05759a5e06e53fe4b14271c4d5 Mon Sep 17 00:00:00 2001 From: juliaaschmidt Date: Thu, 23 Oct 2025 14:51:04 +0100 Subject: [PATCH 03/14] feat: add test for do not remove validator if lamports present 4906cbaaab223dde968e081040d5ade816453082 --- .../tests/update_validator_list_balance.rs | 131 ++++++++++++++++++ 1 file changed, 131 insertions(+) diff --git a/program/tests/update_validator_list_balance.rs b/program/tests/update_validator_list_balance.rs index 82b464a0..6cd94a20 100644 --- a/program/tests/update_validator_list_balance.rs +++ b/program/tests/update_validator_list_balance.rs @@ -3,6 +3,7 @@ mod helpers; use { bincode, + borsh, helpers::*, solana_program::{ borsh1::try_from_slice_unchecked, instruction::InstructionError, program_pack::Pack, @@ -11,6 +12,7 @@ use { solana_sdk::{ hash::Hash, signature::{Keypair, Signer}, stake::state::StakeStateV2, transaction::TransactionError, }, + spl_pod::primitives::PodU64, spl_stake_pool::{ error::StakePoolError, state::{StakePool, StakeStatus, ValidatorList}, @@ -996,3 +998,132 @@ async fn ignores_unusable_stake_accounts_preventing_exploit() { // This test proves fix for "only count active validator stakes if they're usable by the pool" } + +#[tokio::test] +async fn cleanup_does_not_remove_validators_with_remaining_lamports() { + let num_validators = 2; + let ( + mut context, + last_blockhash, + stake_pool_accounts, + stake_accounts, + _, + _validator_lamports, + _reserve_lamports, + _slot, + ) = setup(num_validators).await; + + // Remove 2 validators from the pool + for stake_account in &stake_accounts { + let error = stake_pool_accounts + .remove_validator_from_pool( + &mut context.banks_client, + &context.payer, + &last_blockhash, + &stake_account.stake_account, + &stake_account.transient_stake_account, + ) + .await; + assert!(error.is_none(), "Failed to remove validator: {:?}", error); + } + + // Verify both validators are being deactivated (DeactivatingValidator status) + let validator_list = stake_pool_accounts.get_validator_list(&mut context.banks_client).await; + assert_eq!(validator_list.validators.len(), 2); + for validator_info in &validator_list.validators { + assert_eq!(validator_info.status, StakeStatus::DeactivatingValidator.into()); + } + + // Fast forward one epoch to allow the deactivating stakes to become inactive + let current_slot = context.banks_client.get_root_slot().await.unwrap(); + let slots_per_epoch = context.genesis_config().epoch_schedule.slots_per_epoch; + let next_epoch_slot = current_slot + slots_per_epoch; + context.warp_to_slot(next_epoch_slot).unwrap(); + + // Update validator list balance to process the now-inactive stakes + // This should change status from DeactivatingValidator to ReadyForRemoval + let error = stake_pool_accounts + .update_validator_list_balance( + &mut context.banks_client, + &context.payer, + &last_blockhash, + 2, // Process both validators + false, + ) + .await; + assert!(error.is_none(), "Update should succeed: {:?}", error); + + // Verify both validators are now ReadyForRemoval + let updated_validator_list = stake_pool_accounts.get_validator_list(&mut context.banks_client).await; + assert_eq!(updated_validator_list.validators.len(), 2); + for validator_info in &updated_validator_list.validators { + assert_eq!(validator_info.status, StakeStatus::ReadyForRemoval.into()); + } + + // Now manually modify the first validator to have some remaining active lamports + // This simulates a scenario where cleanup is called but one validator still has funds + let mut modified_validator_list = updated_validator_list.clone(); + modified_validator_list.validators[0].active_stake_lamports = PodU64::from(1_000_000u64); // 1 SOL remaining + modified_validator_list.validators[0].transient_stake_lamports = PodU64::from(0u64); // No transient stake + + // Update the validator list with the modified data + let validator_list_account = context + .banks_client + .get_account(stake_pool_accounts.validator_list.pubkey()) + .await + .unwrap() + .unwrap(); + let mut modified_account = validator_list_account.clone(); + + let mut serialized_data = borsh::to_vec(&modified_validator_list).unwrap(); + serialized_data.resize(modified_account.data.len(), 0); + modified_account.data = serialized_data; + context.set_account(&stake_pool_accounts.validator_list.pubkey(), &modified_account.into()); + + // Verify the validator setup is as intended (2 validators ready for removal, one with remaining lamports) + let pre_cleanup_validator_list = stake_pool_accounts.get_validator_list(&mut context.banks_client).await; + assert_eq!(pre_cleanup_validator_list.validators.len(), 2); + assert_eq!(pre_cleanup_validator_list.validators[0].status, StakeStatus::ReadyForRemoval.into()); + assert_eq!(u64::from(pre_cleanup_validator_list.validators[0].active_stake_lamports), 1_000_000); + assert_eq!(u64::from(pre_cleanup_validator_list.validators[0].transient_stake_lamports), 0); + assert_eq!(pre_cleanup_validator_list.validators[1].status, StakeStatus::ReadyForRemoval.into()); + assert_eq!(u64::from(pre_cleanup_validator_list.validators[1].active_stake_lamports), 0); + assert_eq!(u64::from(pre_cleanup_validator_list.validators[1].transient_stake_lamports), 0); + + // Run cleanup_removed_validator_entries + let error = stake_pool_accounts + .cleanup_removed_validator_entries( + &mut context.banks_client, + &context.payer, + &last_blockhash, + ) + .await; + assert!(error.is_none(), "Cleanup should succeed: {:?}", error); + + // Verify that only the validator with 0 lamports was removed + let post_cleanup_validator_list = stake_pool_accounts.get_validator_list(&mut context.banks_client).await; + + // Should have 1 validator remaining (the one with active lamports) + assert_eq!( + post_cleanup_validator_list.validators.len(), + 1, + "Only validator with 0 lamports should have been removed" + ); + + // The remaining validator should be the one with active lamports + assert_eq!( + u64::from(post_cleanup_validator_list.validators[0].active_stake_lamports), + 1_000_000, + "Validator with remaining lamports should not have been removed" + ); + assert_eq!( + post_cleanup_validator_list.validators[0].status, + StakeStatus::ReadyForRemoval.into(), + "Validator status should remain ReadyForRemoval" + ); + + // This test proves that ValidatorStakeInfo::is_removed() correctly checks both: + // 1. Status == ReadyForRemoval + // 2. Both active_stake_lamports AND transient_stake_lamports == 0 + // Only validators meeting BOTH criteria are removed by cleanup_removed_validator_entries +} From ad8396f1d626d6b8c67a69b4b456f9a9f74cbe76 Mon Sep 17 00:00:00 2001 From: juliaaschmidt Date: Thu, 23 Oct 2025 15:33:59 +0100 Subject: [PATCH 04/14] style: rm superfluous comment --- program/tests/update_validator_list_balance.rs | 5 ----- 1 file changed, 5 deletions(-) diff --git a/program/tests/update_validator_list_balance.rs b/program/tests/update_validator_list_balance.rs index 6cd94a20..bbe6d97d 100644 --- a/program/tests/update_validator_list_balance.rs +++ b/program/tests/update_validator_list_balance.rs @@ -1121,9 +1121,4 @@ async fn cleanup_does_not_remove_validators_with_remaining_lamports() { StakeStatus::ReadyForRemoval.into(), "Validator status should remain ReadyForRemoval" ); - - // This test proves that ValidatorStakeInfo::is_removed() correctly checks both: - // 1. Status == ReadyForRemoval - // 2. Both active_stake_lamports AND transient_stake_lamports == 0 - // Only validators meeting BOTH criteria are removed by cleanup_removed_validator_entries } From 416931280fda26ecb89df5aad9b2cb9e847825c5 Mon Sep 17 00:00:00 2001 From: juliaaschmidt Date: Thu, 23 Oct 2025 16:09:30 +0100 Subject: [PATCH 05/14] feat: ignore Uninitialized accounts during update_validator_list_balance 21d8af09095e7d9a6fad3231306436eb01ffe838 --- .../tests/update_validator_list_balance.rs | 108 ++++++++++++++++++ 1 file changed, 108 insertions(+) diff --git a/program/tests/update_validator_list_balance.rs b/program/tests/update_validator_list_balance.rs index bbe6d97d..25adef73 100644 --- a/program/tests/update_validator_list_balance.rs +++ b/program/tests/update_validator_list_balance.rs @@ -999,6 +999,114 @@ async fn ignores_unusable_stake_accounts_preventing_exploit() { // This test proves fix for "only count active validator stakes if they're usable by the pool" } +#[tokio::test] +async fn update_validator_list_balance_ingores_uninitialized_stake_account_balances() { + let num_validators = 1; + let ( + mut context, + last_blockhash, + stake_pool_accounts, + stake_accounts, + _, + _validator_lamports, + _reserve_lamports, + mut slot, + ) = setup(num_validators).await; + + let validator_stake_account = &stake_accounts[0]; + + // Verify the validator starts as Active + let initial_validator_list = stake_pool_accounts.get_validator_list(&mut context.banks_client).await; + let initial_validator_info = &initial_validator_list.validators[0]; + assert_eq!(initial_validator_info.status, StakeStatus::Active.into()); + assert!(u64::from(initial_validator_info.active_stake_lamports) > 0); + + // First, remove the validator from the pool to trigger deactivation + let error = stake_pool_accounts + .remove_validator_from_pool( + &mut context.banks_client, + &context.payer, + &last_blockhash, + &validator_stake_account.stake_account, + &validator_stake_account.transient_stake_account, + ) + .await; + assert!(error.is_none(), "Failed to remove validator: {:?}", error); + + // Verify validator is now being deactivated + let deactivating_validator_list = stake_pool_accounts.get_validator_list(&mut context.banks_client).await; + let deactivating_validator_info = &deactivating_validator_list.validators[0]; + assert_eq!(deactivating_validator_info.status, StakeStatus::DeactivatingValidator.into()); + + // Fast forward one epoch to allow the stake to deactivate + let slots_per_epoch = context.genesis_config().epoch_schedule.slots_per_epoch; + slot += slots_per_epoch; + context.warp_to_slot(slot).unwrap(); + + let new_blockhash = context + .banks_client + .get_new_latest_blockhash(&last_blockhash) + .await + .unwrap(); + + // Now simulate a scenario where the stake account becomes uninitialized + // This could happen due to cluster restart or other network events + let uninitialized_stake_state = StakeStateV2::Uninitialized; + + let original_stake_account = context + .banks_client + .get_account(validator_stake_account.stake_account) + .await + .unwrap() + .unwrap(); + + let mut modified_account = original_stake_account.clone(); + modified_account.data = bincode::serialize(&uninitialized_stake_state).unwrap(); + context.set_account(&validator_stake_account.stake_account, &modified_account.into()); + + // Run update_validator_list_balance - this should trigger the uninitialized account handling + let error = stake_pool_accounts + .update_validator_list_balance( + &mut context.banks_client, + &context.payer, + &new_blockhash, + 1, + false, + ) + .await; + assert!(error.is_none(), "Update should succeed despite uninitialized account: {:?}", error); + + // Verify the validator status and that the uninitialized account was ignored + let final_validator_list = stake_pool_accounts.get_validator_list(&mut context.banks_client).await; + let final_validator_info = &final_validator_list.validators[0]; + + // The validator should still be in DeactivatingValidator status with 0 active stake + // because the uninitialized account was ignored + assert_eq!(final_validator_info.status, StakeStatus::DeactivatingValidator.into()); + assert_eq!( + u64::from(final_validator_info.active_stake_lamports), + 0, + "Active stake lamports should be 0 because uninitialized account is ignored" + ); + + // Verify the uninitialized stake account still exists + let final_stake_account = context + .banks_client + .get_account(validator_stake_account.stake_account) + .await + .unwrap() + .unwrap(); + + // Verify it's still uninitialized + let final_stake_state: StakeStateV2 = bincode::deserialize(&final_stake_account.data).unwrap(); + matches!(final_stake_state, StakeStateV2::Uninitialized); + + // assert it has a nonzero account balance + assert!(final_stake_account.lamports > 0); + + // This test proves that the uninitialized account scenario can be triggered and is handled correctly +} + #[tokio::test] async fn cleanup_does_not_remove_validators_with_remaining_lamports() { let num_validators = 2; From 4cb27a0d6080a908ae99a9539d2ed60e0cb3cd40 Mon Sep 17 00:00:00 2001 From: juliaaschmidt Date: Thu, 30 Oct 2025 11:53:22 +0000 Subject: [PATCH 06/14] feat: withdraw stake tests with transient lamports 40c260abd4e46ac464c249115c9bf03a998111b4 --- program/tests/withdraw_edge_cases.rs | 289 +++++++++++++++++++++++++++ 1 file changed, 289 insertions(+) diff --git a/program/tests/withdraw_edge_cases.rs b/program/tests/withdraw_edge_cases.rs index 3fcd063d..4f7c9d02 100644 --- a/program/tests/withdraw_edge_cases.rs +++ b/program/tests/withdraw_edge_cases.rs @@ -18,6 +18,137 @@ use { test_case::test_case, }; +#[tokio::test] +async fn fail_remove_validator_blocked_by_transient_stake() { + let ( + mut context, + stake_pool_accounts, + validator_stake, + deposit_info, + user_transfer_authority, + user_stake_recipient, + sol_withdraw_authority, + _, + ) = setup_for_withdraw(spl_token::id(), STAKE_ACCOUNT_RENT_EXEMPTION).await; + + // Step 1: Create transient stake by decreasing some validator stake + let decrease_amount = deposit_info.stake_lamports / 3; // Decrease 1/3 of the stake + let error = stake_pool_accounts + .decrease_validator_stake_either( + &mut context.banks_client, + &context.payer, + &context.last_blockhash, + &validator_stake.stake_account, + &validator_stake.transient_stake_account, + decrease_amount, + validator_stake.transient_stake_seed, + DecreaseInstruction::Additional, // This creates transient stake that won't merge back + ) + .await; + assert!(error.is_none(), "{:?}", error); + + // Step 2: Check the state after creating transient stake + let validator_stake_account_after_decrease = + get_account(&mut context.banks_client, &validator_stake.stake_account).await; + let transient_stake_account_after_decrease = + get_account(&mut context.banks_client, &validator_stake.transient_stake_account).await; + + println!("Validator stake after decrease: {} lamports", validator_stake_account_after_decrease.lamports); + println!("Transient stake after decrease: {} lamports", transient_stake_account_after_decrease.lamports); + + // Step 3: Warp forward to deactivation epoch + let first_normal_slot = context.genesis_config().epoch_schedule.first_normal_slot; + let slot = first_normal_slot + 1; + context.warp_to_slot(slot).unwrap(); + + // Step 4: Update with no_merge=true to keep transient stake separate + let error = stake_pool_accounts + .update_all( + &mut context.banks_client, + &context.payer, + &context.last_blockhash, + true, // no_merge = true to prevent merging transient stake back + ) + .await; + assert!(error.is_none(), "{:?}", error); + + // Step 5: Check final state - validator should have active stake, transient should remain separate + let validator_stake_account_final = + get_account(&mut context.banks_client, &validator_stake.stake_account).await; + let transient_stake_account_final = context + .banks_client + .get_account(validator_stake.transient_stake_account) + .await + .unwrap(); + + println!(" Validator stake: {} lamports", validator_stake_account_final.lamports); + + // Verify transient stake still exists + let transient_account = transient_stake_account_final.expect("Transient stake account should still exist"); + println!(" Transient stake: {} lamports", transient_account.lamports); + + // Step 6: Try to withdraw ALL active stake - this should FAIL because transient stake blocks removal + let rent = context.banks_client.get_rent().await.unwrap(); + let stake_rent = rent.minimum_balance(std::mem::size_of::()); + + let stake_pool = stake_pool_accounts + .get_stake_pool(&mut context.banks_client) + .await; + + // Try to withdraw everything except rent (this should trigger the validator removal check) + let active_stake_to_withdraw = validator_stake_account_final.lamports + .saturating_sub(stake_rent); + + println!("Testing complete withdrawal of active stake {} lamports (should fail due to transient stake)", active_stake_to_withdraw); + + let pool_tokens_all = (active_stake_to_withdraw * stake_pool.pool_token_supply) + .checked_div(stake_pool.total_lamports) + .unwrap(); + let pool_tokens_all_with_fee = stake_pool_accounts.calculate_inverse_withdrawal_fee(pool_tokens_all); + + let new_user_authority_all = Pubkey::new_unique(); + let error_all = stake_pool_accounts + .withdraw_stake( + &mut context.banks_client, + &context.payer, + &context.last_blockhash, + &sol_withdraw_authority, + &user_stake_recipient.pubkey(), + &user_transfer_authority, + &deposit_info.pool_account.pubkey(), + &validator_stake.stake_account, // Withdraw from active stake account + &new_user_authority_all, + pool_tokens_all_with_fee, + ) + .await; + + // Step 7: Verify that the complete withdrawal fails because transient stake prevents validator removal + assert!(error_all.is_some(), "Complete withdrawal should fail because validator has transient stake"); + let transaction_error = error_all.unwrap().unwrap(); + + println!("Complete withdrawal correctly failed with error: {:?}", transaction_error); + // The error should be StakeLamportsNotEqualToMinimum because the program detects + // that this validator cannot be removed due to associated transient lamports + assert_eq!( + transaction_error, + TransactionError::InstructionError( + 0, + InstructionError::Custom(StakePoolError::StakeLamportsNotEqualToMinimum as u32) + ) + ); + + // Step 8: Verify that the transient stake is still there and blocking removal + let final_transient_check = context + .banks_client + .get_account(validator_stake.transient_stake_account) + .await + .unwrap(); + + assert!(final_transient_check.is_some(), "Transient stake account should still exist"); + let final_transient = final_transient_check.unwrap(); + assert!(final_transient.lamports > 0, "Transient stake should still have lamports"); +} + #[tokio::test] async fn fail_remove_validator() { let ( @@ -307,6 +438,164 @@ async fn fail_with_reserve() { ); } +#[tokio::test] +async fn fail_with_transient() { + let ( + mut context, + stake_pool_accounts, + validator_stake, + deposit_info, + user_transfer_authority, + user_stake_recipient, + sol_withdraw_authority, + _, + ) = setup_for_withdraw(spl_token::id(), STAKE_ACCOUNT_RENT_EXEMPTION).await; + + // warp forward to after reward payout + let first_normal_slot = context.genesis_config().epoch_schedule.first_normal_slot; + let mut slot = first_normal_slot + 1; + context.warp_to_slot(slot).unwrap(); + + let error = stake_pool_accounts + .update_all( + &mut context.banks_client, + &context.payer, + &context.last_blockhash, + false, + ) + .await; + assert!(error.is_none(), "{:?}", error); + + let rent = context.banks_client.get_rent().await.unwrap(); + let stake_rent = rent.minimum_balance(std::mem::size_of::()); + let stake_minimum_delegation = + stake_get_minimum_delegation(&mut context.banks_client, &context.payer, &context.last_blockhash) + .await; + + // Calculate how much to decrease to leave only rent + minimum delegation in validator stake account + // This will create a transient stake account with the decreased amount + let validator_stake_account_before = + get_account(&mut context.banks_client, &validator_stake.stake_account).await; + let current_validator_lamports = validator_stake_account_before.lamports; + let target_validator_lamports = stake_rent + stake_minimum_delegation; + let amount_to_decrease = current_validator_lamports - target_validator_lamports; + + println!("Current validator stake: {} lamports", current_validator_lamports); + println!("Target validator stake: {} lamports (rent + min delegation)", target_validator_lamports); + println!("Amount to decrease: {} lamports", amount_to_decrease); + + // Decrease validator stake, creating transient stake + let error = stake_pool_accounts + .decrease_validator_stake_either( + &mut context.banks_client, + &context.payer, + &context.last_blockhash, + &validator_stake.stake_account, + &validator_stake.transient_stake_account, + amount_to_decrease, + validator_stake.transient_stake_seed, + DecreaseInstruction::Additional, + ) + .await; + assert!(error.is_none(), "{:?}", error); + + // Check the state after decrease - validator should have minimum, transient should have the decreased amount + let validator_stake_account_after = + get_account(&mut context.banks_client, &validator_stake.stake_account).await; + let transient_stake_account = + get_account(&mut context.banks_client, &validator_stake.transient_stake_account).await; + + println!("Validator stake after decrease: {} lamports", validator_stake_account_after.lamports); + println!("Transient stake after decrease: {} lamports", transient_stake_account.lamports); + + // warp forward to deactivation epoch + slot += context.genesis_config().epoch_schedule.slots_per_epoch; + context.warp_to_slot(slot).unwrap(); + + let last_blockhash = context + .banks_client + .get_new_latest_blockhash(&context.last_blockhash) + .await + .unwrap(); + + // update to merge deactivated stake into reserve, but use no_merge=true to keep transient stake separate + let error = stake_pool_accounts + .update_all( + &mut context.banks_client, + &context.payer, + &last_blockhash, + true, // no_merge = true to prevent merging transient stake back + ) + .await; + assert!(error.is_none(), "{:?}", error); + + // Check final state before withdrawal + let validator_stake_account_final = + get_account(&mut context.banks_client, &validator_stake.stake_account).await; + + // Check if transient stake account still exists (it might have been merged back) + let transient_stake_account_final = context + .banks_client + .get_account(validator_stake.transient_stake_account) + .await + .unwrap(); + + println!("Validator stake after epoch change: {} lamports", validator_stake_account_final.lamports); + + // Verify transient stake account still exists and has lamports + let transient_account = transient_stake_account_final.expect("Transient stake account should still exist"); + println!("Transient stake after epoch change: {} lamports", transient_account.lamports); + + // Calculate pool tokens needed to withdraw EXACTLY the transient stake amount + let stake_pool = stake_pool_accounts + .get_stake_pool(&mut context.banks_client) + .await; + + // We want to withdraw exactly what's in the transient account (all of it) + let exact_transient_lamports = transient_account.lamports; + + // Calculate the exact pool tokens needed for this withdrawal amount + // Using the same formula as the stake pool: pool_tokens = (lamports * pool_token_supply) / total_lamports + let pool_tokens_post_fee = (exact_transient_lamports * stake_pool.pool_token_supply) + .checked_div(stake_pool.total_lamports) + .unwrap(); + let pool_tokens = stake_pool_accounts.calculate_inverse_withdrawal_fee(pool_tokens_post_fee); + + println!("Transient account has exactly: {} lamports", exact_transient_lamports); + println!("Attempting to withdraw exactly: {} lamports from transient stake", exact_transient_lamports); + println!("Pool tokens calculated: {} (post-fee: {})", pool_tokens, pool_tokens_post_fee); + println!("Stake pool total lamports: {}, pool token supply: {}", stake_pool.total_lamports, stake_pool.pool_token_supply); + + let new_user_authority = Pubkey::new_unique(); + + // Try to withdraw from transient stake account - this should FAIL + let error = stake_pool_accounts + .withdraw_stake( + &mut context.banks_client, + &context.payer, + &last_blockhash, + &sol_withdraw_authority, + &user_stake_recipient.pubkey(), + &user_transfer_authority, + &deposit_info.pool_account.pubkey(), + &validator_stake.transient_stake_account, // ✅ Withdraw from transient stake account + &new_user_authority, + pool_tokens, + ) + .await; + + // Verify that the withdrawal fails with the expected error + assert!(error.is_some(), "Withdrawal should fail"); + let transaction_error = error.unwrap().unwrap(); + assert_eq!( + transaction_error, + TransactionError::InstructionError( + 0, + InstructionError::Custom(StakePoolError::StakeLamportsNotEqualToMinimum as u32) + ) + ); +} + #[tokio::test] async fn success_with_reserve() { let ( From 79f3326b3c3d0cd201eb5ababb02213767bee553 Mon Sep 17 00:00:00 2001 From: juliaaschmidt Date: Thu, 30 Oct 2025 11:53:30 +0000 Subject: [PATCH 07/14] style: comments --- program/tests/update_validator_list_balance.rs | 3 --- 1 file changed, 3 deletions(-) diff --git a/program/tests/update_validator_list_balance.rs b/program/tests/update_validator_list_balance.rs index 25adef73..0af47fc2 100644 --- a/program/tests/update_validator_list_balance.rs +++ b/program/tests/update_validator_list_balance.rs @@ -1169,12 +1169,9 @@ async fn cleanup_does_not_remove_validators_with_remaining_lamports() { } // Now manually modify the first validator to have some remaining active lamports - // This simulates a scenario where cleanup is called but one validator still has funds let mut modified_validator_list = updated_validator_list.clone(); modified_validator_list.validators[0].active_stake_lamports = PodU64::from(1_000_000u64); // 1 SOL remaining modified_validator_list.validators[0].transient_stake_lamports = PodU64::from(0u64); // No transient stake - - // Update the validator list with the modified data let validator_list_account = context .banks_client .get_account(stake_pool_accounts.validator_list.pubkey()) From 7f868a22f8af24903794430d12c01fa954778ba3 Mon Sep 17 00:00:00 2001 From: juliaaschmidt Date: Thu, 30 Oct 2025 12:45:49 +0000 Subject: [PATCH 08/14] feat: add test for reset preferred validator during withdrawal 2ce9b229adfd33b3224225edb67359503e8c54ec --- program/tests/withdraw_edge_cases.rs | 174 +++++++++++++++++++++++++++ 1 file changed, 174 insertions(+) diff --git a/program/tests/withdraw_edge_cases.rs b/program/tests/withdraw_edge_cases.rs index 4f7c9d02..56ffeb9d 100644 --- a/program/tests/withdraw_edge_cases.rs +++ b/program/tests/withdraw_edge_cases.rs @@ -1320,3 +1320,177 @@ async fn fail_overdraw_reserve() { ) ); } + +#[tokio::test] +async fn success_remove_preferred_validator_resets_preference() { + let ( + mut context, + stake_pool_accounts, + validator_stake, + deposit_info, + user_transfer_authority, + user_stake_recipient, + sol_withdraw_authority, + _, + ) = setup_for_withdraw(spl_token::id(), STAKE_ACCOUNT_RENT_EXEMPTION).await; + + let last_blockhash = context.banks_client + .get_new_latest_blockhash(&context.last_blockhash).await + .unwrap(); + + // Set the validator as the preferred withdraw validator + stake_pool_accounts.set_preferred_validator( + &mut context.banks_client, + &context.payer, + &last_blockhash, + instruction::PreferredValidatorType::Withdraw, + Some(validator_stake.vote.pubkey()) + ).await; + + // Also set it as preferred deposit validator to test both reset paths + stake_pool_accounts.set_preferred_validator( + &mut context.banks_client, + &context.payer, + &last_blockhash, + instruction::PreferredValidatorType::Deposit, + Some(validator_stake.vote.pubkey()) + ).await; + + // Verify the preferred deposit and withdraw validators are set + let stake_pool_before = stake_pool_accounts.get_stake_pool(&mut context.banks_client).await; + assert_eq!( + stake_pool_before.preferred_withdraw_validator_vote_address, + Some(validator_stake.vote.pubkey()) + ); + assert_eq!( + stake_pool_before.preferred_deposit_validator_vote_address, + Some(validator_stake.vote.pubkey()) + ); + + println!("Preferred validators set to: {}", validator_stake.vote.pubkey()); + + // Warp forward to after reward payout + let first_normal_slot = context.genesis_config().epoch_schedule.first_normal_slot; + let mut slot = first_normal_slot + 1; + context.warp_to_slot(slot).unwrap(); + + let error = stake_pool_accounts.update_all( + &mut context.banks_client, + &context.payer, + &last_blockhash, + false + ).await; + assert!(error.is_none(), "{:?}", error); + + let rent = context.banks_client.get_rent().await.unwrap(); + let stake_rent = rent.minimum_balance(std::mem::size_of::()); + let stake_pool = stake_pool_accounts.get_stake_pool(&mut context.banks_client).await; + let lamports_per_pool_token = stake_pool.get_lamports_per_pool_token().unwrap(); + + // Decrease all of stake except for exactly lamports_per_pool_token lamports + // This will leave the minimum amount that can be withdrawn completely + let error = stake_pool_accounts.decrease_validator_stake_either( + &mut context.banks_client, + &context.payer, + &last_blockhash, + &validator_stake.stake_account, + &validator_stake.transient_stake_account, + deposit_info.stake_lamports + stake_rent - lamports_per_pool_token, + validator_stake.transient_stake_seed, + DecreaseInstruction::Reserve + ).await; + assert!(error.is_none(), "{:?}", error); + + // Warp forward to deactivation + slot += context.genesis_config().epoch_schedule.slots_per_epoch; + context.warp_to_slot(slot).unwrap(); + + let last_blockhash = context.banks_client + .get_new_latest_blockhash(&last_blockhash).await + .unwrap(); + + // Update to merge deactivated stake into reserve + let error = stake_pool_accounts.update_all( + &mut context.banks_client, + &context.payer, + &last_blockhash, + false + ).await; + assert!(error.is_none(), "{:?}", error); + + let validator_stake_account = get_account( + &mut context.banks_client, + &validator_stake.stake_account + ).await; + let remaining_lamports = validator_stake_account.lamports; + let stake_minimum_delegation = stake_get_minimum_delegation( + &mut context.banks_client, + &context.payer, + &last_blockhash + ).await; + + println!("Remaining lamports in validator: {}", remaining_lamports); + println!("Stake rent: {}", stake_rent); + println!("Minimum delegation: {}", stake_minimum_delegation); + // Make sure it's actually more than the minimum (should be exactly lamports_per_pool_token) + assert!(remaining_lamports > stake_rent + stake_minimum_delegation); + + // Calculate pool tokens needed to withdraw everything (this should remove the validator completely) + let stake_pool_updated = stake_pool_accounts.get_stake_pool(&mut context.banks_client).await; + let pool_tokens_post_fee = (remaining_lamports * stake_pool_updated.pool_token_supply).div_ceil( + stake_pool_updated.total_lamports + ); + let pool_tokens = stake_pool_accounts.calculate_inverse_withdrawal_fee(pool_tokens_post_fee); + + println!("Pool tokens needed for complete withdrawal: {}", pool_tokens); + + let new_user_authority = Pubkey::new_unique(); + + // Perform the complete withdrawal - this should trigger validator removal and reset preferred validators + let error = stake_pool_accounts.withdraw_stake( + &mut context.banks_client, + &context.payer, + &last_blockhash, + &sol_withdraw_authority, + &user_stake_recipient.pubkey(), + &user_transfer_authority, + &deposit_info.pool_account.pubkey(), + &validator_stake.stake_account, + &new_user_authority, + pool_tokens + ).await; + assert!(error.is_none(), "Complete withdrawal should succeed: {:?}", error); + + println!("Complete withdrawal successful - validator should be removed"); + + // Verify validator stake account is gone + let validator_stake_account_after = context.banks_client + .get_account(validator_stake.stake_account).await + .unwrap(); + assert!(validator_stake_account_after.is_none(), "Validator stake account should be removed"); + + // Verify that preferred validators have been reset to None + let stake_pool_after = stake_pool_accounts.get_stake_pool(&mut context.banks_client).await; + + assert_eq!( + stake_pool_after.preferred_withdraw_validator_vote_address, + None, + "Preferred withdraw validator should be reset to None" + ); + assert_eq!( + stake_pool_after.preferred_deposit_validator_vote_address, + None, + "Preferred deposit validator should be reset to None" + ); + + // Verify user received the stake + let user_stake_recipient_account = get_account( + &mut context.banks_client, + &user_stake_recipient.pubkey() + ).await; + assert_eq!( + user_stake_recipient_account.lamports, + remaining_lamports + stake_rent, + "User should receive all lamports from removed validator" + ); +} From 756f88578c5ed9b8ff13d4cfd3a2cd6b8d5ef207 Mon Sep 17 00:00:00 2001 From: juliaaschmidt Date: Thu, 30 Oct 2025 12:51:41 +0000 Subject: [PATCH 09/14] style: rm comment --- program/tests/withdraw_edge_cases.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/program/tests/withdraw_edge_cases.rs b/program/tests/withdraw_edge_cases.rs index 56ffeb9d..ef3c8f54 100644 --- a/program/tests/withdraw_edge_cases.rs +++ b/program/tests/withdraw_edge_cases.rs @@ -578,7 +578,7 @@ async fn fail_with_transient() { &user_stake_recipient.pubkey(), &user_transfer_authority, &deposit_info.pool_account.pubkey(), - &validator_stake.transient_stake_account, // ✅ Withdraw from transient stake account + &validator_stake.transient_stake_account, &new_user_authority, pool_tokens, ) From 9de28ee09f24c483b248d88cf3b64db4fc9511d0 Mon Sep 17 00:00:00 2001 From: juliaaschmidt Date: Thu, 30 Oct 2025 16:38:23 +0000 Subject: [PATCH 10/14] feat: Set status to DeactivatingAll if transient lamports 82e9fa9bd3ff4b226b5d1be6816b4470886fd6f --- program/tests/vsa_remove.rs | 89 +++++++++++++++++++++++++++++++++++++ 1 file changed, 89 insertions(+) diff --git a/program/tests/vsa_remove.rs b/program/tests/vsa_remove.rs index 21b07f58..dfc9f5c8 100644 --- a/program/tests/vsa_remove.rs +++ b/program/tests/vsa_remove.rs @@ -948,3 +948,92 @@ async fn update_no_merge_after_removal() { }; assert_eq!(validator_list, expected_list); } + +#[tokio::test] +async fn success_remove_validator_with_transient_stake_triggers_deactivating_all() { + let (mut context, stake_pool_accounts, validator_stake) = setup().await; + + // First, increase validator stake to create a transient stake account + let increase_amount = TEST_STAKE_AMOUNT; + let error = stake_pool_accounts + .increase_validator_stake( + &mut context.banks_client, + &context.payer, + &context.last_blockhash, + &validator_stake.transient_stake_account, + &validator_stake.stake_account, + &validator_stake.vote.pubkey(), + increase_amount, + validator_stake.transient_stake_seed, + ) + .await; + assert!(error.is_none(), "{:?}", error); + + // Update the validator list to register the transient stake + let error = stake_pool_accounts + .update_validator_list_balance( + &mut context.banks_client, + &context.payer, + &context.last_blockhash, + 1, + true, // no_merge = true to keep transient stake separate + ) + .await; + assert!(error.is_none(), "{:?}", error); + + // Verify the validator has both active and transient stake + let validator_list_before = stake_pool_accounts + .get_validator_list(&mut context.banks_client) + .await; + let validator_info_before = &validator_list_before.validators[0]; + assert!(u64::from(validator_info_before.active_stake_lamports) > 0); + assert!(u64::from(validator_info_before.transient_stake_lamports) > 0); + let status_before: state::StakeStatus = validator_info_before.status.try_into().unwrap(); + assert_eq!(status_before, state::StakeStatus::Active); + + // Now remove the validator - this should trigger DeactivatingAll status + // because the validator has transient_stake_lamports > 0 + let error = stake_pool_accounts + .remove_validator_from_pool( + &mut context.banks_client, + &context.payer, + &context.last_blockhash, + &validator_stake.stake_account, + &validator_stake.transient_stake_account, + ) + .await; + assert!(error.is_none(), "{:?}", error); + + // Verify the validator status is now DeactivatingAll + let validator_list_after = stake_pool_accounts + .get_validator_list(&mut context.banks_client) + .await; + let validator_info_after = &validator_list_after.validators[0]; + let status: state::StakeStatus = validator_info_after.status.try_into().unwrap(); + assert_eq!( + status, + state::StakeStatus::DeactivatingAll, + "Validator with transient stake should be marked as DeactivatingAll" + ); + + // Verify both active and transient stake accounts are being deactivated + let active_stake_account = get_account(&mut context.banks_client, &validator_stake.stake_account).await; + let active_stake_state = deserialize::(&active_stake_account.data).unwrap(); + if let stake::state::StakeStateV2::Stake(_, active_stake, _) = active_stake_state { + assert_ne!( + active_stake.delegation.deactivation_epoch, + u64::MAX, + "Active stake should be deactivating" + ); + } + + let transient_stake_account = get_account(&mut context.banks_client, &validator_stake.transient_stake_account).await; + let transient_stake_state = deserialize::(&transient_stake_account.data).unwrap(); + if let stake::state::StakeStateV2::Stake(_, transient_stake, _) = transient_stake_state { + assert_ne!( + transient_stake.delegation.deactivation_epoch, + u64::MAX, + "Transient stake should be deactivating" + ); + } +} From 4d8d2f747f2905ea38a81bdd4227795926ba69dc Mon Sep 17 00:00:00 2001 From: juliaaschmidt Date: Mon, 15 Dec 2025 13:11:12 +0100 Subject: [PATCH 11/14] gpg test --- program/tests/withdraw_edge_cases.rs | 6 ------ 1 file changed, 6 deletions(-) diff --git a/program/tests/withdraw_edge_cases.rs b/program/tests/withdraw_edge_cases.rs index ef3c8f54..39f78383 100644 --- a/program/tests/withdraw_edge_cases.rs +++ b/program/tests/withdraw_edge_cases.rs @@ -27,7 +27,6 @@ async fn fail_remove_validator_blocked_by_transient_stake() { deposit_info, user_transfer_authority, user_stake_recipient, - sol_withdraw_authority, _, ) = setup_for_withdraw(spl_token::id(), STAKE_ACCOUNT_RENT_EXEMPTION).await; @@ -112,7 +111,6 @@ async fn fail_remove_validator_blocked_by_transient_stake() { &mut context.banks_client, &context.payer, &context.last_blockhash, - &sol_withdraw_authority, &user_stake_recipient.pubkey(), &user_transfer_authority, &deposit_info.pool_account.pubkey(), @@ -447,7 +445,6 @@ async fn fail_with_transient() { deposit_info, user_transfer_authority, user_stake_recipient, - sol_withdraw_authority, _, ) = setup_for_withdraw(spl_token::id(), STAKE_ACCOUNT_RENT_EXEMPTION).await; @@ -574,7 +571,6 @@ async fn fail_with_transient() { &mut context.banks_client, &context.payer, &last_blockhash, - &sol_withdraw_authority, &user_stake_recipient.pubkey(), &user_transfer_authority, &deposit_info.pool_account.pubkey(), @@ -1330,7 +1326,6 @@ async fn success_remove_preferred_validator_resets_preference() { deposit_info, user_transfer_authority, user_stake_recipient, - sol_withdraw_authority, _, ) = setup_for_withdraw(spl_token::id(), STAKE_ACCOUNT_RENT_EXEMPTION).await; @@ -1451,7 +1446,6 @@ async fn success_remove_preferred_validator_resets_preference() { &mut context.banks_client, &context.payer, &last_blockhash, - &sol_withdraw_authority, &user_stake_recipient.pubkey(), &user_transfer_authority, &deposit_info.pool_account.pubkey(), From c22bb6cab78921dd7a8ff3ab18c594a7a4dd5c67 Mon Sep 17 00:00:00 2001 From: juliaaschmidt Date: Mon, 15 Dec 2025 13:11:25 +0100 Subject: [PATCH 12/14] gpg test From e3d12df79784a77ee79f5bd7360d92c7dc580749 Mon Sep 17 00:00:00 2001 From: juliaaschmidt Date: Mon, 15 Dec 2025 13:11:36 +0100 Subject: [PATCH 13/14] fix: rm sol withdraw authority from tests --- .../tests/update_validator_list_balance.rs | 250 ++++++++++++------ 1 file changed, 166 insertions(+), 84 deletions(-) diff --git a/program/tests/update_validator_list_balance.rs b/program/tests/update_validator_list_balance.rs index 0af47fc2..cdc0fc2b 100644 --- a/program/tests/update_validator_list_balance.rs +++ b/program/tests/update_validator_list_balance.rs @@ -2,15 +2,17 @@ mod helpers; use { - bincode, - borsh, + bincode, borsh, helpers::*, solana_program::{ borsh1::try_from_slice_unchecked, instruction::InstructionError, program_pack::Pack, }, solana_program_test::*, solana_sdk::{ - hash::Hash, signature::{Keypair, Signer}, stake::state::StakeStateV2, transaction::TransactionError, + hash::Hash, + signature::{Keypair, Signer}, + stake::state::StakeStateV2, + transaction::TransactionError, }, spl_pod::primitives::PodU64, spl_stake_pool::{ @@ -798,27 +800,31 @@ async fn updates_validator_status_after_cluster_restart_merge() { ) = setup(num_validators).await; let validator_stake_account = &stake_accounts[0]; - + // Get initial validator list state - should be Active - let initial_validator_list = stake_pool_accounts.get_validator_list(&mut context.banks_client).await; + let initial_validator_list = stake_pool_accounts + .get_validator_list(&mut context.banks_client) + .await; let initial_validator_info = &initial_validator_list.validators[0]; assert_eq!(initial_validator_info.status, StakeStatus::Active.into()); - + // Simulate cluster restart scenario where stake account gets reset to Initialized - let stake_pool = stake_pool_accounts.get_stake_pool(&mut context.banks_client).await; - + let stake_pool = stake_pool_accounts + .get_stake_pool(&mut context.banks_client) + .await; + // Create an Initialized stake account (as would happen during cluster restart) let initialized_meta = solana_stake_interface::state::Meta { rent_exempt_reserve: 2282880, // Standard rent exemption authorized: solana_stake_interface::state::Authorized { - staker: stake_pool_accounts.withdraw_authority, // Correct authorities + staker: stake_pool_accounts.withdraw_authority, // Correct authorities withdrawer: stake_pool_accounts.withdraw_authority, }, lockup: stake_pool.lockup, // Correct lockup }; - + let initialized_stake_state = StakeStateV2::Initialized(initialized_meta); - + // Set the stake account to Initialized state (simulating cluster restart) let mut stake_account_data = context .banks_client @@ -826,10 +832,13 @@ async fn updates_validator_status_after_cluster_restart_merge() { .await .unwrap() .unwrap(); - + stake_account_data.data = bincode::serialize(&initialized_stake_state).unwrap(); - context.set_account(&validator_stake_account.stake_account, &stake_account_data.into()); - + context.set_account( + &validator_stake_account.stake_account, + &stake_account_data.into(), + ); + // Run update_validator_list_balance - this should merge the Initialized account into reserve // and CRITICALLY update the validator status to ReadyForRemoval let error = stake_pool_accounts @@ -842,17 +851,19 @@ async fn updates_validator_status_after_cluster_restart_merge() { ) .await; assert!(error.is_none(), "Update should succeed: {:?}", error); - + // MAIN TEST: Verify the validator status was properly updated to ReadyForRemoval // This is the critical fix - without it, the status would remain Active - let post_merge_validator_list = stake_pool_accounts.get_validator_list(&mut context.banks_client).await; + let post_merge_validator_list = stake_pool_accounts + .get_validator_list(&mut context.banks_client) + .await; let post_merge_validator_info = &post_merge_validator_list.validators[0]; assert_eq!( - post_merge_validator_info.status, + post_merge_validator_info.status, StakeStatus::ReadyForRemoval.into(), "Validator status should be updated to ReadyForRemoval after merging Initialized stake" ); - + // Verify the funds were properly merged into reserve let post_merge_reserve = context .banks_client @@ -864,14 +875,14 @@ async fn updates_validator_status_after_cluster_restart_merge() { post_merge_reserve.lamports > reserve_lamports, "Reserve should have absorbed the Initialized stake funds" ); - + // Verify active stake lamports are 0 (since account was merged) assert_eq!( u64::from(post_merge_validator_info.active_stake_lamports), 0, "Active stake lamports should be 0 after merging Initialized account" ); - + // This test proves that Fix #2 works: "update the validator status correctly after merging the inactive stake into the reserve" } @@ -890,34 +901,40 @@ async fn ignores_unusable_stake_accounts_preventing_exploit() { ) = setup(num_validators).await; let validator_stake_account = &stake_accounts[0]; - + // Verify the validator starts as Active with proper stake - let initial_validator_list = stake_pool_accounts.get_validator_list(&mut context.banks_client).await; + let initial_validator_list = stake_pool_accounts + .get_validator_list(&mut context.banks_client) + .await; let initial_validator_info = &initial_validator_list.validators[0]; assert_eq!(initial_validator_info.status, StakeStatus::Active.into()); assert!(u64::from(initial_validator_info.active_stake_lamports) > 0); - let initial_stake_pool = stake_pool_accounts.get_stake_pool(&mut context.banks_client).await; - + let initial_stake_pool = stake_pool_accounts + .get_stake_pool(&mut context.banks_client) + .await; + // Now simulate the attack - malicious actor takes control of the stake account // This could happen after a cluster restart where the account gets reset to Initialized // and then the attacker manages to gain control before the pool processes it let malicious_authority = Keypair::new(); let extra_malicious_lamports = 1_000_000_000; // 1 SOL extra - let stake_pool = stake_pool_accounts.get_stake_pool(&mut context.banks_client).await; - + let stake_pool = stake_pool_accounts + .get_stake_pool(&mut context.banks_client) + .await; + let malicious_meta = solana_stake_interface::state::Meta { rent_exempt_reserve: 2282880, authorized: solana_stake_interface::state::Authorized { - staker: malicious_authority.pubkey(), // WRONG - should be pool's withdraw authority + staker: malicious_authority.pubkey(), // WRONG - should be pool's withdraw authority withdrawer: malicious_authority.pubkey(), // WRONG - should be pool's withdraw authority }, lockup: solana_stake_interface::state::Lockup { - custodian: malicious_authority.pubkey(), // WRONG - different from pool's lockup + custodian: malicious_authority.pubkey(), // WRONG - different from pool's lockup epoch: stake_pool.lockup.epoch + 100, // WRONG - different epoch ..stake_pool.lockup }, }; - + // Create a malicious delegation with the original validator + extra stake let malicious_delegation = solana_stake_interface::state::Delegation { voter_pubkey: validator_stake_account.vote.pubkey(), @@ -926,18 +943,18 @@ async fn ignores_unusable_stake_accounts_preventing_exploit() { deactivation_epoch: u64::MAX, ..Default::default() }; - + let malicious_stake = solana_stake_interface::state::Stake { delegation: malicious_delegation, credits_observed: 0, }; - + let malicious_stake_state = StakeStateV2::Stake( - malicious_meta, - malicious_stake, - solana_stake_interface::stake_flags::StakeFlags::empty() + malicious_meta, + malicious_stake, + solana_stake_interface::stake_flags::StakeFlags::empty(), ); - + // Get original stake account for comparison let original_stake_account = context .banks_client @@ -946,13 +963,16 @@ async fn ignores_unusable_stake_accounts_preventing_exploit() { .unwrap() .unwrap(); let original_lamports = original_stake_account.lamports; - + // Replace the legitimate stake account with the malicious one let mut malicious_account = original_stake_account.clone(); malicious_account.lamports = original_lamports + extra_malicious_lamports; malicious_account.data = bincode::serialize(&malicious_stake_state).unwrap(); - context.set_account(&validator_stake_account.stake_account, &malicious_account.into()); - + context.set_account( + &validator_stake_account.stake_account, + &malicious_account.into(), + ); + // Run update_validator_list_balance - this should succeed butignore the malicious account let error = stake_pool_accounts .update_validator_list_balance( @@ -963,11 +983,17 @@ async fn ignores_unusable_stake_accounts_preventing_exploit() { false, ) .await; - assert!(error.is_none(), "Update should succeed but ignore malicious account: {:?}", error); - - let final_validator_list = stake_pool_accounts.get_validator_list(&mut context.banks_client).await; + assert!( + error.is_none(), + "Update should succeed but ignore malicious account: {:?}", + error + ); + + let final_validator_list = stake_pool_accounts + .get_validator_list(&mut context.banks_client) + .await; let final_validator_info = &final_validator_list.validators[0]; - + // The validator should still be marked as Active but with 0 active stake // because the malicious account was ignored assert_eq!(final_validator_info.status, StakeStatus::Active.into()); @@ -976,7 +1002,7 @@ async fn ignores_unusable_stake_accounts_preventing_exploit() { 0, "Active stake lamports should be 0 because malicious account is ignored" ); - + // The malicious stake account should still exist with all its lamports // (proving it was ignored, not merged or processed) let final_malicious_account = context @@ -985,17 +1011,22 @@ async fn ignores_unusable_stake_accounts_preventing_exploit() { .await .unwrap() .unwrap(); - + assert_eq!( final_malicious_account.lamports, original_lamports + extra_malicious_lamports, "Malicious account should retain all its lamports since it was ignored" ); - + // Verify the pool's total assets remained unchanged (malicious account was ignored) - let final_stake_pool = stake_pool_accounts.get_stake_pool(&mut context.banks_client).await; - assert_eq!(final_stake_pool.total_lamports, initial_stake_pool.total_lamports); - + let final_stake_pool = stake_pool_accounts + .get_stake_pool(&mut context.banks_client) + .await; + assert_eq!( + final_stake_pool.total_lamports, + initial_stake_pool.total_lamports + ); + // This test proves fix for "only count active validator stakes if they're usable by the pool" } @@ -1014,13 +1045,15 @@ async fn update_validator_list_balance_ingores_uninitialized_stake_account_balan ) = setup(num_validators).await; let validator_stake_account = &stake_accounts[0]; - + // Verify the validator starts as Active - let initial_validator_list = stake_pool_accounts.get_validator_list(&mut context.banks_client).await; + let initial_validator_list = stake_pool_accounts + .get_validator_list(&mut context.banks_client) + .await; let initial_validator_info = &initial_validator_list.validators[0]; assert_eq!(initial_validator_info.status, StakeStatus::Active.into()); assert!(u64::from(initial_validator_info.active_stake_lamports) > 0); - + // First, remove the validator from the pool to trigger deactivation let error = stake_pool_accounts .remove_validator_from_pool( @@ -1032,38 +1065,46 @@ async fn update_validator_list_balance_ingores_uninitialized_stake_account_balan ) .await; assert!(error.is_none(), "Failed to remove validator: {:?}", error); - + // Verify validator is now being deactivated - let deactivating_validator_list = stake_pool_accounts.get_validator_list(&mut context.banks_client).await; + let deactivating_validator_list = stake_pool_accounts + .get_validator_list(&mut context.banks_client) + .await; let deactivating_validator_info = &deactivating_validator_list.validators[0]; - assert_eq!(deactivating_validator_info.status, StakeStatus::DeactivatingValidator.into()); - + assert_eq!( + deactivating_validator_info.status, + StakeStatus::DeactivatingValidator.into() + ); + // Fast forward one epoch to allow the stake to deactivate let slots_per_epoch = context.genesis_config().epoch_schedule.slots_per_epoch; slot += slots_per_epoch; context.warp_to_slot(slot).unwrap(); - + let new_blockhash = context .banks_client .get_new_latest_blockhash(&last_blockhash) .await .unwrap(); - + // Now simulate a scenario where the stake account becomes uninitialized // This could happen due to cluster restart or other network events let uninitialized_stake_state = StakeStateV2::Uninitialized; - + let original_stake_account = context .banks_client .get_account(validator_stake_account.stake_account) .await .unwrap() .unwrap(); - + let mut modified_account = original_stake_account.clone(); modified_account.data = bincode::serialize(&uninitialized_stake_state).unwrap(); - context.set_account(&validator_stake_account.stake_account, &modified_account.into()); - + context.set_account( + &validator_stake_account.stake_account, + &modified_account.into(), + ); + // Run update_validator_list_balance - this should trigger the uninitialized account handling let error = stake_pool_accounts .update_validator_list_balance( @@ -1074,21 +1115,30 @@ async fn update_validator_list_balance_ingores_uninitialized_stake_account_balan false, ) .await; - assert!(error.is_none(), "Update should succeed despite uninitialized account: {:?}", error); - + assert!( + error.is_none(), + "Update should succeed despite uninitialized account: {:?}", + error + ); + // Verify the validator status and that the uninitialized account was ignored - let final_validator_list = stake_pool_accounts.get_validator_list(&mut context.banks_client).await; + let final_validator_list = stake_pool_accounts + .get_validator_list(&mut context.banks_client) + .await; let final_validator_info = &final_validator_list.validators[0]; - + // The validator should still be in DeactivatingValidator status with 0 active stake // because the uninitialized account was ignored - assert_eq!(final_validator_info.status, StakeStatus::DeactivatingValidator.into()); + assert_eq!( + final_validator_info.status, + StakeStatus::DeactivatingValidator.into() + ); assert_eq!( u64::from(final_validator_info.active_stake_lamports), 0, "Active stake lamports should be 0 because uninitialized account is ignored" ); - + // Verify the uninitialized stake account still exists let final_stake_account = context .banks_client @@ -1096,14 +1146,14 @@ async fn update_validator_list_balance_ingores_uninitialized_stake_account_balan .await .unwrap() .unwrap(); - + // Verify it's still uninitialized let final_stake_state: StakeStateV2 = bincode::deserialize(&final_stake_account.data).unwrap(); matches!(final_stake_state, StakeStateV2::Uninitialized); // assert it has a nonzero account balance assert!(final_stake_account.lamports > 0); - + // This test proves that the uninitialized account scenario can be triggered and is handled correctly } @@ -1136,10 +1186,15 @@ async fn cleanup_does_not_remove_validators_with_remaining_lamports() { } // Verify both validators are being deactivated (DeactivatingValidator status) - let validator_list = stake_pool_accounts.get_validator_list(&mut context.banks_client).await; + let validator_list = stake_pool_accounts + .get_validator_list(&mut context.banks_client) + .await; assert_eq!(validator_list.validators.len(), 2); for validator_info in &validator_list.validators { - assert_eq!(validator_info.status, StakeStatus::DeactivatingValidator.into()); + assert_eq!( + validator_info.status, + StakeStatus::DeactivatingValidator.into() + ); } // Fast forward one epoch to allow the deactivating stakes to become inactive @@ -1162,7 +1217,9 @@ async fn cleanup_does_not_remove_validators_with_remaining_lamports() { assert!(error.is_none(), "Update should succeed: {:?}", error); // Verify both validators are now ReadyForRemoval - let updated_validator_list = stake_pool_accounts.get_validator_list(&mut context.banks_client).await; + let updated_validator_list = stake_pool_accounts + .get_validator_list(&mut context.banks_client) + .await; assert_eq!(updated_validator_list.validators.len(), 2); for validator_info in &updated_validator_list.validators { assert_eq!(validator_info.status, StakeStatus::ReadyForRemoval.into()); @@ -1183,17 +1240,40 @@ async fn cleanup_does_not_remove_validators_with_remaining_lamports() { let mut serialized_data = borsh::to_vec(&modified_validator_list).unwrap(); serialized_data.resize(modified_account.data.len(), 0); modified_account.data = serialized_data; - context.set_account(&stake_pool_accounts.validator_list.pubkey(), &modified_account.into()); + context.set_account( + &stake_pool_accounts.validator_list.pubkey(), + &modified_account.into(), + ); // Verify the validator setup is as intended (2 validators ready for removal, one with remaining lamports) - let pre_cleanup_validator_list = stake_pool_accounts.get_validator_list(&mut context.banks_client).await; + let pre_cleanup_validator_list = stake_pool_accounts + .get_validator_list(&mut context.banks_client) + .await; assert_eq!(pre_cleanup_validator_list.validators.len(), 2); - assert_eq!(pre_cleanup_validator_list.validators[0].status, StakeStatus::ReadyForRemoval.into()); - assert_eq!(u64::from(pre_cleanup_validator_list.validators[0].active_stake_lamports), 1_000_000); - assert_eq!(u64::from(pre_cleanup_validator_list.validators[0].transient_stake_lamports), 0); - assert_eq!(pre_cleanup_validator_list.validators[1].status, StakeStatus::ReadyForRemoval.into()); - assert_eq!(u64::from(pre_cleanup_validator_list.validators[1].active_stake_lamports), 0); - assert_eq!(u64::from(pre_cleanup_validator_list.validators[1].transient_stake_lamports), 0); + assert_eq!( + pre_cleanup_validator_list.validators[0].status, + StakeStatus::ReadyForRemoval.into() + ); + assert_eq!( + u64::from(pre_cleanup_validator_list.validators[0].active_stake_lamports), + 1_000_000 + ); + assert_eq!( + u64::from(pre_cleanup_validator_list.validators[0].transient_stake_lamports), + 0 + ); + assert_eq!( + pre_cleanup_validator_list.validators[1].status, + StakeStatus::ReadyForRemoval.into() + ); + assert_eq!( + u64::from(pre_cleanup_validator_list.validators[1].active_stake_lamports), + 0 + ); + assert_eq!( + u64::from(pre_cleanup_validator_list.validators[1].transient_stake_lamports), + 0 + ); // Run cleanup_removed_validator_entries let error = stake_pool_accounts @@ -1206,15 +1286,17 @@ async fn cleanup_does_not_remove_validators_with_remaining_lamports() { assert!(error.is_none(), "Cleanup should succeed: {:?}", error); // Verify that only the validator with 0 lamports was removed - let post_cleanup_validator_list = stake_pool_accounts.get_validator_list(&mut context.banks_client).await; - + let post_cleanup_validator_list = stake_pool_accounts + .get_validator_list(&mut context.banks_client) + .await; + // Should have 1 validator remaining (the one with active lamports) assert_eq!( - post_cleanup_validator_list.validators.len(), + post_cleanup_validator_list.validators.len(), 1, "Only validator with 0 lamports should have been removed" ); - + // The remaining validator should be the one with active lamports assert_eq!( u64::from(post_cleanup_validator_list.validators[0].active_stake_lamports), From 52ff256c9a6988a740d062e1b07d1ac079b9ac0f Mon Sep 17 00:00:00 2001 From: juliaaschmidt Date: Mon, 15 Dec 2025 13:47:47 +0100 Subject: [PATCH 14/14] fix: rustfmt + clippy program --- .../tests/update_validator_list_balance.rs | 1 - program/tests/vsa_remove.rs | 15 +- program/tests/withdraw_edge_cases.rs | 348 +++++++++++------- 3 files changed, 232 insertions(+), 132 deletions(-) diff --git a/program/tests/update_validator_list_balance.rs b/program/tests/update_validator_list_balance.rs index cdc0fc2b..8e64e7b1 100644 --- a/program/tests/update_validator_list_balance.rs +++ b/program/tests/update_validator_list_balance.rs @@ -2,7 +2,6 @@ mod helpers; use { - bincode, borsh, helpers::*, solana_program::{ borsh1::try_from_slice_unchecked, instruction::InstructionError, program_pack::Pack, diff --git a/program/tests/vsa_remove.rs b/program/tests/vsa_remove.rs index dfc9f5c8..b7cd7d37 100644 --- a/program/tests/vsa_remove.rs +++ b/program/tests/vsa_remove.rs @@ -1017,8 +1017,10 @@ async fn success_remove_validator_with_transient_stake_triggers_deactivating_all ); // Verify both active and transient stake accounts are being deactivated - let active_stake_account = get_account(&mut context.banks_client, &validator_stake.stake_account).await; - let active_stake_state = deserialize::(&active_stake_account.data).unwrap(); + let active_stake_account = + get_account(&mut context.banks_client, &validator_stake.stake_account).await; + let active_stake_state = + deserialize::(&active_stake_account.data).unwrap(); if let stake::state::StakeStateV2::Stake(_, active_stake, _) = active_stake_state { assert_ne!( active_stake.delegation.deactivation_epoch, @@ -1027,8 +1029,13 @@ async fn success_remove_validator_with_transient_stake_triggers_deactivating_all ); } - let transient_stake_account = get_account(&mut context.banks_client, &validator_stake.transient_stake_account).await; - let transient_stake_state = deserialize::(&transient_stake_account.data).unwrap(); + let transient_stake_account = get_account( + &mut context.banks_client, + &validator_stake.transient_stake_account, + ) + .await; + let transient_stake_state = + deserialize::(&transient_stake_account.data).unwrap(); if let stake::state::StakeStateV2::Stake(_, transient_stake, _) = transient_stake_state { assert_ne!( transient_stake.delegation.deactivation_epoch, diff --git a/program/tests/withdraw_edge_cases.rs b/program/tests/withdraw_edge_cases.rs index 39f78383..c38057b2 100644 --- a/program/tests/withdraw_edge_cases.rs +++ b/program/tests/withdraw_edge_cases.rs @@ -49,11 +49,20 @@ async fn fail_remove_validator_blocked_by_transient_stake() { // Step 2: Check the state after creating transient stake let validator_stake_account_after_decrease = get_account(&mut context.banks_client, &validator_stake.stake_account).await; - let transient_stake_account_after_decrease = - get_account(&mut context.banks_client, &validator_stake.transient_stake_account).await; - - println!("Validator stake after decrease: {} lamports", validator_stake_account_after_decrease.lamports); - println!("Transient stake after decrease: {} lamports", transient_stake_account_after_decrease.lamports); + let transient_stake_account_after_decrease = get_account( + &mut context.banks_client, + &validator_stake.transient_stake_account, + ) + .await; + + println!( + "Validator stake after decrease: {} lamports", + validator_stake_account_after_decrease.lamports + ); + println!( + "Transient stake after decrease: {} lamports", + transient_stake_account_after_decrease.lamports + ); // Step 3: Warp forward to deactivation epoch let first_normal_slot = context.genesis_config().epoch_schedule.first_normal_slot; @@ -79,31 +88,37 @@ async fn fail_remove_validator_blocked_by_transient_stake() { .get_account(validator_stake.transient_stake_account) .await .unwrap(); - - println!(" Validator stake: {} lamports", validator_stake_account_final.lamports); - + + println!( + " Validator stake: {} lamports", + validator_stake_account_final.lamports + ); + // Verify transient stake still exists - let transient_account = transient_stake_account_final.expect("Transient stake account should still exist"); + let transient_account = + transient_stake_account_final.expect("Transient stake account should still exist"); println!(" Transient stake: {} lamports", transient_account.lamports); // Step 6: Try to withdraw ALL active stake - this should FAIL because transient stake blocks removal let rent = context.banks_client.get_rent().await.unwrap(); let stake_rent = rent.minimum_balance(std::mem::size_of::()); - + let stake_pool = stake_pool_accounts .get_stake_pool(&mut context.banks_client) .await; - + // Try to withdraw everything except rent (this should trigger the validator removal check) - let active_stake_to_withdraw = validator_stake_account_final.lamports + let active_stake_to_withdraw = validator_stake_account_final + .lamports .saturating_sub(stake_rent); - + println!("Testing complete withdrawal of active stake {} lamports (should fail due to transient stake)", active_stake_to_withdraw); - + let pool_tokens_all = (active_stake_to_withdraw * stake_pool.pool_token_supply) .checked_div(stake_pool.total_lamports) .unwrap(); - let pool_tokens_all_with_fee = stake_pool_accounts.calculate_inverse_withdrawal_fee(pool_tokens_all); + let pool_tokens_all_with_fee = + stake_pool_accounts.calculate_inverse_withdrawal_fee(pool_tokens_all); let new_user_authority_all = Pubkey::new_unique(); let error_all = stake_pool_accounts @@ -121,10 +136,16 @@ async fn fail_remove_validator_blocked_by_transient_stake() { .await; // Step 7: Verify that the complete withdrawal fails because transient stake prevents validator removal - assert!(error_all.is_some(), "Complete withdrawal should fail because validator has transient stake"); + assert!( + error_all.is_some(), + "Complete withdrawal should fail because validator has transient stake" + ); let transaction_error = error_all.unwrap().unwrap(); - println!("Complete withdrawal correctly failed with error: {:?}", transaction_error); + println!( + "Complete withdrawal correctly failed with error: {:?}", + transaction_error + ); // The error should be StakeLamportsNotEqualToMinimum because the program detects // that this validator cannot be removed due to associated transient lamports assert_eq!( @@ -141,10 +162,16 @@ async fn fail_remove_validator_blocked_by_transient_stake() { .get_account(validator_stake.transient_stake_account) .await .unwrap(); - - assert!(final_transient_check.is_some(), "Transient stake account should still exist"); + + assert!( + final_transient_check.is_some(), + "Transient stake account should still exist" + ); let final_transient = final_transient_check.unwrap(); - assert!(final_transient.lamports > 0, "Transient stake should still have lamports"); + assert!( + final_transient.lamports > 0, + "Transient stake should still have lamports" + ); } #[tokio::test] @@ -465,9 +492,12 @@ async fn fail_with_transient() { let rent = context.banks_client.get_rent().await.unwrap(); let stake_rent = rent.minimum_balance(std::mem::size_of::()); - let stake_minimum_delegation = - stake_get_minimum_delegation(&mut context.banks_client, &context.payer, &context.last_blockhash) - .await; + let stake_minimum_delegation = stake_get_minimum_delegation( + &mut context.banks_client, + &context.payer, + &context.last_blockhash, + ) + .await; // Calculate how much to decrease to leave only rent + minimum delegation in validator stake account // This will create a transient stake account with the decreased amount @@ -477,8 +507,14 @@ async fn fail_with_transient() { let target_validator_lamports = stake_rent + stake_minimum_delegation; let amount_to_decrease = current_validator_lamports - target_validator_lamports; - println!("Current validator stake: {} lamports", current_validator_lamports); - println!("Target validator stake: {} lamports (rent + min delegation)", target_validator_lamports); + println!( + "Current validator stake: {} lamports", + current_validator_lamports + ); + println!( + "Target validator stake: {} lamports (rent + min delegation)", + target_validator_lamports + ); println!("Amount to decrease: {} lamports", amount_to_decrease); // Decrease validator stake, creating transient stake @@ -499,11 +535,20 @@ async fn fail_with_transient() { // Check the state after decrease - validator should have minimum, transient should have the decreased amount let validator_stake_account_after = get_account(&mut context.banks_client, &validator_stake.stake_account).await; - let transient_stake_account = - get_account(&mut context.banks_client, &validator_stake.transient_stake_account).await; - - println!("Validator stake after decrease: {} lamports", validator_stake_account_after.lamports); - println!("Transient stake after decrease: {} lamports", transient_stake_account.lamports); + let transient_stake_account = get_account( + &mut context.banks_client, + &validator_stake.transient_stake_account, + ) + .await; + + println!( + "Validator stake after decrease: {} lamports", + validator_stake_account_after.lamports + ); + println!( + "Transient stake after decrease: {} lamports", + transient_stake_account.lamports + ); // warp forward to deactivation epoch slot += context.genesis_config().epoch_schedule.slots_per_epoch; @@ -529,42 +574,61 @@ async fn fail_with_transient() { // Check final state before withdrawal let validator_stake_account_final = get_account(&mut context.banks_client, &validator_stake.stake_account).await; - + // Check if transient stake account still exists (it might have been merged back) let transient_stake_account_final = context .banks_client .get_account(validator_stake.transient_stake_account) .await .unwrap(); - - println!("Validator stake after epoch change: {} lamports", validator_stake_account_final.lamports); - + + println!( + "Validator stake after epoch change: {} lamports", + validator_stake_account_final.lamports + ); + // Verify transient stake account still exists and has lamports - let transient_account = transient_stake_account_final.expect("Transient stake account should still exist"); - println!("Transient stake after epoch change: {} lamports", transient_account.lamports); - + let transient_account = + transient_stake_account_final.expect("Transient stake account should still exist"); + println!( + "Transient stake after epoch change: {} lamports", + transient_account.lamports + ); + // Calculate pool tokens needed to withdraw EXACTLY the transient stake amount let stake_pool = stake_pool_accounts .get_stake_pool(&mut context.banks_client) .await; - + // We want to withdraw exactly what's in the transient account (all of it) let exact_transient_lamports = transient_account.lamports; - + // Calculate the exact pool tokens needed for this withdrawal amount // Using the same formula as the stake pool: pool_tokens = (lamports * pool_token_supply) / total_lamports let pool_tokens_post_fee = (exact_transient_lamports * stake_pool.pool_token_supply) .checked_div(stake_pool.total_lamports) .unwrap(); let pool_tokens = stake_pool_accounts.calculate_inverse_withdrawal_fee(pool_tokens_post_fee); - - println!("Transient account has exactly: {} lamports", exact_transient_lamports); - println!("Attempting to withdraw exactly: {} lamports from transient stake", exact_transient_lamports); - println!("Pool tokens calculated: {} (post-fee: {})", pool_tokens, pool_tokens_post_fee); - println!("Stake pool total lamports: {}, pool token supply: {}", stake_pool.total_lamports, stake_pool.pool_token_supply); + + println!( + "Transient account has exactly: {} lamports", + exact_transient_lamports + ); + println!( + "Attempting to withdraw exactly: {} lamports from transient stake", + exact_transient_lamports + ); + println!( + "Pool tokens calculated: {} (post-fee: {})", + pool_tokens, pool_tokens_post_fee + ); + println!( + "Stake pool total lamports: {}, pool token supply: {}", + stake_pool.total_lamports, stake_pool.pool_token_supply + ); let new_user_authority = Pubkey::new_unique(); - + // Try to withdraw from transient stake account - this should FAIL let error = stake_pool_accounts .withdraw_stake( @@ -579,7 +643,7 @@ async fn fail_with_transient() { pool_tokens, ) .await; - + // Verify that the withdrawal fails with the expected error assert!(error.is_some(), "Withdrawal should fail"); let transaction_error = error.unwrap().unwrap(); @@ -1329,30 +1393,38 @@ async fn success_remove_preferred_validator_resets_preference() { _, ) = setup_for_withdraw(spl_token::id(), STAKE_ACCOUNT_RENT_EXEMPTION).await; - let last_blockhash = context.banks_client - .get_new_latest_blockhash(&context.last_blockhash).await + let last_blockhash = context + .banks_client + .get_new_latest_blockhash(&context.last_blockhash) + .await .unwrap(); // Set the validator as the preferred withdraw validator - stake_pool_accounts.set_preferred_validator( - &mut context.banks_client, - &context.payer, - &last_blockhash, - instruction::PreferredValidatorType::Withdraw, - Some(validator_stake.vote.pubkey()) - ).await; + stake_pool_accounts + .set_preferred_validator( + &mut context.banks_client, + &context.payer, + &last_blockhash, + instruction::PreferredValidatorType::Withdraw, + Some(validator_stake.vote.pubkey()), + ) + .await; // Also set it as preferred deposit validator to test both reset paths - stake_pool_accounts.set_preferred_validator( - &mut context.banks_client, - &context.payer, - &last_blockhash, - instruction::PreferredValidatorType::Deposit, - Some(validator_stake.vote.pubkey()) - ).await; + stake_pool_accounts + .set_preferred_validator( + &mut context.banks_client, + &context.payer, + &last_blockhash, + instruction::PreferredValidatorType::Deposit, + Some(validator_stake.vote.pubkey()), + ) + .await; // Verify the preferred deposit and withdraw validators are set - let stake_pool_before = stake_pool_accounts.get_stake_pool(&mut context.banks_client).await; + let stake_pool_before = stake_pool_accounts + .get_stake_pool(&mut context.banks_client) + .await; assert_eq!( stake_pool_before.preferred_withdraw_validator_vote_address, Some(validator_stake.vote.pubkey()) @@ -1362,67 +1434,76 @@ async fn success_remove_preferred_validator_resets_preference() { Some(validator_stake.vote.pubkey()) ); - println!("Preferred validators set to: {}", validator_stake.vote.pubkey()); + println!( + "Preferred validators set to: {}", + validator_stake.vote.pubkey() + ); // Warp forward to after reward payout let first_normal_slot = context.genesis_config().epoch_schedule.first_normal_slot; let mut slot = first_normal_slot + 1; context.warp_to_slot(slot).unwrap(); - let error = stake_pool_accounts.update_all( - &mut context.banks_client, - &context.payer, - &last_blockhash, - false - ).await; + let error = stake_pool_accounts + .update_all( + &mut context.banks_client, + &context.payer, + &last_blockhash, + false, + ) + .await; assert!(error.is_none(), "{:?}", error); let rent = context.banks_client.get_rent().await.unwrap(); let stake_rent = rent.minimum_balance(std::mem::size_of::()); - let stake_pool = stake_pool_accounts.get_stake_pool(&mut context.banks_client).await; + let stake_pool = stake_pool_accounts + .get_stake_pool(&mut context.banks_client) + .await; let lamports_per_pool_token = stake_pool.get_lamports_per_pool_token().unwrap(); // Decrease all of stake except for exactly lamports_per_pool_token lamports // This will leave the minimum amount that can be withdrawn completely - let error = stake_pool_accounts.decrease_validator_stake_either( - &mut context.banks_client, - &context.payer, - &last_blockhash, - &validator_stake.stake_account, - &validator_stake.transient_stake_account, - deposit_info.stake_lamports + stake_rent - lamports_per_pool_token, - validator_stake.transient_stake_seed, - DecreaseInstruction::Reserve - ).await; + let error = stake_pool_accounts + .decrease_validator_stake_either( + &mut context.banks_client, + &context.payer, + &last_blockhash, + &validator_stake.stake_account, + &validator_stake.transient_stake_account, + deposit_info.stake_lamports + stake_rent - lamports_per_pool_token, + validator_stake.transient_stake_seed, + DecreaseInstruction::Reserve, + ) + .await; assert!(error.is_none(), "{:?}", error); // Warp forward to deactivation slot += context.genesis_config().epoch_schedule.slots_per_epoch; context.warp_to_slot(slot).unwrap(); - let last_blockhash = context.banks_client - .get_new_latest_blockhash(&last_blockhash).await + let last_blockhash = context + .banks_client + .get_new_latest_blockhash(&last_blockhash) + .await .unwrap(); // Update to merge deactivated stake into reserve - let error = stake_pool_accounts.update_all( - &mut context.banks_client, - &context.payer, - &last_blockhash, - false - ).await; + let error = stake_pool_accounts + .update_all( + &mut context.banks_client, + &context.payer, + &last_blockhash, + false, + ) + .await; assert!(error.is_none(), "{:?}", error); - let validator_stake_account = get_account( - &mut context.banks_client, - &validator_stake.stake_account - ).await; + let validator_stake_account = + get_account(&mut context.banks_client, &validator_stake.stake_account).await; let remaining_lamports = validator_stake_account.lamports; - let stake_minimum_delegation = stake_get_minimum_delegation( - &mut context.banks_client, - &context.payer, - &last_blockhash - ).await; + let stake_minimum_delegation = + stake_get_minimum_delegation(&mut context.banks_client, &context.payer, &last_blockhash) + .await; println!("Remaining lamports in validator: {}", remaining_lamports); println!("Stake rent: {}", stake_rent); @@ -1431,57 +1512,70 @@ async fn success_remove_preferred_validator_resets_preference() { assert!(remaining_lamports > stake_rent + stake_minimum_delegation); // Calculate pool tokens needed to withdraw everything (this should remove the validator completely) - let stake_pool_updated = stake_pool_accounts.get_stake_pool(&mut context.banks_client).await; - let pool_tokens_post_fee = (remaining_lamports * stake_pool_updated.pool_token_supply).div_ceil( - stake_pool_updated.total_lamports - ); + let stake_pool_updated = stake_pool_accounts + .get_stake_pool(&mut context.banks_client) + .await; + let pool_tokens_post_fee = (remaining_lamports * stake_pool_updated.pool_token_supply) + .div_ceil(stake_pool_updated.total_lamports); let pool_tokens = stake_pool_accounts.calculate_inverse_withdrawal_fee(pool_tokens_post_fee); - println!("Pool tokens needed for complete withdrawal: {}", pool_tokens); + println!( + "Pool tokens needed for complete withdrawal: {}", + pool_tokens + ); let new_user_authority = Pubkey::new_unique(); // Perform the complete withdrawal - this should trigger validator removal and reset preferred validators - let error = stake_pool_accounts.withdraw_stake( - &mut context.banks_client, - &context.payer, - &last_blockhash, - &user_stake_recipient.pubkey(), - &user_transfer_authority, - &deposit_info.pool_account.pubkey(), - &validator_stake.stake_account, - &new_user_authority, - pool_tokens - ).await; - assert!(error.is_none(), "Complete withdrawal should succeed: {:?}", error); + let error = stake_pool_accounts + .withdraw_stake( + &mut context.banks_client, + &context.payer, + &last_blockhash, + &user_stake_recipient.pubkey(), + &user_transfer_authority, + &deposit_info.pool_account.pubkey(), + &validator_stake.stake_account, + &new_user_authority, + pool_tokens, + ) + .await; + assert!( + error.is_none(), + "Complete withdrawal should succeed: {:?}", + error + ); println!("Complete withdrawal successful - validator should be removed"); // Verify validator stake account is gone - let validator_stake_account_after = context.banks_client - .get_account(validator_stake.stake_account).await + let validator_stake_account_after = context + .banks_client + .get_account(validator_stake.stake_account) + .await .unwrap(); - assert!(validator_stake_account_after.is_none(), "Validator stake account should be removed"); + assert!( + validator_stake_account_after.is_none(), + "Validator stake account should be removed" + ); // Verify that preferred validators have been reset to None - let stake_pool_after = stake_pool_accounts.get_stake_pool(&mut context.banks_client).await; + let stake_pool_after = stake_pool_accounts + .get_stake_pool(&mut context.banks_client) + .await; assert_eq!( - stake_pool_after.preferred_withdraw_validator_vote_address, - None, + stake_pool_after.preferred_withdraw_validator_vote_address, None, "Preferred withdraw validator should be reset to None" ); assert_eq!( - stake_pool_after.preferred_deposit_validator_vote_address, - None, + stake_pool_after.preferred_deposit_validator_vote_address, None, "Preferred deposit validator should be reset to None" ); // Verify user received the stake - let user_stake_recipient_account = get_account( - &mut context.banks_client, - &user_stake_recipient.pubkey() - ).await; + let user_stake_recipient_account = + get_account(&mut context.banks_client, &user_stake_recipient.pubkey()).await; assert_eq!( user_stake_recipient_account.lamports, remaining_lamports + stake_rent,