add an option to realloc account when configuring account with registry

Author: samkim-crypto

token/client/src/token.rs

@@ -1980,6 +1980,7 @@ where
         &self,
         account: &Pubkey,
         elgamal_registry_account: &Pubkey,
+        payer: Option<&Pubkey>,
     ) -> TokenResult<T::Output> {
         self.process_ixs::<[&dyn Signer; 0]>(
             &[
@@ -1988,6 +1989,7 @@ where
                     account,
                     &self.pubkey,
                     elgamal_registry_account,
+                    payer,
                 )?,
             ],
             &[],

token/program-2022-test/tests/confidential_transfer.rs

@@ -2892,6 +2892,7 @@ async fn confidential_transfer_configure_token_account_with_registry() {
         ProofData::InstructionData(&proof_data),
     );
 
+    let payer_pubkey = ctx.payer.pubkey();
     let instructions =
         spl_elgamal_registry::instruction::update_registry(&alice.pubkey(), proof_location)
             .unwrap();
@@ -2911,7 +2912,7 @@ async fn confidential_transfer_configure_token_account_with_registry() {
         .create_auxiliary_token_account_with_extension_space(
             &alice_account_keypair,
             &alice.pubkey(),
-            vec![ExtensionType::ConfidentialTransferAccount],
+            vec![], // do not allocate space for confidential transfers
         )
         .await
         .unwrap();
@@ -2920,6 +2921,7 @@ async fn confidential_transfer_configure_token_account_with_registry() {
         .confidential_transfer_configure_token_account_with_registry(
             &alice_account_keypair.pubkey(),
             &elgamal_registry_address,
+            Some(&payer_pubkey), // test account allocation
         )
         .await
         .unwrap();

token/program-2022/src/extension/confidential_transfer/instruction.rs

@@ -18,7 +18,7 @@ use {
         instruction::{AccountMeta, Instruction},
         program_error::ProgramError,
         pubkey::Pubkey,
-        sysvar,
+        system_program, sysvar,
     },
     spl_token_confidential_transfer_proof_extraction::instruction::{ProofData, ProofLocation},
 };
@@ -484,15 +484,21 @@ pub enum ConfidentialTransferInstruction {
     /// then the program skips the verification of the ElGamal pubkey
     /// validity proof as well as the token owner signature.
     ///
+    /// If the token account is not large enough to include the new cconfidential transfer
+    /// extension, then optionally reallocate the account to increase the data size.
+    ///
     /// Accounts expected by this instruction:
     ///
     ///   * Single owner/delegate
     ///   0. `[writable]` The SPL Token account.
     ///   1. `[]` The corresponding SPL Token mint.
     ///   2. `[]` The ElGamal registry account.
+    ///   3. `[signer, writable]` (Optional) The payer account to fund
+    ///      reallocation
+    ///   4. `[]` (Optional) System program for reallocation funding
     ///
     /// Data expected by this instruction:
-    ///   None
+    ///   `ConfigureAccountWithRegistryInstructionData`
     ConfigureAccountWithRegistry,
 }
 
@@ -669,6 +675,16 @@ pub struct TransferWithFeeInstructionData {
     pub range_proof_instruction_offset: i8,
 }
 
+/// Data expected by `ConfidentialTransferInstruction::ConfigureAccountWithRegistry`
+#[cfg_attr(feature = "serde-traits", derive(Serialize, Deserialize))]
+#[derive(Clone, Copy, Debug, PartialEq, Pod, Zeroable)]
+#[repr(C)]
+pub struct ConfigureAccountWithRegistryInstructionData {
+    /// Reallocate token account if it is not large enough for the `ConfidentialTransfer`
+    /// extension.
+    pub reallocate_account: PodBool,
+}
+
 /// Create a `InitializeMint` instruction
 pub fn initialize_mint(
     token_program_id: &Pubkey,
@@ -1730,19 +1746,29 @@ pub fn configure_account_with_registry(
     token_account: &Pubkey,
     mint: &Pubkey,
     elgamal_registry_account: &Pubkey,
+    payer: Option<&Pubkey>,
 ) -> Result<Instruction, ProgramError> {
     check_program_account(token_program_id)?;
-    let accounts = vec![
+    let mut accounts = vec![
         AccountMeta::new(*token_account, false),
         AccountMeta::new_readonly(*mint, false),
         AccountMeta::new_readonly(*elgamal_registry_account, false),
     ];
+    let reallocate_account = if let Some(payer) = payer {
+        accounts.push(AccountMeta::new(*payer, true));
+        accounts.push(AccountMeta::new_readonly(system_program::id(), false));
+        true
+    } else {
+        false
+    };
 
     Ok(encode_instruction(
         token_program_id,
         accounts,
         TokenInstruction::ConfidentialTransferExtension,
         ConfidentialTransferInstruction::ConfigureAccountWithRegistry,
-        &(),
+        &ConfigureAccountWithRegistryInstructionData {
+            reallocate_account: reallocate_account.into(),
+        },
     ))
 }

token/program-2022/src/extension/confidential_transfer/processor.rs

@@ -15,22 +15,26 @@ use {
                 EncryptedWithheldAmount,
             },
             memo_transfer::{check_previous_sibling_instruction_is_memo, memo_required},
+            set_account_type,
             transfer_fee::TransferFeeConfig,
             transfer_hook, BaseStateWithExtensions, BaseStateWithExtensionsMut,
-            PodStateWithExtensions, PodStateWithExtensionsMut,
+            PodStateWithExtensions, PodStateWithExtensionsMut, StateWithExtensions,
         },
         instruction::{decode_instruction_data, decode_instruction_type},
         pod::{PodAccount, PodMint},
         processor::Processor,
+        state::Account,
     },
     solana_program::{
         account_info::{next_account_info, AccountInfo},
         clock::Clock,
         entrypoint::ProgramResult,
         msg,
+        program::invoke,
         program_error::ProgramError,
         pubkey::Pubkey,
-        sysvar::Sysvar,
+        system_instruction,
+        sysvar::{rent::Rent, Sysvar},
     },
     spl_elgamal_registry::state::ElGamalRegistry,
     spl_pod::bytemuck::pod_from_bytes,
@@ -103,10 +107,25 @@ enum ElGamalPubkeySource<'a> {
 fn process_configure_account_with_registry(
     program_id: &Pubkey,
     accounts: &[AccountInfo],
+    reallocate: bool,
 ) -> ProgramResult {
-    let elgamal_registry_account = accounts.get(2).ok_or(ProgramError::NotEnoughAccountKeys)?;
+    let account_info_iter = &mut accounts.iter();
+    let token_account_info = next_account_info(account_info_iter)?;
+    let _mint_info = next_account_info(account_info_iter)?;
+    let elgamal_registry_account = next_account_info(account_info_iter)?;
+
     check_elgamal_registry_program_account(elgamal_registry_account.owner)?;
 
+    if reallocate {
+        let payer_info = next_account_info(account_info_iter)?;
+        let system_program_info = next_account_info(account_info_iter)?;
+        reallocate_for_configure_account_with_registry(
+            token_account_info,
+            payer_info,
+            system_program_info,
+        )?;
+    }
+
     let elgamal_registry_account_data = &elgamal_registry_account.data.borrow();
     let elgamal_registry_account =
         pod_from_bytes::<ElGamalRegistry>(elgamal_registry_account_data)?;
@@ -124,6 +143,56 @@ fn process_configure_account_with_registry(
     )
 }
 
+fn reallocate_for_configure_account_with_registry<'a>(
+    token_account_info: &AccountInfo<'a>,
+    payer_info: &AccountInfo<'a>,
+    system_program_info: &AccountInfo<'a>,
+) -> ProgramResult {
+    let mut current_extension_types = {
+        let token_account = token_account_info.data.borrow();
+        let account = StateWithExtensions::<Account>::unpack(&token_account)?;
+        account.get_extension_types()?
+    };
+    current_extension_types.push(ExtensionType::ConfidentialTransferAccount);
+    let needed_account_len =
+        ExtensionType::try_calculate_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_rent_exempt_reserve = rent.minimum_balance(needed_account_len);
+
+    let current_lamport_reserve = token_account_info.lamports();
+    let lamports_diff = new_rent_exempt_reserve.saturating_sub(current_lamport_reserve);
+    if lamports_diff > 0 {
+        invoke(
+            &system_instruction::transfer(payer_info.key, token_account_info.key, lamports_diff),
+            &[
+                payer_info.clone(),
+                token_account_info.clone(),
+                system_program_info.clone(),
+            ],
+        )?;
+    }
+
+    // set account_type, if needed
+    let mut token_account_data = token_account_info.data.borrow_mut();
+    set_account_type::<Account>(&mut token_account_data)?;
+
+    Ok(())
+}
+
 /// Processes a [ConfigureAccount] instruction.
 fn process_configure_account(
     program_id: &Pubkey,
@@ -1272,7 +1341,13 @@ pub(crate) fn process_instruction(
         }
         ConfidentialTransferInstruction::ConfigureAccountWithRegistry => {
             msg!("ConfidentialTransferInstruction::ConfigureAccountWithRegistry");
-            process_configure_account_with_registry(program_id, accounts)
+            let data =
+                decode_instruction_data::<ConfigureAccountWithRegistryInstructionData>(input)?;
+            process_configure_account_with_registry(
+                program_id,
+                accounts,
+                data.reallocate_account.into(),
+            )
         }
     }
 }