Skip to content
This repository has been archived by the owner on Jan 22, 2025. It is now read-only.

Commit

Permalink
Support spl-token-2022 properly in pre-/postTokenBalances (backport #…
Browse files Browse the repository at this point in the history 
…26759) (#26775)

Support spl-token-2022 properly in pre-/postTokenBalances (#26759)

* Add failing test

* Use program-2022 unpacking to support token balances

* Add comments to make test results more clear

(cherry picked from commit a55e34d)

Co-authored-by: Tyera Eulberg <[email protected]>
  • Loading branch information
@mergify
mergify[bot] and Tyera Eulberg authored Jul 25, 2022
1 parent ae37e5e commit de3218f
Showing 1 changed file with 211 additions and 14 deletions.
225 changes: 211 additions & 14 deletions transaction-status/src/token_balances.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,8 +8,8 @@ use {
solana_metrics::datapoint_debug,
solana_runtime::{bank::Bank, transaction_batch::TransactionBatch},
solana_sdk::{account::ReadableAccount, pubkey::Pubkey},
spl_token::{
solana_program::program_pack::Pack,
spl_token_2022::{
extension::StateWithExtensions,
state::{Account as TokenAccount, Mint},
},
std::collections::HashMap,
Expand Down Expand Up @@ -45,8 +45,8 @@ fn get_mint_decimals(bank: &Bank, mint: &Pubkey) -> Option<u8> {
return None;
}

let decimals = Mint::unpack(mint_account.data())
.map(|mint| mint.decimals)
let decimals = StateWithExtensions::<Mint>::unpack(mint_account.data())
.map(|mint| mint.base.decimals)
.ok()?;

Some(decimals)
Expand Down Expand Up @@ -118,8 +118,8 @@ fn collect_token_balance_from_account(
return None;
}

let token_account = TokenAccount::unpack(account.data()).ok()?;
let mint = pubkey_from_spl_token(&token_account.mint);
let token_account = StateWithExtensions::<TokenAccount>::unpack(account.data()).ok()?;
let mint = pubkey_from_spl_token(&token_account.base.mint);

let decimals = mint_decimals.get(&mint).cloned().or_else(|| {
let decimals = get_mint_decimals(bank, &mint)?;
Expand All @@ -128,9 +128,9 @@ fn collect_token_balance_from_account(
})?;

Some(TokenBalanceData {
mint: token_account.mint.to_string(),
owner: token_account.owner.to_string(),
ui_token_amount: token_amount_to_ui_amount(token_account.amount, decimals),
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(),
})
}
Expand All @@ -141,7 +141,14 @@ mod test {
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,
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,
};

Expand Down Expand Up @@ -173,7 +180,7 @@ mod test {
let other_mint = Account {
lamports: 100,
data: data.to_vec(),
owner: Pubkey::new_unique(),
owner: Pubkey::new_unique(), // !is_known_spl_token_id
executable: false,
rent_epoch: 0,
};
Expand All @@ -184,7 +191,7 @@ mod test {
owner: spl_token_pubkey(&token_owner),
amount: 42,
delegate: COption::None,
state: spl_token::state::AccountState::Initialized,
state: spl_token_2022::state::AccountState::Initialized,
is_native: COption::Some(100),
delegated_amount: 0,
close_authority: COption::None,
Expand All @@ -202,7 +209,7 @@ mod test {
let other_account = Account {
lamports: 100,
data: data.to_vec(),
owner: Pubkey::new_unique(),
owner: Pubkey::new_unique(), // !is_known_spl_token_id
executable: false,
rent_epoch: 0,
};
Expand All @@ -212,7 +219,7 @@ mod test {
owner: spl_token_pubkey(&token_owner),
amount: 42,
delegate: COption::None,
state: spl_token::state::AccountState::Initialized,
state: spl_token_2022::state::AccountState::Initialized,
is_native: COption::Some(100),
delegated_amount: 0,
close_authority: COption::None,
Expand Down Expand Up @@ -246,16 +253,19 @@ mod test {
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,
Expand All @@ -275,11 +285,198 @@ mod test {
})
);

// 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>().unwrap();
mint_close_authority.close_authority =
OptionalNonZeroPubkey::try_from(Some(spl_token_pubkey(&mint_authority))).unwrap();

let mint_pubkey = Pubkey::new_unique();
let mint = Account {
lamports: 100,
data: mint_data.to_vec(),
owner: pubkey_from_spl_token(&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: spl_token_pubkey(&mint_pubkey),
owner: spl_token_pubkey(&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>().unwrap();
let mut memo_transfer = account_state.init_extension::<MemoTransfer>().unwrap();
memo_transfer.require_incoming_transfer_memos = true.into();

let spl_token_account = Account {
lamports: 100,
data: account_data.to_vec(),
owner: pubkey_from_spl_token(&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: spl_token_pubkey(&other_mint_pubkey),
owner: spl_token_pubkey(&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>().unwrap();
let mut memo_transfer = account_state.init_extension::<MemoTransfer>().unwrap();
memo_transfer.require_incoming_transfer_memos = true.into();

let other_mint_token_account = Account {
lamports: 100,
data: account_data.to_vec(),
owner: pubkey_from_spl_token(&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,
Expand Down

0 comments on commit de3218f

Please sign in to comment.