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 + ); +} diff --git a/program/tests/update_validator_list_balance.rs b/program/tests/update_validator_list_balance.rs index c05defef..8e64e7b1 100644 --- a/program/tests/update_validator_list_balance.rs +++ b/program/tests/update_validator_list_balance.rs @@ -8,8 +8,12 @@ use { }, 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_pod::primitives::PodU64, spl_stake_pool::{ error::StakePoolError, state::{StakePool, StakeStatus, ValidatorList}, @@ -779,3 +783,528 @@ 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" +} + +#[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; + 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 + 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 + 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" + ); +} diff --git a/program/tests/vsa_remove.rs b/program/tests/vsa_remove.rs index 21b07f58..b7cd7d37 100644 --- a/program/tests/vsa_remove.rs +++ b/program/tests/vsa_remove.rs @@ -948,3 +948,99 @@ 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" + ); + } +} diff --git a/program/tests/withdraw_edge_cases.rs b/program/tests/withdraw_edge_cases.rs index 3fcd063d..c38057b2 100644 --- a/program/tests/withdraw_edge_cases.rs +++ b/program/tests/withdraw_edge_cases.rs @@ -18,6 +18,162 @@ 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, + _, + ) = 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, + &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 +463,199 @@ 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, + _, + ) = 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, + &user_stake_recipient.pubkey(), + &user_transfer_authority, + &deposit_info.pool_account.pubkey(), + &validator_stake.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 ( @@ -1031,3 +1380,205 @@ 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, + _, + ) = 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, + &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" + ); +}