diff --git a/transaction-status/src/token_balances.rs b/transaction-status/src/token_balances.rs index fc4f2ad9e..649aaa364 100644 --- a/transaction-status/src/token_balances.rs +++ b/transaction-status/src/token_balances.rs @@ -41,6 +41,10 @@ fn get_mint_decimals(bank: &Bank, mint: &Pubkey) -> Option { } else { let mint_account = bank.get_account(mint)?; + if !is_known_spl_token_id(&mint_account.owner()) { + return None; + } + let decimals = Mint::unpack(mint_account.data()) .map(|mint| mint.decimals) .ok()?; @@ -93,6 +97,7 @@ pub fn collect_token_balances( balances } +#[derive(Debug, PartialEq)] struct TokenBalanceData { mint: String, owner: String, @@ -106,6 +111,10 @@ fn collect_token_balance_from_account( ) -> Option { let account = bank.get_account(account_id)?; + if !is_known_spl_token_id(&account.owner()) { + return None; + } + let token_account = TokenAccount::unpack(account.data()).ok()?; let mint = pubkey_from_spl_token(&token_account.mint); @@ -121,3 +130,158 @@ fn collect_token_balance_from_account( ui_token_amount: token_amount_to_ui_amount(token_account.amount, decimals), }) } + +#[cfg(test)] +mod test { + use { + super::*, + solana_account_decoder::parse_token::{pubkey_from_spl_token, spl_token_pubkey}, + solana_sdk::{account::Account, genesis_config::create_genesis_config}, + spl_token::solana_program::program_option::COption, + std::collections::BTreeMap, + }; + + #[test] + fn test_collect_token_balance_from_account() { + let (mut genesis_config, _mint_keypair) = create_genesis_config(500); + + // Add a variety of accounts, token and not + let account = Account::new(42, 55, &Pubkey::new_unique()); + + let mint_data = Mint { + mint_authority: COption::None, + supply: 4242, + decimals: 2, + is_initialized: true, + freeze_authority: COption::None, + }; + let mut data = [0; Mint::LEN]; + Mint::pack(mint_data, &mut data).unwrap(); + let mint_pubkey = Pubkey::new_unique(); + let mint = Account { + lamports: 100, + data: data.to_vec(), + owner: pubkey_from_spl_token(&spl_token::id()), + executable: false, + rent_epoch: 0, + }; + let other_mint_pubkey = Pubkey::new_unique(); + let other_mint = Account { + lamports: 100, + data: data.to_vec(), + owner: Pubkey::new_unique(), + executable: false, + rent_epoch: 0, + }; + + let token_owner = Pubkey::new_unique(); + let token_data = TokenAccount { + mint: spl_token_pubkey(&mint_pubkey), + owner: spl_token_pubkey(&token_owner), + amount: 42, + delegate: COption::None, + state: spl_token::state::AccountState::Initialized, + is_native: COption::Some(100), + delegated_amount: 0, + close_authority: COption::None, + }; + let mut data = [0; TokenAccount::LEN]; + TokenAccount::pack(token_data, &mut data).unwrap(); + + let spl_token_account = Account { + lamports: 100, + data: data.to_vec(), + owner: pubkey_from_spl_token(&spl_token::id()), + executable: false, + rent_epoch: 0, + }; + let other_account = Account { + lamports: 100, + data: data.to_vec(), + owner: Pubkey::new_unique(), + executable: false, + rent_epoch: 0, + }; + + let other_mint_data = TokenAccount { + mint: spl_token_pubkey(&other_mint_pubkey), + owner: spl_token_pubkey(&token_owner), + amount: 42, + delegate: COption::None, + state: spl_token::state::AccountState::Initialized, + is_native: COption::Some(100), + delegated_amount: 0, + close_authority: COption::None, + }; + let mut data = [0; TokenAccount::LEN]; + TokenAccount::pack(other_mint_data, &mut data).unwrap(); + + let other_mint_token_account = Account { + lamports: 100, + data: data.to_vec(), + owner: pubkey_from_spl_token(&spl_token::id()), + executable: false, + rent_epoch: 0, + }; + + let mut accounts = BTreeMap::new(); + + let account_pubkey = Pubkey::new_unique(); + accounts.insert(account_pubkey, account); + accounts.insert(mint_pubkey, mint); + accounts.insert(other_mint_pubkey, other_mint); + let spl_token_account_pubkey = Pubkey::new_unique(); + accounts.insert(spl_token_account_pubkey, spl_token_account); + let other_account_pubkey = Pubkey::new_unique(); + accounts.insert(other_account_pubkey, other_account); + let other_mint_account_pubkey = Pubkey::new_unique(); + accounts.insert(other_mint_account_pubkey, other_mint_token_account); + + genesis_config.accounts = accounts; + + let bank = Bank::new_for_tests(&genesis_config); + let mut mint_decimals = HashMap::new(); + + assert_eq!( + collect_token_balance_from_account(&bank, &account_pubkey, &mut mint_decimals), + None + ); + + assert_eq!( + collect_token_balance_from_account(&bank, &mint_pubkey, &mut mint_decimals), + None + ); + + assert_eq!( + collect_token_balance_from_account( + &bank, + &spl_token_account_pubkey, + &mut mint_decimals + ), + Some(TokenBalanceData { + mint: mint_pubkey.to_string(), + owner: token_owner.to_string(), + ui_token_amount: UiTokenAmount { + ui_amount: Some(0.42), + decimals: 2, + amount: "42".to_string(), + ui_amount_string: "0.42".to_string(), + } + }) + ); + + assert_eq!( + collect_token_balance_from_account(&bank, &other_account_pubkey, &mut mint_decimals), + None + ); + + assert_eq!( + collect_token_balance_from_account( + &bank, + &other_mint_account_pubkey, + &mut mint_decimals + ), + None + ); + } +}