Skip to content
This repository was archived by the owner on Mar 11, 2025. It is now read-only.

Commit 809c1a1

Browse files
mvinesmvines
committed
GetAccountDataSize instruction now accepts user-provided extension types
1 parent 4ddd6c8 commit 809c1a1

File tree

3 files changed

+86
-17
lines changed

3 files changed

+86
-17
lines changed

‎associated-token-account/program/src/tools/account.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -78,7 +78,7 @@ pub fn get_account_len<'a>(
7878
spl_token_program: &AccountInfo<'a>,
7979
) -> Result<usize, ProgramError> {
8080
invoke(
81-
&spl_token::instruction::get_account_data_size(spl_token_program.key, mint.key)?,
81+
&spl_token::instruction::get_account_data_size(spl_token_program.key, mint.key, vec![])?,
8282
&[mint.clone(), spl_token_program.clone()],
8383
)?;
8484
get_return_data()

‎token/program-2022/src/instruction.rs

Lines changed: 37 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -2,8 +2,9 @@
22
33
use {
44
crate::{
5-
check_program_account, error::TokenError,
6-
extension::transfer_fee::instruction::TransferFeeInstruction,
5+
check_program_account,
6+
error::TokenError,
7+
extension::{transfer_fee::instruction::TransferFeeInstruction, ExtensionType},
78
},
89
solana_program::{
910
instruction::{AccountMeta, Instruction},
@@ -451,7 +452,10 @@ pub enum TokenInstruction {
451452
/// Accounts expected by this instruction:
452453
///
453454
/// 0. `[]` The mint to calculate for
454-
GetAccountDataSize,
455+
GetAccountDataSize {
456+
/// Additional extension types to include in the returned account size
457+
extension_types: Vec<ExtensionType>,
458+
},
455459
/// Initialize the close account authority on a new mint.
456460
///
457461
/// Fails if the mint has already been initialized, so must be called before
@@ -576,7 +580,13 @@ impl TokenInstruction {
576580
decimals,
577581
}
578582
}
579-
21 => Self::GetAccountDataSize,
583+
21 => {
584+
let mut extension_types = vec![];
585+
for chunk in rest.chunks(size_of::<ExtensionType>()) {
586+
extension_types.push(chunk.try_into()?);
587+
}
588+
Self::GetAccountDataSize { extension_types }
589+
}
580590
22 => {
581591
let (close_authority, _rest) = Self::unpack_pubkey_option(rest)?;
582592
Self::InitializeMintCloseAuthority { close_authority }
@@ -684,8 +694,13 @@ impl TokenInstruction {
684694
buf.extend_from_slice(mint_authority.as_ref());
685695
Self::pack_pubkey_option(freeze_authority, &mut buf);
686696
}
687-
&Self::GetAccountDataSize => {
697+
&Self::GetAccountDataSize {
698+
ref extension_types,
699+
} => {
688700
buf.push(21);
701+
for extension_type in extension_types {
702+
buf.extend_from_slice(&<[u8; 2]>::from(*extension_type));
703+
}
689704
}
690705
&Self::InitializeMintCloseAuthority {
691706
ref close_authority,
@@ -1416,12 +1431,13 @@ pub fn sync_native(
14161431
pub fn get_account_data_size(
14171432
token_program_id: &Pubkey,
14181433
mint_pubkey: &Pubkey,
1434+
extension_types: Vec<ExtensionType>,
14191435
) -> Result<Instruction, ProgramError> {
14201436
check_program_account(token_program_id)?;
14211437
Ok(Instruction {
14221438
program_id: *token_program_id,
14231439
accounts: vec![AccountMeta::new_readonly(*mint_pubkey, false)],
1424-
data: TokenInstruction::GetAccountDataSize.pack(),
1440+
data: TokenInstruction::GetAccountDataSize { extension_types }.pack(),
14251441
})
14261442
}
14271443

@@ -1662,13 +1678,27 @@ mod test {
16621678
let unpacked = TokenInstruction::unpack(&expect).unwrap();
16631679
assert_eq!(unpacked, check);
16641680

1665-
let check = TokenInstruction::GetAccountDataSize;
1681+
let check = TokenInstruction::GetAccountDataSize {
1682+
extension_types: vec![],
1683+
};
16661684
let packed = check.pack();
16671685
let expect = [21u8];
16681686
assert_eq!(packed, &[21u8]);
16691687
let unpacked = TokenInstruction::unpack(&expect).unwrap();
16701688
assert_eq!(unpacked, check);
16711689

1690+
let check = TokenInstruction::GetAccountDataSize {
1691+
extension_types: vec![
1692+
ExtensionType::TransferFeeConfig,
1693+
ExtensionType::TransferFeeAmount,
1694+
],
1695+
};
1696+
let packed = check.pack();
1697+
let expect = [21u8, 1, 0, 2, 0];
1698+
assert_eq!(packed, &[21u8, 1, 0, 2, 0]);
1699+
let unpacked = TokenInstruction::unpack(&expect).unwrap();
1700+
assert_eq!(unpacked, check);
1701+
16721702
let check = TokenInstruction::InitializeMintCloseAuthority {
16731703
close_authority: COption::Some(Pubkey::new(&[10u8; 32])),
16741704
};

‎token/program-2022/src/processor.rs

Lines changed: 48 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -962,11 +962,20 @@ impl Processor {
962962
}
963963

964964
/// Processes a [GetAccountDataSize](enum.TokenInstruction.html) instruction
965-
pub fn process_get_account_data_size(accounts: &[AccountInfo]) -> ProgramResult {
965+
pub fn process_get_account_data_size(
966+
accounts: &[AccountInfo],
967+
extension_types: Vec<ExtensionType>,
968+
) -> ProgramResult {
966969
let account_info_iter = &mut accounts.iter();
967970
let mint_account_info = next_account_info(account_info_iter)?;
968971

969-
let account_extensions = Self::get_required_account_extensions(mint_account_info)?;
972+
let mut account_extensions = Self::get_required_account_extensions(mint_account_info)?;
973+
974+
for extension_type in extension_types {
975+
if !account_extensions.contains(&extension_type) {
976+
account_extensions.push(extension_type);
977+
}
978+
}
970979

971980
let account_len = ExtensionType::get_account_len::<Account>(&account_extensions);
972981
set_return_data(&account_len.to_le_bytes());
@@ -1075,9 +1084,9 @@ impl Processor {
10751084
msg!("Instruction: SyncNative");
10761085
Self::process_sync_native(accounts)
10771086
}
1078-
TokenInstruction::GetAccountDataSize => {
1087+
TokenInstruction::GetAccountDataSize { extension_types } => {
10791088
msg!("Instruction: GetAccountDataSize");
1080-
Self::process_get_account_data_size(accounts)
1089+
Self::process_get_account_data_size(accounts, extension_types)
10811090
}
10821091
TokenInstruction::InitializeMintCloseAuthority { close_authority } => {
10831092
msg!("Instruction: InitializeMintCloseAuthority");
@@ -6725,7 +6734,26 @@ mod tests {
67256734
.to_vec(),
67266735
);
67276736
do_process_instruction(
6728-
get_account_data_size(&program_id, &mint_key).unwrap(),
6737+
get_account_data_size(&program_id, &mint_key, vec![]).unwrap(),
6738+
vec![&mut mint_account],
6739+
)
6740+
.unwrap();
6741+
6742+
set_expected_data(
6743+
ExtensionType::get_account_len::<Account>(&[ExtensionType::TransferFeeAmount])
6744+
.to_le_bytes()
6745+
.to_vec(),
6746+
);
6747+
do_process_instruction(
6748+
get_account_data_size(
6749+
&program_id,
6750+
&mint_key,
6751+
vec![
6752+
ExtensionType::TransferFeeAmount,
6753+
ExtensionType::TransferFeeAmount, // Duplicate user input ignored...
6754+
],
6755+
)
6756+
.unwrap(),
67296757
vec![&mut mint_account],
67306758
)
67316759
.unwrap();
@@ -6738,7 +6766,7 @@ mod tests {
67386766
.to_vec(),
67396767
);
67406768
do_process_instruction(
6741-
get_account_data_size(&program_id, &mint_key).unwrap(),
6769+
get_account_data_size(&program_id, &mint_key, vec![]).unwrap(),
67426770
vec![&mut mint_account],
67436771
)
67446772
.unwrap();
@@ -6769,7 +6797,18 @@ mod tests {
67696797
.to_vec(),
67706798
);
67716799
do_process_instruction(
6772-
get_account_data_size(&program_id, &mint_key).unwrap(),
6800+
get_account_data_size(&program_id, &mint_key, vec![]).unwrap(),
6801+
vec![&mut extended_mint_account],
6802+
)
6803+
.unwrap();
6804+
6805+
do_process_instruction(
6806+
get_account_data_size(
6807+
&program_id,
6808+
&mint_key,
6809+
vec![ExtensionType::TransferFeeAmount], // User extension that's also added by the mint ignored...
6810+
)
6811+
.unwrap(),
67736812
vec![&mut extended_mint_account],
67746813
)
67756814
.unwrap();
@@ -6794,7 +6833,7 @@ mod tests {
67946833

67956834
assert_eq!(
67966835
do_process_instruction(
6797-
get_account_data_size(&program_id, &invalid_mint_key).unwrap(),
6836+
get_account_data_size(&program_id, &invalid_mint_key, vec![]).unwrap(),
67986837
vec![&mut invalid_mint_account],
67996838
),
68006839
Err(TokenError::InvalidMint.into())
@@ -6819,7 +6858,7 @@ mod tests {
68196858

68206859
assert_eq!(
68216860
do_process_instruction(
6822-
get_account_data_size(&program_id, &invalid_mint_key).unwrap(),
6861+
get_account_data_size(&program_id, &invalid_mint_key, vec![]).unwrap(),
68236862
vec![&mut invalid_mint_account],
68246863
),
68256864
Err(ProgramError::IncorrectProgramId)

0 commit comments

Comments
 (0)