476 lines
16 KiB
Rust
476 lines
16 KiB
Rust
use {
|
|
solana_account_decoder::parse_token::{
|
|
is_known_spl_token_id, token_amount_to_ui_amount, UiTokenAmount,
|
|
},
|
|
solana_measure::measure::Measure,
|
|
solana_metrics::datapoint_debug,
|
|
solana_runtime::{bank::Bank, transaction_batch::TransactionBatch},
|
|
solana_sdk::{account::ReadableAccount, pubkey::Pubkey},
|
|
solana_transaction_status::{
|
|
token_balances::TransactionTokenBalances, TransactionTokenBalance,
|
|
},
|
|
spl_token_2022::{
|
|
extension::StateWithExtensions,
|
|
state::{Account as TokenAccount, Mint},
|
|
},
|
|
std::collections::HashMap,
|
|
};
|
|
|
|
fn get_mint_decimals(bank: &Bank, mint: &Pubkey) -> Option<u8> {
|
|
if mint == &spl_token::native_mint::id() {
|
|
Some(spl_token::native_mint::DECIMALS)
|
|
} else {
|
|
let mint_account = bank.get_account(mint)?;
|
|
|
|
if !is_known_spl_token_id(mint_account.owner()) {
|
|
return None;
|
|
}
|
|
|
|
let decimals = StateWithExtensions::<Mint>::unpack(mint_account.data())
|
|
.map(|mint| mint.base.decimals)
|
|
.ok()?;
|
|
|
|
Some(decimals)
|
|
}
|
|
}
|
|
|
|
pub fn collect_token_balances(
|
|
bank: &Bank,
|
|
batch: &TransactionBatch,
|
|
mint_decimals: &mut HashMap<Pubkey, u8>,
|
|
) -> TransactionTokenBalances {
|
|
let mut balances: TransactionTokenBalances = vec![];
|
|
let mut collect_time = Measure::start("collect_token_balances");
|
|
|
|
for transaction in batch.sanitized_transactions() {
|
|
let account_keys = transaction.message().account_keys();
|
|
let has_token_program = account_keys.iter().any(is_known_spl_token_id);
|
|
|
|
let mut transaction_balances: Vec<TransactionTokenBalance> = vec![];
|
|
if has_token_program {
|
|
for (index, account_id) in account_keys.iter().enumerate() {
|
|
if transaction.message().is_invoked(index) || is_known_spl_token_id(account_id) {
|
|
continue;
|
|
}
|
|
|
|
if let Some(TokenBalanceData {
|
|
mint,
|
|
ui_token_amount,
|
|
owner,
|
|
program_id,
|
|
}) = collect_token_balance_from_account(bank, account_id, mint_decimals)
|
|
{
|
|
transaction_balances.push(TransactionTokenBalance {
|
|
account_index: index as u8,
|
|
mint,
|
|
ui_token_amount,
|
|
owner,
|
|
program_id,
|
|
});
|
|
}
|
|
}
|
|
}
|
|
balances.push(transaction_balances);
|
|
}
|
|
collect_time.stop();
|
|
datapoint_debug!(
|
|
"collect_token_balances",
|
|
("collect_time_us", collect_time.as_us(), i64),
|
|
);
|
|
balances
|
|
}
|
|
|
|
#[derive(Debug, PartialEq)]
|
|
struct TokenBalanceData {
|
|
mint: String,
|
|
owner: String,
|
|
ui_token_amount: UiTokenAmount,
|
|
program_id: String,
|
|
}
|
|
|
|
fn collect_token_balance_from_account(
|
|
bank: &Bank,
|
|
account_id: &Pubkey,
|
|
mint_decimals: &mut HashMap<Pubkey, u8>,
|
|
) -> Option<TokenBalanceData> {
|
|
let account = bank.get_account(account_id)?;
|
|
|
|
if !is_known_spl_token_id(account.owner()) {
|
|
return None;
|
|
}
|
|
|
|
let token_account = StateWithExtensions::<TokenAccount>::unpack(account.data()).ok()?;
|
|
let mint = token_account.base.mint;
|
|
|
|
let decimals = mint_decimals.get(&mint).cloned().or_else(|| {
|
|
let decimals = get_mint_decimals(bank, &mint)?;
|
|
mint_decimals.insert(mint, decimals);
|
|
Some(decimals)
|
|
})?;
|
|
|
|
Some(TokenBalanceData {
|
|
mint: token_account.base.mint.to_string(),
|
|
owner: token_account.base.owner.to_string(),
|
|
ui_token_amount: token_amount_to_ui_amount(token_account.base.amount, decimals),
|
|
program_id: account.owner().to_string(),
|
|
})
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod test {
|
|
use {
|
|
super::*,
|
|
solana_sdk::{account::Account, genesis_config::create_genesis_config},
|
|
spl_token_2022::{
|
|
extension::{
|
|
immutable_owner::ImmutableOwner, memo_transfer::MemoTransfer,
|
|
mint_close_authority::MintCloseAuthority, ExtensionType, StateWithExtensionsMut,
|
|
},
|
|
pod::OptionalNonZeroPubkey,
|
|
solana_program::{program_option::COption, program_pack::Pack},
|
|
},
|
|
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: 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(), // !is_known_spl_token_id
|
|
executable: false,
|
|
rent_epoch: 0,
|
|
};
|
|
|
|
let token_owner = Pubkey::new_unique();
|
|
let token_data = TokenAccount {
|
|
mint: mint_pubkey,
|
|
owner: token_owner,
|
|
amount: 42,
|
|
delegate: COption::None,
|
|
state: spl_token_2022::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: spl_token::id(),
|
|
executable: false,
|
|
rent_epoch: 0,
|
|
};
|
|
let other_account = Account {
|
|
lamports: 100,
|
|
data: data.to_vec(),
|
|
owner: Pubkey::new_unique(), // !is_known_spl_token_id
|
|
executable: false,
|
|
rent_epoch: 0,
|
|
};
|
|
|
|
let other_mint_data = TokenAccount {
|
|
mint: other_mint_pubkey,
|
|
owner: token_owner,
|
|
amount: 42,
|
|
delegate: COption::None,
|
|
state: spl_token_2022::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: 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();
|
|
|
|
// Account is not owned by spl_token (nor does it have TokenAccount state)
|
|
assert_eq!(
|
|
collect_token_balance_from_account(&bank, &account_pubkey, &mut mint_decimals),
|
|
None
|
|
);
|
|
|
|
// Mint does not have TokenAccount state
|
|
assert_eq!(
|
|
collect_token_balance_from_account(&bank, &mint_pubkey, &mut mint_decimals),
|
|
None
|
|
);
|
|
|
|
// TokenAccount owned by spl_token::id() works
|
|
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(),
|
|
},
|
|
program_id: spl_token::id().to_string(),
|
|
})
|
|
);
|
|
|
|
// TokenAccount is not owned by known spl-token program_id
|
|
assert_eq!(
|
|
collect_token_balance_from_account(&bank, &other_account_pubkey, &mut mint_decimals),
|
|
None
|
|
);
|
|
|
|
// TokenAccount's mint is not owned by known spl-token program_id
|
|
assert_eq!(
|
|
collect_token_balance_from_account(
|
|
&bank,
|
|
&other_mint_account_pubkey,
|
|
&mut mint_decimals
|
|
),
|
|
None
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_collect_token_balance_from_spl_token_2022_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_authority = Pubkey::new_unique();
|
|
let mint_size =
|
|
ExtensionType::get_account_len::<Mint>(&[ExtensionType::MintCloseAuthority]);
|
|
let mint_base = Mint {
|
|
mint_authority: COption::None,
|
|
supply: 4242,
|
|
decimals: 2,
|
|
is_initialized: true,
|
|
freeze_authority: COption::None,
|
|
};
|
|
let mut mint_data = vec![0; mint_size];
|
|
let mut mint_state =
|
|
StateWithExtensionsMut::<Mint>::unpack_uninitialized(&mut mint_data).unwrap();
|
|
mint_state.base = mint_base;
|
|
mint_state.pack_base();
|
|
mint_state.init_account_type().unwrap();
|
|
let mut mint_close_authority = mint_state
|
|
.init_extension::<MintCloseAuthority>(true)
|
|
.unwrap();
|
|
mint_close_authority.close_authority =
|
|
OptionalNonZeroPubkey::try_from(Some(mint_authority)).unwrap();
|
|
|
|
let mint_pubkey = Pubkey::new_unique();
|
|
let mint = Account {
|
|
lamports: 100,
|
|
data: mint_data.to_vec(),
|
|
owner: spl_token_2022::id(),
|
|
executable: false,
|
|
rent_epoch: 0,
|
|
};
|
|
let other_mint_pubkey = Pubkey::new_unique();
|
|
let other_mint = Account {
|
|
lamports: 100,
|
|
data: mint_data.to_vec(),
|
|
owner: Pubkey::new_unique(),
|
|
executable: false,
|
|
rent_epoch: 0,
|
|
};
|
|
|
|
let token_owner = Pubkey::new_unique();
|
|
let token_base = TokenAccount {
|
|
mint: mint_pubkey,
|
|
owner: token_owner,
|
|
amount: 42,
|
|
delegate: COption::None,
|
|
state: spl_token_2022::state::AccountState::Initialized,
|
|
is_native: COption::Some(100),
|
|
delegated_amount: 0,
|
|
close_authority: COption::None,
|
|
};
|
|
let account_size = ExtensionType::get_account_len::<TokenAccount>(&[
|
|
ExtensionType::ImmutableOwner,
|
|
ExtensionType::MemoTransfer,
|
|
]);
|
|
let mut account_data = vec![0; account_size];
|
|
let mut account_state =
|
|
StateWithExtensionsMut::<TokenAccount>::unpack_uninitialized(&mut account_data)
|
|
.unwrap();
|
|
account_state.base = token_base;
|
|
account_state.pack_base();
|
|
account_state.init_account_type().unwrap();
|
|
account_state
|
|
.init_extension::<ImmutableOwner>(true)
|
|
.unwrap();
|
|
let mut memo_transfer = account_state.init_extension::<MemoTransfer>(true).unwrap();
|
|
memo_transfer.require_incoming_transfer_memos = true.into();
|
|
|
|
let spl_token_account = Account {
|
|
lamports: 100,
|
|
data: account_data.to_vec(),
|
|
owner: spl_token_2022::id(),
|
|
executable: false,
|
|
rent_epoch: 0,
|
|
};
|
|
let other_account = Account {
|
|
lamports: 100,
|
|
data: account_data.to_vec(),
|
|
owner: Pubkey::new_unique(),
|
|
executable: false,
|
|
rent_epoch: 0,
|
|
};
|
|
|
|
let other_mint_token_base = TokenAccount {
|
|
mint: other_mint_pubkey,
|
|
owner: token_owner,
|
|
amount: 42,
|
|
delegate: COption::None,
|
|
state: spl_token_2022::state::AccountState::Initialized,
|
|
is_native: COption::Some(100),
|
|
delegated_amount: 0,
|
|
close_authority: COption::None,
|
|
};
|
|
let account_size = ExtensionType::get_account_len::<TokenAccount>(&[
|
|
ExtensionType::ImmutableOwner,
|
|
ExtensionType::MemoTransfer,
|
|
]);
|
|
let mut account_data = vec![0; account_size];
|
|
let mut account_state =
|
|
StateWithExtensionsMut::<TokenAccount>::unpack_uninitialized(&mut account_data)
|
|
.unwrap();
|
|
account_state.base = other_mint_token_base;
|
|
account_state.pack_base();
|
|
account_state.init_account_type().unwrap();
|
|
account_state
|
|
.init_extension::<ImmutableOwner>(true)
|
|
.unwrap();
|
|
let mut memo_transfer = account_state.init_extension::<MemoTransfer>(true).unwrap();
|
|
memo_transfer.require_incoming_transfer_memos = true.into();
|
|
|
|
let other_mint_token_account = Account {
|
|
lamports: 100,
|
|
data: account_data.to_vec(),
|
|
owner: spl_token_2022::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();
|
|
|
|
// Account is not owned by spl_token (nor does it have TokenAccount state)
|
|
assert_eq!(
|
|
collect_token_balance_from_account(&bank, &account_pubkey, &mut mint_decimals),
|
|
None
|
|
);
|
|
|
|
// Mint does not have TokenAccount state
|
|
assert_eq!(
|
|
collect_token_balance_from_account(&bank, &mint_pubkey, &mut mint_decimals),
|
|
None
|
|
);
|
|
|
|
// TokenAccount owned by spl_token_2022::id() works
|
|
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(),
|
|
},
|
|
program_id: spl_token_2022::id().to_string(),
|
|
})
|
|
);
|
|
|
|
// TokenAccount is not owned by known spl-token program_id
|
|
assert_eq!(
|
|
collect_token_balance_from_account(&bank, &other_account_pubkey, &mut mint_decimals),
|
|
None
|
|
);
|
|
|
|
// TokenAccount's mint is not owned by known spl-token program_id
|
|
assert_eq!(
|
|
collect_token_balance_from_account(
|
|
&bank,
|
|
&other_mint_account_pubkey,
|
|
&mut mint_decimals
|
|
),
|
|
None
|
|
);
|
|
}
|
|
}
|