From e168a5db92f4048879755fae2bc283bb7a3ec9bc Mon Sep 17 00:00:00 2001 From: jon-chuang <9093549+jon-chuang@users.noreply.github.com> Date: Wed, 18 Aug 2021 04:32:47 +0800 Subject: [PATCH] stake-pool: skip manager fee if invalid account, check new manager fee owned by `spl_token` (#2277) * fix * add check if new manager fee is spl token program * Convert manager fee info check to fail only inside `Result` box * update also checks validity of manager fee info * clippy --- stake-pool/program/src/processor.rs | 19 ++++- stake-pool/program/src/state.rs | 18 +++++ stake-pool/program/tests/helpers/mod.rs | 48 ++++++++++++ stake-pool/program/tests/withdraw.rs | 97 ++++++++++++++++++++++--- 4 files changed, 167 insertions(+), 15 deletions(-) diff --git a/stake-pool/program/src/processor.rs b/stake-pool/program/src/processor.rs index 1ccb1a3a..e4186f12 100644 --- a/stake-pool/program/src/processor.rs +++ b/stake-pool/program/src/processor.rs @@ -1642,9 +1642,15 @@ impl Processor { } let reward_lamports = total_stake_lamports.saturating_sub(previous_lamports); - let fee = stake_pool - .calc_epoch_fee_amount(reward_lamports) - .ok_or(StakePoolError::CalculationFailure)?; + + // If the manager fee info is invalid, they don't deserve to receive the fee. + let fee = if stake_pool.check_manager_fee_info(manager_fee_info).is_ok() { + stake_pool + .calc_epoch_fee_amount(reward_lamports) + .ok_or(StakePoolError::CalculationFailure)? + } else { + 0 + }; if fee > 0 { Self::token_mint_to( @@ -2205,7 +2211,11 @@ impl Processor { return Err(StakePoolError::InvalidState.into()); } - let pool_tokens_fee = if stake_pool.manager_fee_account == *burn_from_pool_info.key { + // To prevent a faulty manager fee account from preventing withdrawals + // if the token program does not own the account, or if the account is not initialized + let pool_tokens_fee = if stake_pool.manager_fee_account == *burn_from_pool_info.key + || stake_pool.check_manager_fee_info(manager_fee_info).is_err() + { 0 } else { stake_pool @@ -2393,6 +2403,7 @@ impl Processor { check_account_owner(stake_pool_info, program_id)?; let mut stake_pool = try_from_slice_unchecked::(&stake_pool_info.data.borrow())?; + check_account_owner(new_manager_fee_info, &stake_pool.token_program_id)?; if !stake_pool.is_valid() { return Err(StakePoolError::InvalidState.into()); } diff --git a/stake-pool/program/src/state.rs b/stake-pool/program/src/state.rs index 96b5c853..e7ec27b8 100644 --- a/stake-pool/program/src/state.rs +++ b/stake-pool/program/src/state.rs @@ -1,5 +1,6 @@ //! State transition types +use spl_token::state::{Account, AccountState}; use { crate::{ big_vec::BigVec, error::StakePoolError, stake_program::Lockup, MAX_WITHDRAWAL_FEE_INCREASE, @@ -261,6 +262,23 @@ impl StakePool { } } + /// Check if the manager fee info is a valid token program account + /// capable of receiving tokens from the mint. + pub(crate) fn check_manager_fee_info( + &self, + manager_fee_info: &AccountInfo, + ) -> Result<(), ProgramError> { + let token_account = Account::unpack(&manager_fee_info.data.borrow())?; + if manager_fee_info.owner != &self.token_program_id + || token_account.state != AccountState::Initialized + || token_account.mint != self.pool_mint + { + msg!("Manager fee account is not owned by token program, is not initialized, or does not match stake pool's mint"); + return Err(StakePoolError::InvalidFeeAccount.into()); + } + Ok(()) + } + /// Checks that the withdraw authority is valid #[inline] pub(crate) fn check_authority_withdraw( diff --git a/stake-pool/program/tests/helpers/mod.rs b/stake-pool/program/tests/helpers/mod.rs index e95ae5c4..c39bbd07 100644 --- a/stake-pool/program/tests/helpers/mod.rs +++ b/stake-pool/program/tests/helpers/mod.rs @@ -160,6 +160,54 @@ pub async fn create_token_account( Ok(()) } +pub async fn close_token_account( + banks_client: &mut BanksClient, + payer: &Keypair, + recent_blockhash: &Hash, + account: &Pubkey, + lamports_destination: &Pubkey, + manager: &Keypair, +) -> Result<(), TransportError> { + let mut transaction = Transaction::new_with_payer( + &[spl_token::instruction::close_account( + &spl_token::id(), + &account, + &lamports_destination, + &manager.pubkey(), + &[], + ) + .unwrap()], + Some(&payer.pubkey()), + ); + transaction.sign(&[payer, manager], *recent_blockhash); + banks_client.process_transaction(transaction).await?; + Ok(()) +} + +pub async fn freeze_token_account( + banks_client: &mut BanksClient, + payer: &Keypair, + recent_blockhash: &Hash, + account: &Pubkey, + pool_mint: &Pubkey, + manager: &Keypair, +) -> Result<(), TransportError> { + let mut transaction = Transaction::new_with_payer( + &[spl_token::instruction::freeze_account( + &spl_token::id(), + &account, + pool_mint, + &manager.pubkey(), + &[], + ) + .unwrap()], + Some(&payer.pubkey()), + ); + transaction.sign(&[payer, manager], *recent_blockhash); + banks_client.process_transaction(transaction).await?; + Ok(()) +} + pub async fn mint_tokens( banks_client: &mut BanksClient, payer: &Keypair, diff --git a/stake-pool/program/tests/withdraw.rs b/stake-pool/program/tests/withdraw.rs index 1af022bf..db34fe55 100644 --- a/stake-pool/program/tests/withdraw.rs +++ b/stake-pool/program/tests/withdraw.rs @@ -102,6 +102,20 @@ async fn setup() -> ( #[tokio::test] async fn success() { + _success(SuccessTestType::Success).await; +} + +#[tokio::test] +async fn success_with_closed_manager_fee_account() { + _success(SuccessTestType::UninitializedManagerFee).await; +} + +enum SuccessTestType { + Success, + UninitializedManagerFee, +} + +async fn _success(test_type: SuccessTestType) { let ( mut banks_client, payer, @@ -148,6 +162,60 @@ async fn success() { ) .await; + let destination_keypair = Keypair::new(); + create_token_account( + &mut banks_client, + &payer, + &recent_blockhash, + &destination_keypair, + &stake_pool_accounts.pool_mint.pubkey(), + &Keypair::new().pubkey(), + ) + .await + .unwrap(); + + if let SuccessTestType::UninitializedManagerFee = test_type { + transfer_spl_tokens( + &mut banks_client, + &payer, + &recent_blockhash, + &stake_pool_accounts.pool_fee_account.pubkey(), + &destination_keypair.pubkey(), + &stake_pool_accounts.manager, + pool_fee_balance_before, + ) + .await; + // Check that the account cannot be frozen due to lack of + // freeze authority. + let transaction_error = freeze_token_account( + &mut banks_client, + &payer, + &recent_blockhash, + &stake_pool_accounts.pool_fee_account.pubkey(), + &stake_pool_accounts.pool_mint.pubkey(), + &stake_pool_accounts.manager, + ) + .await + .unwrap_err(); + + match transaction_error { + TransportError::TransactionError(TransactionError::InstructionError(_, error)) => { + assert_eq!(error, InstructionError::Custom(0x10)); + } + _ => panic!("Wrong error occurs while try to withdraw with wrong stake program ID"), + } + close_token_account( + &mut banks_client, + &payer, + &recent_blockhash, + &stake_pool_accounts.pool_fee_account.pubkey(), + &destination_keypair.pubkey(), + &stake_pool_accounts.manager, + ) + .await + .unwrap(); + } + let new_authority = Pubkey::new_unique(); let error = stake_pool_accounts .withdraw_stake( @@ -169,7 +237,12 @@ async fn success() { let stake_pool = try_from_slice_unchecked::(&stake_pool.data.as_slice()).unwrap(); // first and only deposit, lamports:pool 1:1 - let tokens_withdrawal_fee = stake_pool_accounts.calculate_withdrawal_fee(tokens_to_withdraw); + let tokens_withdrawal_fee = match test_type { + SuccessTestType::Success => { + stake_pool_accounts.calculate_withdrawal_fee(tokens_to_withdraw) + } + _ => 0, + }; let tokens_burnt = tokens_to_withdraw - tokens_withdrawal_fee; assert_eq!( stake_pool.total_stake_lamports, @@ -180,16 +253,18 @@ async fn success() { stake_pool_before.pool_token_supply - tokens_burnt ); - // Check manager received withdrawal fee - let pool_fee_balance = get_token_balance( - &mut banks_client, - &stake_pool_accounts.pool_fee_account.pubkey(), - ) - .await; - assert_eq!( - pool_fee_balance, - pool_fee_balance_before + tokens_withdrawal_fee, - ); + if let SuccessTestType::Success = test_type { + // Check manager received withdrawal fee + let pool_fee_balance = get_token_balance( + &mut banks_client, + &stake_pool_accounts.pool_fee_account.pubkey(), + ) + .await; + assert_eq!( + pool_fee_balance, + pool_fee_balance_before + tokens_withdrawal_fee, + ); + } // Check validator stake list storage let validator_list = get_account(