token-2022: Add Reallocate instruction (#2864)

* Pass by ref

* Dedupe ExtensionTypes in length apis

* Depend on deduplication in ExtensionType::get_account_len

* Add Reallocate instruction

* Remove unneeded api

* Add set_account_type helper and remove unneeded StateWithExtensionsMut api

Author: Tyera Eulberg

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(

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<ExtensionType>,
+    extension_types: &[ExtensionType],
 ) -> Result<usize, ProgramError> {
     invoke(
         &spl_token::instruction::get_account_data_size(

token/program-2022/src/extension/mod.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<ExtensionType>,
+) -> 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::<Account>::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::<Account>(&current_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::<Account>(&mut token_account)?;
+
+    Ok(())
+}

token/program-2022/src/extension/reallocate.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<ExtensionType>,
+    },
 }
 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::<ExtensionType>()) {
+                    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<ExtensionType>,
+    extension_types: &[ExtensionType],
 ) -> Result<Instruction, ProgramError> {
     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<Instruction, ProgramError> {
+    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)

token/program-2022/src/instruction.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,
         },
@@ -992,18 +993,15 @@ impl Processor {
     /// Processes a [GetAccountDataSize](enum.TokenInstruction.html) instruction
     pub fn process_get_account_data_size(
         accounts: &[AccountInfo],
-        extension_types: Vec<ExtensionType>,
+        new_extension_types: Vec<ExtensionType>,
     ) -> 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>(&account_extensions);
         set_return_data(&account_len.to_le_bytes());
@@ -1151,6 +1149,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)
+            }
         }
     }
 
@@ -6919,7 +6921,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();
@@ -6933,7 +6935,7 @@ mod tests {
             get_account_data_size(
                 &program_id,
                 &mint_key,
-                vec![
+                &[
                     ExtensionType::TransferFeeAmount,
                     ExtensionType::TransferFeeAmount, // Duplicate user input ignored...
                 ],
@@ -6951,7 +6953,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();
@@ -6982,7 +6984,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();
@@ -6991,7 +6993,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],
@@ -7018,7 +7020,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())
@@ -7043,7 +7045,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)