diff --git a/associated-token-account/program/src/processor.rs b/associated-token-account/program/src/processor.rs index 3f157ca1b2f..d9f32435326 100644 --- a/associated-token-account/program/src/processor.rs +++ b/associated-token-account/program/src/processor.rs @@ -79,7 +79,7 @@ pub fn process_create_associated_token_account( let account_len = get_account_len( spl_token_mint_info, spl_token_program_info, - vec![spl_token::extension::ExtensionType::ImmutableOwner], + &[spl_token::extension::ExtensionType::ImmutableOwner], )?; create_pda_account( diff --git a/associated-token-account/program/src/tools/account.rs b/associated-token-account/program/src/tools/account.rs index 4f8ae81b364..f1a7d1bde63 100644 --- a/associated-token-account/program/src/tools/account.rs +++ b/associated-token-account/program/src/tools/account.rs @@ -76,7 +76,7 @@ pub fn create_pda_account<'a>( pub fn get_account_len<'a>( mint: &AccountInfo<'a>, spl_token_program: &AccountInfo<'a>, - extension_types: Vec, + extension_types: &[ExtensionType], ) -> Result { invoke( &spl_token::instruction::get_account_data_size( diff --git a/token/program-2022/src/extension/mod.rs b/token/program-2022/src/extension/mod.rs index 65f31245b40..63895fc2fa8 100644 --- a/token/program-2022/src/extension/mod.rs +++ b/token/program-2022/src/extension/mod.rs @@ -33,6 +33,8 @@ pub mod default_account_state; pub mod immutable_owner; /// Mint Close Authority extension pub mod mint_close_authority; +/// Utility to reallocate token accounts +pub mod reallocate; /// Transfer Fee extension pub mod transfer_fee; @@ -326,28 +328,13 @@ impl<'data, S: BaseState> StateWithExtensionsMut<'data, S> { /// /// Fails if the base state is not initialized. pub fn unpack(input: &'data mut [u8]) -> Result { - Self::_unpack(input, false) - } - /// Unpack base state, leaving the extension data as a mutable slice - /// Checks the account_type, and initializes it if Uninitialized - /// - /// Fails if the base state is not initialized. - pub fn unpack_after_realloc(input: &'data mut [u8]) -> Result { - Self::_unpack(input, true) - } - - fn _unpack(input: &'data mut [u8], init_account_type: bool) -> Result { check_min_len_and_not_multisig(input, S::LEN)?; let (base_data, rest) = input.split_at_mut(S::LEN); let base = S::unpack(base_data)?; if let Some((account_type_index, tlv_start_index)) = type_and_tlv_indices::(rest)? { // type_and_tlv_indices() checks that returned indexes are within range - let mut account_type = AccountType::try_from(rest[account_type_index]) + let account_type = AccountType::try_from(rest[account_type_index]) .map_err(|_| ProgramError::InvalidAccountData)?; - if init_account_type && account_type == AccountType::Uninitialized { - rest[account_type_index] = S::ACCOUNT_TYPE.into(); - account_type = S::ACCOUNT_TYPE; - } check_account_type::(account_type)?; let (account_type, tlv_data) = rest.split_at_mut(tlv_start_index); Ok(Self { @@ -511,31 +498,24 @@ impl<'data, S: BaseState> StateWithExtensionsMut<'data, S> { fn get_first_extension_type(&self) -> Result, ProgramError> { get_first_extension_type(self.tlv_data) } +} - /// Compares the length of an extension with the currently used TLV buffer to determine if - /// reallocation is needed. If so, returns Some(v), where v is the difference between current - /// space and needed. - #[allow(dead_code)] - pub(crate) fn realloc_needed( - &self, - new_extension: ExtensionType, - ) -> Result, ProgramError> { - let mut extensions = self.get_extension_types()?; - if !extensions.contains(&new_extension) { - extensions.push(new_extension); - } - let new_needed_tlv_len = ExtensionType::get_total_tlv_len(&extensions); - if self.tlv_data.len() >= new_needed_tlv_len { - Ok(None) - } else { - let mut diff = new_needed_tlv_len - self.tlv_data.len(); // arithmetic safe because of if clause - if self.account_type.is_empty() { - diff = diff - .saturating_add(size_of::()) - .saturating_add(BASE_ACCOUNT_LENGTH.saturating_sub(S::LEN)); - } - Ok(Some(diff)) +/// If AccountType is uninitialized, set it to the BaseState's ACCOUNT_TYPE; +/// if AccountType is already set, check is set correctly for BaseState +pub fn set_account_type(input: &mut [u8]) -> Result<(), ProgramError> { + check_min_len_and_not_multisig(input, S::LEN)?; + let (_base_data, rest) = input.split_at_mut(S::LEN); + if let Some((account_type_index, _tlv_start_index)) = type_and_tlv_indices::(rest)? { + let mut account_type = AccountType::try_from(rest[account_type_index]) + .map_err(|_| ProgramError::InvalidAccountData)?; + if account_type == AccountType::Uninitialized { + rest[account_type_index] = S::ACCOUNT_TYPE.into(); + account_type = S::ACCOUNT_TYPE; } + check_account_type::(account_type)?; + Ok(()) + } else { + Err(ProgramError::InvalidAccountData) } } @@ -634,7 +614,14 @@ impl ExtensionType { /// Get the TLV length for a set of ExtensionTypes fn get_total_tlv_len(extension_types: &[Self]) -> usize { - let tlv_len: usize = extension_types.iter().map(|e| e.get_tlv_len()).sum(); + // dedupe extensions + let mut extensions = vec![]; + for extension_type in extension_types { + if !extensions.contains(&extension_type) { + extensions.push(extension_type); + } + } + let tlv_len: usize = extensions.iter().map(|e| e.get_tlv_len()).sum(); if tlv_len == Multisig::LEN .saturating_sub(BASE_ACCOUNT_LENGTH) @@ -1286,140 +1273,94 @@ mod test { } #[test] - fn test_unpack_after_realloc() { - // account + fn test_set_account_type() { + // account with buffer big enough for AccountType and Extension let mut buffer = TEST_ACCOUNT_SLICE.to_vec(); - let state = StateWithExtensionsMut::::unpack(&mut buffer).unwrap(); - let realloc = state - .realloc_needed(ExtensionType::ImmutableOwner) - .unwrap() - .unwrap(); - drop(state); - buffer.append(&mut vec![0; realloc]); - let mut state = - StateWithExtensionsMut::::unpack_after_realloc(&mut buffer).unwrap(); + let needed_len = + ExtensionType::get_account_len::(&[ExtensionType::ImmutableOwner]) + - buffer.len(); + buffer.append(&mut vec![0; needed_len]); + let err = StateWithExtensionsMut::::unpack(&mut buffer).unwrap_err(); + assert_eq!(err, ProgramError::InvalidAccountData); + set_account_type::(&mut buffer).unwrap(); + // unpack is viable after manual set_account_type + let mut state = StateWithExtensionsMut::::unpack(&mut buffer).unwrap(); assert_eq!(state.base, TEST_ACCOUNT); assert_eq!(state.account_type[0], AccountType::Account as u8); - state.init_extension::().unwrap(); + state.init_extension::().unwrap(); // just confirming initialization works - // account with AccountType + // account with buffer big enough for AccountType only let mut buffer = TEST_ACCOUNT_SLICE.to_vec(); - buffer.append(&mut vec![2, 0]); + buffer.append(&mut vec![0; 2]); + let err = StateWithExtensionsMut::::unpack(&mut buffer).unwrap_err(); + assert_eq!(err, ProgramError::InvalidAccountData); + set_account_type::(&mut buffer).unwrap(); + // unpack is viable after manual set_account_type let state = StateWithExtensionsMut::::unpack(&mut buffer).unwrap(); assert_eq!(state.base, TEST_ACCOUNT); assert_eq!(state.account_type[0], AccountType::Account as u8); - let realloc = state - .realloc_needed(ExtensionType::ImmutableOwner) - .unwrap() - .unwrap(); - drop(state); - buffer.append(&mut vec![0; realloc]); - let mut state = - StateWithExtensionsMut::::unpack_after_realloc(&mut buffer).unwrap(); + + // account with AccountType already set => noop + let mut buffer = TEST_ACCOUNT_SLICE.to_vec(); + buffer.append(&mut vec![2, 0]); + let _ = StateWithExtensionsMut::::unpack(&mut buffer).unwrap(); + set_account_type::(&mut buffer).unwrap(); + let state = StateWithExtensionsMut::::unpack(&mut buffer).unwrap(); assert_eq!(state.base, TEST_ACCOUNT); assert_eq!(state.account_type[0], AccountType::Account as u8); - state.init_extension::().unwrap(); - // account with wrong AccountType + // account with wrong AccountType fails let mut buffer = TEST_ACCOUNT_SLICE.to_vec(); buffer.append(&mut vec![1, 0]); - let err = StateWithExtensionsMut::::unpack_after_realloc(&mut buffer).unwrap_err(); + let err = StateWithExtensionsMut::::unpack(&mut buffer).unwrap_err(); + assert_eq!(err, ProgramError::InvalidAccountData); + let err = set_account_type::(&mut buffer).unwrap_err(); assert_eq!(err, ProgramError::InvalidAccountData); - // account with pre-existing extension - let account_size = - ExtensionType::get_account_len::(&[ExtensionType::ImmutableOwner]); - let mut buffer = vec![0; account_size]; - let mut state = - StateWithExtensionsMut::::unpack_uninitialized(&mut buffer).unwrap(); - state.base = TEST_ACCOUNT; - state.pack_base(); - state.init_extension::().unwrap(); - state.init_account_type().unwrap(); - let realloc = state - .realloc_needed(ExtensionType::TransferFeeAmount) - .unwrap() - .unwrap(); - drop(state); - buffer.append(&mut vec![0; realloc]); - let mut state = - StateWithExtensionsMut::::unpack_after_realloc(&mut buffer).unwrap(); - assert_eq!(state.base, TEST_ACCOUNT); - assert_eq!(state.account_type[0], AccountType::Account as u8); - state.init_extension::().unwrap(); - assert_eq!( - state.get_extension_types().unwrap(), - vec![ - ExtensionType::ImmutableOwner, - ExtensionType::TransferFeeAmount - ] - ); - - // mint + // mint with buffer big enough for AccountType and Extension let mut buffer = TEST_MINT_SLICE.to_vec(); - let state = StateWithExtensionsMut::::unpack(&mut buffer).unwrap(); - let realloc = state - .realloc_needed(ExtensionType::MintCloseAuthority) - .unwrap() - .unwrap(); - drop(state); - buffer.append(&mut vec![0; realloc]); - let mut state = StateWithExtensionsMut::::unpack_after_realloc(&mut buffer).unwrap(); + let needed_len = + ExtensionType::get_account_len::(&[ExtensionType::MintCloseAuthority]) + - buffer.len(); + buffer.append(&mut vec![0; needed_len]); + let err = StateWithExtensionsMut::::unpack(&mut buffer).unwrap_err(); + assert_eq!(err, ProgramError::InvalidAccountData); + set_account_type::(&mut buffer).unwrap(); + // unpack is viable after manual set_account_type + let mut state = StateWithExtensionsMut::::unpack(&mut buffer).unwrap(); assert_eq!(state.base, TEST_MINT); assert_eq!(state.account_type[0], AccountType::Mint as u8); state.init_extension::().unwrap(); - // mint with AccountType + // mint with buffer big enough for AccountType only let mut buffer = TEST_MINT_SLICE.to_vec(); buffer.append(&mut vec![0; Account::LEN - Mint::LEN]); - buffer.append(&mut vec![1, 0]); + buffer.append(&mut vec![0; 2]); + let err = StateWithExtensionsMut::::unpack(&mut buffer).unwrap_err(); + assert_eq!(err, ProgramError::InvalidAccountData); + set_account_type::(&mut buffer).unwrap(); + // unpack is viable after manual set_account_type let state = StateWithExtensionsMut::::unpack(&mut buffer).unwrap(); assert_eq!(state.base, TEST_MINT); assert_eq!(state.account_type[0], AccountType::Mint as u8); - let realloc = state - .realloc_needed(ExtensionType::MintCloseAuthority) - .unwrap() - .unwrap(); - drop(state); - buffer.append(&mut vec![0; realloc]); - let mut state = StateWithExtensionsMut::::unpack_after_realloc(&mut buffer).unwrap(); + + // mint with AccountType already set => noop + let mut buffer = TEST_MINT_SLICE.to_vec(); + buffer.append(&mut vec![0; Account::LEN - Mint::LEN]); + buffer.append(&mut vec![1, 0]); + set_account_type::(&mut buffer).unwrap(); + let state = StateWithExtensionsMut::::unpack(&mut buffer).unwrap(); assert_eq!(state.base, TEST_MINT); assert_eq!(state.account_type[0], AccountType::Mint as u8); - state.init_extension::().unwrap(); - // mint with wrong AccountType + // mint with wrong AccountType fails let mut buffer = TEST_MINT_SLICE.to_vec(); buffer.append(&mut vec![0; Account::LEN - Mint::LEN]); buffer.append(&mut vec![2, 0]); - let err = StateWithExtensionsMut::::unpack_after_realloc(&mut buffer).unwrap_err(); + let err = StateWithExtensionsMut::::unpack(&mut buffer).unwrap_err(); + assert_eq!(err, ProgramError::InvalidAccountData); + let err = set_account_type::(&mut buffer).unwrap_err(); assert_eq!(err, ProgramError::InvalidAccountData); - - // mint with pre-existing extension - let mut buffer = MINT_WITH_EXTENSION.to_vec(); - let state = StateWithExtensionsMut::::unpack(&mut buffer).unwrap(); - assert_eq!(state.base, TEST_MINT); - assert_eq!(state.account_type[0], AccountType::Mint as u8); - assert_eq!( - state.get_extension_types().unwrap(), - vec![ExtensionType::MintCloseAuthority] - ); - let realloc = state - .realloc_needed(ExtensionType::TransferFeeConfig) - .unwrap() - .unwrap(); - drop(state); - buffer.append(&mut vec![0; realloc]); - let mut state = StateWithExtensionsMut::::unpack_after_realloc(&mut buffer).unwrap(); - assert_eq!(state.base, TEST_MINT); - assert_eq!(state.account_type[0], AccountType::Mint as u8); - state.init_extension::().unwrap(); - assert_eq!( - state.get_extension_types().unwrap(), - vec![ - ExtensionType::MintCloseAuthority, - ExtensionType::TransferFeeConfig - ] - ); } #[test] @@ -1535,155 +1476,6 @@ mod test { assert_eq!(state.get_extension_types().unwrap(), vec![]); } - #[test] - fn test_realloc_needed() { - // buffer exact size of base-state account - let account_size = ExtensionType::get_account_len::(&[]); - let mut buffer = vec![0; account_size]; - let mut state = - StateWithExtensionsMut::::unpack_uninitialized(&mut buffer).unwrap(); - state.base = TEST_ACCOUNT; - state.pack_base(); - state.init_account_type().unwrap(); - let realloc = state.realloc_needed(ExtensionType::ImmutableOwner).unwrap(); - assert_eq!( - realloc, - Some(ExtensionType::ImmutableOwner.get_tlv_len() + size_of::()) - ); - assert_eq!( - account_size + realloc.unwrap(), - ExtensionType::get_account_len::(&[ExtensionType::ImmutableOwner]) - ); - let mut buffer = vec![0; account_size + realloc.unwrap()]; - let mut state = - StateWithExtensionsMut::::unpack_uninitialized(&mut buffer).unwrap(); - state.base = TEST_ACCOUNT; - state.pack_base(); - state.init_account_type().unwrap(); - state.init_extension::().unwrap(); - - // buffer exact size of base-state mint - let account_size = ExtensionType::get_account_len::(&[]); - let mut buffer = vec![0; account_size]; - let mut state = StateWithExtensionsMut::::unpack_uninitialized(&mut buffer).unwrap(); - state.base = TEST_MINT; - state.pack_base(); - state.init_account_type().unwrap(); - let realloc = state - .realloc_needed(ExtensionType::MintCloseAuthority) - .unwrap(); - assert_eq!( - realloc, - Some( - ExtensionType::MintCloseAuthority.get_tlv_len() - + size_of::() - + (Account::LEN - Mint::LEN) - ) - ); - assert_eq!( - account_size + realloc.unwrap(), - ExtensionType::get_account_len::(&[ExtensionType::MintCloseAuthority]) - ); - let mut buffer = vec![0; account_size + realloc.unwrap()]; - let mut state = StateWithExtensionsMut::::unpack_uninitialized(&mut buffer).unwrap(); - state.base = TEST_MINT; - state.pack_base(); - state.init_account_type().unwrap(); - state.init_extension::().unwrap(); - - // buffer exact size of existing extension - let mint_size = ExtensionType::get_account_len::(&[ExtensionType::TransferFeeConfig]); - let mut buffer = vec![0; mint_size]; - let mut state = StateWithExtensionsMut::::unpack_uninitialized(&mut buffer).unwrap(); - state.base = TEST_MINT; - state.pack_base(); - state.init_account_type().unwrap(); - assert_eq!( - state - .realloc_needed(ExtensionType::TransferFeeConfig) - .unwrap(), - None - ); - state.init_extension::().unwrap(); - let realloc = state - .realloc_needed(ExtensionType::MintCloseAuthority) - .unwrap(); - assert_eq!( - realloc, - Some(ExtensionType::MintCloseAuthority.get_tlv_len()) - ); - assert_eq!( - mint_size + realloc.unwrap(), - ExtensionType::get_account_len::(&[ - ExtensionType::TransferFeeConfig, - ExtensionType::MintCloseAuthority - ]) - ); - let mut buffer = vec![0; mint_size + realloc.unwrap()]; - let mut state = StateWithExtensionsMut::::unpack_uninitialized(&mut buffer).unwrap(); - state.base = TEST_MINT; - state.pack_base(); - state.init_account_type().unwrap(); - state.init_extension::().unwrap(); - state.init_extension::().unwrap(); - - // buffer with multisig len - let mint_size = ExtensionType::get_account_len::(&[ExtensionType::MintPaddingTest]); - let mut buffer = vec![0; mint_size]; - let mut state = StateWithExtensionsMut::::unpack_uninitialized(&mut buffer).unwrap(); - state.base = TEST_MINT; - state.pack_base(); - state.init_account_type().unwrap(); - assert_eq!( - state - .realloc_needed(ExtensionType::MintPaddingTest) - .unwrap(), - None - ); - state.init_extension::().unwrap(); - let realloc = state - .realloc_needed(ExtensionType::MintCloseAuthority) - .unwrap(); - assert_eq!( - realloc, - Some(ExtensionType::MintCloseAuthority.get_tlv_len() - size_of::()) - ); - assert_eq!( - mint_size + realloc.unwrap(), - ExtensionType::get_account_len::(&[ - ExtensionType::MintPaddingTest, - ExtensionType::MintCloseAuthority - ]) - ); - let mut buffer = vec![0; mint_size + realloc.unwrap()]; - let mut state = StateWithExtensionsMut::::unpack_uninitialized(&mut buffer).unwrap(); - state.base = TEST_MINT; - state.pack_base(); - state.init_account_type().unwrap(); - state.init_extension::().unwrap(); - state.init_extension::().unwrap(); - - // huge buffer - let mut buffer = vec![0; u16::MAX.into()]; - let mut state = StateWithExtensionsMut::::unpack_uninitialized(&mut buffer).unwrap(); - state.base = TEST_MINT; - state.pack_base(); - state.init_account_type().unwrap(); - assert_eq!( - state - .realloc_needed(ExtensionType::TransferFeeConfig) - .unwrap(), - None - ); - state.init_extension::().unwrap(); - assert_eq!( - state - .realloc_needed(ExtensionType::MintCloseAuthority) - .unwrap(), - None - ); - } - #[test] fn test_extension_with_no_data() { let account_size = diff --git a/token/program-2022/src/extension/reallocate.rs b/token/program-2022/src/extension/reallocate.rs new file mode 100644 index 00000000000..7c2bf555ad0 --- /dev/null +++ b/token/program-2022/src/extension/reallocate.rs @@ -0,0 +1,87 @@ +use { + crate::{ + error::TokenError, + extension::{set_account_type, AccountType, ExtensionType, StateWithExtensions}, + processor::Processor, + state::Account, + }, + solana_program::{ + account_info::{next_account_info, AccountInfo}, + entrypoint::ProgramResult, + msg, + program::invoke, + pubkey::Pubkey, + system_instruction, + sysvar::{rent::Rent, Sysvar}, + }, +}; + +/// Processes a [Reallocate](enum.TokenInstruction.html) instruction +pub fn process_reallocate( + program_id: &Pubkey, + accounts: &[AccountInfo], + new_extension_types: Vec, +) -> ProgramResult { + let account_info_iter = &mut accounts.iter(); + let token_account_info = next_account_info(account_info_iter)?; + let payer_info = next_account_info(account_info_iter)?; + let system_program_info = next_account_info(account_info_iter)?; + let authority_info = next_account_info(account_info_iter)?; + let authority_info_data_len = authority_info.data_len(); + + // check that account is the right type and validate owner + let mut current_extension_types = { + let token_account = token_account_info.data.borrow(); + let account = StateWithExtensions::::unpack(&token_account)?; + Processor::validate_owner( + program_id, + &account.base.owner, + authority_info, + authority_info_data_len, + account_info_iter.as_slice(), + )?; + account.get_extension_types()? + }; + + // check that all desired extensions are for the right account type + if new_extension_types + .iter() + .any(|extension_type| extension_type.get_account_type() != AccountType::Account) + { + return Err(TokenError::InvalidState.into()); + } + // ExtensionType::get_account_len() dedupes types, so just a dumb concatenation is fine here + current_extension_types.extend_from_slice(&new_extension_types); + let needed_account_len = ExtensionType::get_account_len::(¤t_extension_types); + + // if account is already large enough, return early + if token_account_info.data_len() >= needed_account_len { + return Ok(()); + } + + // reallocate + msg!( + "account needs realloc, +{:?} bytes", + needed_account_len - token_account_info.data_len() + ); + token_account_info.realloc(needed_account_len, false)?; + + // if additional lamports needed to remain rent-exempt, transfer them + let rent = Rent::get()?; + let new_minimum_balance = rent.minimum_balance(needed_account_len); + let lamports_diff = new_minimum_balance.saturating_sub(token_account_info.lamports()); + invoke( + &system_instruction::transfer(payer_info.key, token_account_info.key, lamports_diff), + &[ + payer_info.clone(), + token_account_info.clone(), + system_program_info.clone(), + ], + )?; + + // unpack to set account_type, if needed + let mut token_account = token_account_info.data.borrow_mut(); + set_account_type::(&mut token_account)?; + + Ok(()) +} diff --git a/token/program-2022/src/instruction.rs b/token/program-2022/src/instruction.rs index a6ad8476272..37d2bf9b383 100644 --- a/token/program-2022/src/instruction.rs +++ b/token/program-2022/src/instruction.rs @@ -11,7 +11,7 @@ use { program_error::ProgramError, program_option::COption, pubkey::{Pubkey, PUBKEY_BYTES}, - sysvar, + system_program, sysvar, }, std::{convert::TryInto, mem::size_of}, }; @@ -495,11 +495,33 @@ pub enum TokenInstruction { /// Accounts expected by this instruction: /// /// 0. `[writable]` The account to initialize. - // + /// /// Data expected by this instruction: /// None /// InitializeImmutableOwner, + /// Check to see if a token account is large enough for a list of ExtensionTypes, and if not, + /// use reallocation to increase the data size. + /// + /// Accounts expected by this instruction: + /// + /// * Single owner + /// 0. `[writable]` The account to reallocate. + /// 1. `[signer, writable]` The payer account to fund reallocation + /// 2. `[]` System program for reallocation funding + /// 3. `[signer]` The account's owner. + /// + /// * Multisignature owner + /// 0. `[writable]` The account to reallocate. + /// 1. `[signer, writable]` The payer account to fund reallocation + /// 2. `[]` System program for reallocation funding + /// 3. `[]` The account's multisignature owner/delegate. + /// 4. ..4+M `[signer]` M signer accounts. + /// + Reallocate { + /// New extension types to include in the reallocated account + extension_types: Vec, + }, } impl TokenInstruction { /// Unpacks a byte buffer into a [TokenInstruction](enum.TokenInstruction.html). @@ -611,6 +633,13 @@ impl TokenInstruction { 24 => Self::ConfidentialTransferExtension, 25 => Self::DefaultAccountStateExtension, 26 => Self::InitializeImmutableOwner, + 27 => { + let mut extension_types = vec![]; + for chunk in rest.chunks(size_of::()) { + extension_types.push(chunk.try_into()?); + } + Self::Reallocate { extension_types } + } _ => return Err(TokenError::InvalidInstruction.into()), }) } @@ -735,6 +764,14 @@ impl TokenInstruction { &Self::InitializeImmutableOwner => { buf.push(26); } + &Self::Reallocate { + ref extension_types, + } => { + buf.push(27); + for extension_type in extension_types { + buf.extend_from_slice(&<[u8; 2]>::from(*extension_type)); + } + } }; buf } @@ -1448,13 +1485,16 @@ pub fn sync_native( pub fn get_account_data_size( token_program_id: &Pubkey, mint_pubkey: &Pubkey, - extension_types: Vec, + extension_types: &[ExtensionType], ) -> Result { check_program_account(token_program_id)?; Ok(Instruction { program_id: *token_program_id, accounts: vec![AccountMeta::new_readonly(*mint_pubkey, false)], - data: TokenInstruction::GetAccountDataSize { extension_types }.pack(), + data: TokenInstruction::GetAccountDataSize { + extension_types: extension_types.to_vec(), + } + .pack(), }) } @@ -1486,6 +1526,39 @@ pub fn initialize_immutable_owner( }) } +/// Creates a `Reallocate` instruction +pub fn reallocate( + token_program_id: &Pubkey, + account_pubkey: &Pubkey, + payer: &Pubkey, + owner_pubkey: &Pubkey, + signer_pubkeys: &[&Pubkey], + extension_types: &[ExtensionType], +) -> Result { + check_program_account(token_program_id)?; + + let mut accounts = Vec::with_capacity(4 + signer_pubkeys.len()); + accounts.push(AccountMeta::new(*account_pubkey, false)); + accounts.push(AccountMeta::new(*payer, true)); + accounts.push(AccountMeta::new_readonly(system_program::id(), false)); + accounts.push(AccountMeta::new_readonly( + *owner_pubkey, + signer_pubkeys.is_empty(), + )); + for signer_pubkey in signer_pubkeys.iter() { + accounts.push(AccountMeta::new_readonly(**signer_pubkey, true)); + } + + Ok(Instruction { + program_id: *token_program_id, + accounts, + data: TokenInstruction::Reallocate { + extension_types: extension_types.to_vec(), + } + .pack(), + }) +} + /// Utility function that checks index is between MIN_SIGNERS and MAX_SIGNERS pub fn is_valid_signer_index(index: usize) -> bool { (MIN_SIGNERS..=MAX_SIGNERS).contains(&index) diff --git a/token/program-2022/src/processor.rs b/token/program-2022/src/processor.rs index ca1435916d4..3d3a9320872 100644 --- a/token/program-2022/src/processor.rs +++ b/token/program-2022/src/processor.rs @@ -9,6 +9,7 @@ use { default_account_state::{self, DefaultAccountState}, immutable_owner::ImmutableOwner, mint_close_authority::MintCloseAuthority, + reallocate, transfer_fee::{self, TransferFeeAmount, TransferFeeConfig}, ExtensionType, StateWithExtensions, StateWithExtensionsMut, }, @@ -969,18 +970,15 @@ impl Processor { /// Processes a [GetAccountDataSize](enum.TokenInstruction.html) instruction pub fn process_get_account_data_size( accounts: &[AccountInfo], - extension_types: Vec, + new_extension_types: Vec, ) -> ProgramResult { let account_info_iter = &mut accounts.iter(); let mint_account_info = next_account_info(account_info_iter)?; let mut account_extensions = Self::get_required_account_extensions(mint_account_info)?; - - for extension_type in extension_types { - if !account_extensions.contains(&extension_type) { - account_extensions.push(extension_type); - } - } + // ExtensionType::get_account_len() dedupes types, so just a dumb concatenation is fine + // here + account_extensions.extend_from_slice(&new_extension_types); let account_len = ExtensionType::get_account_len::(&account_extensions); set_return_data(&account_len.to_le_bytes()); @@ -1128,6 +1126,10 @@ impl Processor { msg!("Instruction: InitializeImmutableOwner"); Self::process_initialize_immutable_owner(accounts) } + TokenInstruction::Reallocate { extension_types } => { + msg!("Instruction: Reallocate"); + reallocate::process_reallocate(program_id, accounts, extension_types) + } } } @@ -6825,7 +6827,7 @@ mod tests { .to_vec(), ); do_process_instruction( - get_account_data_size(&program_id, &mint_key, vec![]).unwrap(), + get_account_data_size(&program_id, &mint_key, &[]).unwrap(), vec![&mut mint_account], ) .unwrap(); @@ -6839,7 +6841,7 @@ mod tests { get_account_data_size( &program_id, &mint_key, - vec![ + &[ ExtensionType::TransferFeeAmount, ExtensionType::TransferFeeAmount, // Duplicate user input ignored... ], @@ -6857,7 +6859,7 @@ mod tests { .to_vec(), ); do_process_instruction( - get_account_data_size(&program_id, &mint_key, vec![]).unwrap(), + get_account_data_size(&program_id, &mint_key, &[]).unwrap(), vec![&mut mint_account], ) .unwrap(); @@ -6888,7 +6890,7 @@ mod tests { .to_vec(), ); do_process_instruction( - get_account_data_size(&program_id, &mint_key, vec![]).unwrap(), + get_account_data_size(&program_id, &mint_key, &[]).unwrap(), vec![&mut extended_mint_account], ) .unwrap(); @@ -6897,7 +6899,7 @@ mod tests { get_account_data_size( &program_id, &mint_key, - vec![ExtensionType::TransferFeeAmount], // User extension that's also added by the mint ignored... + &[ExtensionType::TransferFeeAmount], // User extension that's also added by the mint ignored... ) .unwrap(), vec![&mut extended_mint_account], @@ -6924,7 +6926,7 @@ mod tests { assert_eq!( do_process_instruction( - get_account_data_size(&program_id, &invalid_mint_key, vec![]).unwrap(), + get_account_data_size(&program_id, &invalid_mint_key, &[]).unwrap(), vec![&mut invalid_mint_account], ), Err(TokenError::InvalidMint.into()) @@ -6949,7 +6951,7 @@ mod tests { assert_eq!( do_process_instruction( - get_account_data_size(&program_id, &invalid_mint_key, vec![]).unwrap(), + get_account_data_size(&program_id, &invalid_mint_key, &[]).unwrap(), vec![&mut invalid_mint_account], ), Err(ProgramError::IncorrectProgramId) diff --git a/token/program-2022/tests/reallocate.rs b/token/program-2022/tests/reallocate.rs new file mode 100644 index 00000000000..f53c36c43de --- /dev/null +++ b/token/program-2022/tests/reallocate.rs @@ -0,0 +1,164 @@ +#![cfg(feature = "test-bpf")] + +mod program_test; +use { + program_test::{TestContext, TokenContext}, + solana_program_test::tokio, + solana_sdk::{ + instruction::InstructionError, program_option::COption, pubkey::Pubkey, signature::Signer, + signer::keypair::Keypair, transaction::TransactionError, transport::TransportError, + }, + spl_token_2022::{error::TokenError, extension::ExtensionType, state::Account}, + spl_token_client::token::{ExtensionInitializationParams, TokenError as TokenClientError}, + std::convert::TryInto, +}; + +#[tokio::test] +async fn reallocate() { + let mut context = TestContext::new().await; + context.init_token_with_mint(vec![]).await.unwrap(); + let TokenContext { + token, + alice, + mint_authority, + .. + } = context.token_context.unwrap(); + + // reallocate fails on wrong account type + let error = token + .reallocate( + token.get_address(), + &mint_authority, + &[ExtensionType::ImmutableOwner], + ) + .await + .unwrap_err(); + assert_eq!( + error, + TokenClientError::Client(Box::new(TransportError::TransactionError( + TransactionError::InstructionError(0, InstructionError::InvalidAccountData) + ))) + ); + + // create account just large enough for base + let alice_account = Keypair::new(); + let alice_account = token + .create_auxiliary_token_account(&alice_account, &alice.pubkey()) + .await + .unwrap(); + + // reallocate fails on invalid extension type + let error = token + .reallocate(&alice_account, &alice, &[ExtensionType::MintCloseAuthority]) + .await + .unwrap_err(); + assert_eq!( + error, + TokenClientError::Client(Box::new(TransportError::TransactionError( + TransactionError::InstructionError( + 0, + InstructionError::Custom(TokenError::InvalidState as u32) + ) + ))) + ); + + // reallocate fails on invalid authority + let error = token + .reallocate( + &alice_account, + &mint_authority, + &[ExtensionType::ImmutableOwner], + ) + .await + .unwrap_err(); + assert_eq!( + error, + TokenClientError::Client(Box::new(TransportError::TransactionError( + TransactionError::InstructionError( + 0, + InstructionError::Custom(TokenError::OwnerMismatch as u32) + ) + ))) + ); + + // reallocate succeeds + token + .reallocate(&alice_account, &alice, &[ExtensionType::ImmutableOwner]) + .await + .unwrap(); + let account = token.get_account(&alice_account).await.unwrap(); + assert_eq!( + account.data.len(), + ExtensionType::get_account_len::(&[ExtensionType::ImmutableOwner]) + ); + + // reallocate succeeds with noop if account is already large enough + token + .reallocate(&alice_account, &alice, &[ExtensionType::ImmutableOwner]) + .await + .unwrap(); + let account = token.get_account(&alice_account).await.unwrap(); + assert_eq!( + account.data.len(), + ExtensionType::get_account_len::(&[ExtensionType::ImmutableOwner]) + ); + + // reallocate only reallocates enough for new extension, and dedupes extensions + token + .reallocate( + &alice_account, + &alice, + &[ + ExtensionType::ImmutableOwner, + ExtensionType::ImmutableOwner, + ExtensionType::TransferFeeAmount, + ExtensionType::TransferFeeAmount, + ], + ) + .await + .unwrap(); + let account = token.get_account(&alice_account).await.unwrap(); + assert_eq!( + account.data.len(), + ExtensionType::get_account_len::(&[ + ExtensionType::ImmutableOwner, + ExtensionType::TransferFeeAmount + ]) + ); +} + +#[tokio::test] +async fn reallocate_without_current_extension_knowledge() { + let mut context = TestContext::new().await; + context + .init_token_with_mint(vec![ExtensionInitializationParams::TransferFeeConfig { + transfer_fee_config_authority: COption::Some(Pubkey::new_unique()).try_into().unwrap(), + withdraw_withheld_authority: COption::Some(Pubkey::new_unique()).try_into().unwrap(), + transfer_fee_basis_points: 250, + maximum_fee: 10_000_000, + }]) + .await + .unwrap(); + let TokenContext { token, alice, .. } = context.token_context.unwrap(); + + // create account just large enough for TransferFeeAmount extension + let alice_account = Keypair::new(); + let alice_account = token + .create_auxiliary_token_account(&alice_account, &alice.pubkey()) + .await + .unwrap(); + + // reallocate resizes account to accommodate new and existing extensions + token + .reallocate(&alice_account, &alice, &[ExtensionType::ImmutableOwner]) + .await + .unwrap(); + let account = token.get_account(&alice_account).await.unwrap(); + assert_eq!( + account.data.len(), + ExtensionType::get_account_len::(&[ + ExtensionType::TransferFeeAmount, + ExtensionType::ImmutableOwner + ]) + ); +} diff --git a/token/rust/src/token.rs b/token/rust/src/token.rs index bfbfa9d15cf..4ea0d1bf492 100644 --- a/token/rust/src/token.rs +++ b/token/rust/src/token.rs @@ -620,4 +620,25 @@ where ) .await } + + /// Reallocate a token account to be large enough for a set of ExtensionTypes + pub async fn reallocate( + &self, + account: &Pubkey, + authority: &S2, + extension_types: &[ExtensionType], + ) -> TokenResult { + self.process_ixs( + &[instruction::reallocate( + &self.program_id, + account, + &self.payer.pubkey(), + &authority.pubkey(), + &[], + extension_types, + )?], + &[authority], + ) + .await + } }