Add ImmutableOwner extension to block ATA owner authority changes

Author: mvines

associated-token-account/program/src/processor.rs

@@ -76,7 +76,11 @@ pub fn process_create_associated_token_account(
         &[bump_seed],
     ];
 
-    let account_len = get_account_len(spl_token_mint_info, spl_token_program_info)?;
+    let account_len = get_account_len(
+        spl_token_mint_info,
+        spl_token_program_info,
+        vec![spl_token::extension::ExtensionType::ImmutableOwner],
+    )?;
 
     create_pda_account(
         funder_info,
@@ -89,6 +93,16 @@ pub fn process_create_associated_token_account(
     )?;
 
     msg!("Initialize the associated token account");
+    invoke(
+        &spl_token::instruction::initialize_immutable_owner(
+            spl_token_program_id,
+            associated_token_account_info.key,
+        )?,
+        &[
+            associated_token_account_info.clone(),
+            spl_token_program_info.clone(),
+        ],
+    )?;
     invoke(
         &spl_token::instruction::initialize_account3(
             spl_token_program_id,

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

@@ -10,7 +10,7 @@ use {
         rent::Rent,
         system_instruction,
     },
-    spl_token::check_program_account,
+    spl_token::{check_program_account, extension::ExtensionType},
     std::convert::TryInto,
 };
 
@@ -76,9 +76,14 @@ pub fn create_pda_account<'a>(
 pub fn get_account_len<'a>(
     mint: &AccountInfo<'a>,
     spl_token_program: &AccountInfo<'a>,
+    extension_types: Vec<ExtensionType>,
 ) -> Result<usize, ProgramError> {
     invoke(
-        &spl_token::instruction::get_account_data_size(spl_token_program.key, mint.key, vec![])?,
+        &spl_token::instruction::get_account_data_size(
+            spl_token_program.key,
+            mint.key,
+            extension_types,
+        )?,
         &[mint.clone(), spl_token_program.clone()],
     )?;
     get_return_data()

associated-token-account/program/tests/process_create_associated_token_account.rs

@@ -5,9 +5,7 @@ mod program_test;
 
 use {
     program_test::program_test,
-    solana_program::{
-        instruction::*, program_pack::Pack, pubkey::Pubkey, system_instruction, sysvar,
-    },
+    solana_program::{instruction::*, pubkey::Pubkey, system_instruction, sysvar},
     solana_program_test::*,
     solana_sdk::{
         signature::Signer,
@@ -16,6 +14,7 @@ use {
     spl_associated_token_account::{
         get_associated_token_address, instruction::create_associated_token_account,
     },
+    spl_token::extension::ExtensionType,
 };
 
 #[allow(deprecated)]
@@ -31,7 +30,12 @@ async fn test_associated_token_address() {
     let (mut banks_client, payer, recent_blockhash) =
         program_test(token_mint_address, true).start().await;
     let rent = banks_client.get_rent().await.unwrap();
-    let expected_token_account_balance = rent.minimum_balance(spl_token::state::Account::LEN);
+
+    let expected_token_account_len =
+        ExtensionType::get_account_len::<spl_token::state::Account>(&[
+            ExtensionType::ImmutableOwner,
+        ]);
+    let expected_token_account_balance = rent.minimum_balance(expected_token_account_len);
 
     // Associated account does not exist
     assert_eq!(
@@ -60,10 +64,7 @@ async fn test_associated_token_address() {
         .await
         .expect("get_account")
         .expect("associated_account not none");
-    assert_eq!(
-        associated_account.data.len(),
-        spl_token::state::Account::LEN
-    );
+    assert_eq!(associated_account.data.len(), expected_token_account_len,);
     assert_eq!(associated_account.owner, spl_token::id());
     assert_eq!(associated_account.lamports, expected_token_account_balance);
 }
@@ -78,7 +79,11 @@ async fn test_create_with_fewer_lamports() {
     let (mut banks_client, payer, recent_blockhash) =
         program_test(token_mint_address, true).start().await;
     let rent = banks_client.get_rent().await.unwrap();
-    let expected_token_account_balance = rent.minimum_balance(spl_token::state::Account::LEN);
+    let expected_token_account_len =
+        ExtensionType::get_account_len::<spl_token::state::Account>(&[
+            ExtensionType::ImmutableOwner,
+        ]);
+    let expected_token_account_balance = rent.minimum_balance(expected_token_account_len);
 
     // Transfer lamports into `associated_token_address` before creating it - enough to be
     // rent-exempt for 0 data, but not for an initialized token account
@@ -133,7 +138,12 @@ async fn test_create_with_excess_lamports() {
     let (mut banks_client, payer, recent_blockhash) =
         program_test(token_mint_address, true).start().await;
     let rent = banks_client.get_rent().await.unwrap();
-    let expected_token_account_balance = rent.minimum_balance(spl_token::state::Account::LEN);
+
+    let expected_token_account_len =
+        ExtensionType::get_account_len::<spl_token::state::Account>(&[
+            ExtensionType::ImmutableOwner,
+        ]);
+    let expected_token_account_balance = rent.minimum_balance(expected_token_account_len);
 
     // Transfer 1 lamport into `associated_token_address` before creating it
     let mut transaction = Transaction::new_with_payer(
@@ -255,7 +265,11 @@ async fn test_create_associated_token_account_using_legacy_implicit_instruction(
     let (mut banks_client, payer, recent_blockhash) =
         program_test(token_mint_address, true).start().await;
     let rent = banks_client.get_rent().await.unwrap();
-    let expected_token_account_balance = rent.minimum_balance(spl_token::state::Account::LEN);
+    let expected_token_account_len =
+        ExtensionType::get_account_len::<spl_token::state::Account>(&[
+            ExtensionType::ImmutableOwner,
+        ]);
+    let expected_token_account_balance = rent.minimum_balance(expected_token_account_len);
 
     // Associated account does not exist
     assert_eq!(
@@ -290,10 +304,7 @@ async fn test_create_associated_token_account_using_legacy_implicit_instruction(
         .await
         .expect("get_account")
         .expect("associated_account not none");
-    assert_eq!(
-        associated_account.data.len(),
-        spl_token::state::Account::LEN
-    );
+    assert_eq!(associated_account.data.len(), expected_token_account_len);
     assert_eq!(associated_account.owner, spl_token::id());
     assert_eq!(associated_account.lamports, expected_token_account_balance);
 }
@@ -308,7 +319,11 @@ async fn test_create_associated_token_account_using_deprecated_instruction_creat
     let (mut banks_client, payer, recent_blockhash) =
         program_test(token_mint_address, true).start().await;
     let rent = banks_client.get_rent().await.unwrap();
-    let expected_token_account_balance = rent.minimum_balance(spl_token::state::Account::LEN);
+    let expected_token_account_len =
+        ExtensionType::get_account_len::<spl_token::state::Account>(&[
+            ExtensionType::ImmutableOwner,
+        ]);
+    let expected_token_account_balance = rent.minimum_balance(expected_token_account_len);
 
     // Associated account does not exist
     assert_eq!(
@@ -338,10 +353,7 @@ async fn test_create_associated_token_account_using_deprecated_instruction_creat
         .await
         .expect("get_account")
         .expect("associated_account not none");
-    assert_eq!(
-        associated_account.data.len(),
-        spl_token::state::Account::LEN
-    );
+    assert_eq!(associated_account.data.len(), expected_token_account_len);
     assert_eq!(associated_account.owner, spl_token::id());
     assert_eq!(associated_account.lamports, expected_token_account_balance);
 }

token/program-2022/src/error.rs

@@ -121,6 +121,9 @@ pub enum TokenError {
     /// Calculated fee does not match expected fee
     #[error("Calculated fee does not match expected fee")]
     FeeMismatch,
+    /// The owner authority cannot be changed
+    #[error("The owner authority cannot be changed")]
+    ImmutableOwner,
 }
 impl From<TokenError> for ProgramError {
     fn from(e: TokenError) -> Self {

token/program-2022/src/extension/immutable_owner.rs

@@ -0,0 +1,13 @@
+use {
+    crate::extension::{Extension, ExtensionType},
+    bytemuck::{Pod, Zeroable},
+};
+
+/// Indicates that the Account owner authority cannot be changed
+#[derive(Clone, Copy, Debug, Default, PartialEq, Pod, Zeroable)]
+#[repr(transparent)]
+pub struct ImmutableOwner;
+
+impl Extension for ImmutableOwner {
+    const TYPE: ExtensionType = ExtensionType::ImmutableOwner;
+}

token/program-2022/src/extension/mod.rs

@@ -6,6 +6,7 @@ use {
         extension::{
             confidential_transfer::{ConfidentialTransferAccount, ConfidentialTransferMint},
             default_account_state::DefaultAccountState,
+            immutable_owner::ImmutableOwner,
             mint_close_authority::MintCloseAuthority,
             transfer_fee::{TransferFeeAmount, TransferFeeConfig},
         },
@@ -28,6 +29,8 @@ use {
 pub mod confidential_transfer;
 /// Default Account State extension
 pub mod default_account_state;
+/// Immutable Owner extension
+pub mod immutable_owner;
 /// Mint Close Authority extension
 pub mod mint_close_authority;
 /// Transfer Fee extension
@@ -64,6 +67,7 @@ fn get_tlv_indices(type_start: usize) -> TlvIndices {
 
 /// Helper struct for returning the indices of the type, length, and value in
 /// a TLV entry
+#[derive(Debug)]
 struct TlvIndices {
     pub type_start: usize,
     pub length_start: usize,
@@ -77,7 +81,7 @@ fn get_extension_indices<V: Extension>(
     let v_account_type = V::TYPE.get_account_type();
     while start_index < tlv_data.len() {
         let tlv_indices = get_tlv_indices(start_index);
-        if tlv_data.len() <= tlv_indices.value_start {
+        if tlv_data.len() < tlv_indices.value_start {
             return Err(ProgramError::InvalidAccountData);
         }
         let extension_type =
@@ -397,6 +401,7 @@ impl<'data, S: BaseState> StateWithExtensionsMut<'data, S> {
             length_start,
             value_start,
         } = get_extension_indices::<V>(self.tlv_data, init)?;
+
         if self.tlv_data[type_start..].len() < V::TYPE.get_tlv_len() {
             return Err(ProgramError::InvalidAccountData);
         }
@@ -551,6 +556,8 @@ pub enum ExtensionType {
     ConfidentialTransferAccount,
     /// Specifies the default Account::state for new Accounts
     DefaultAccountState,
+    /// Indicates that the Account owner authority cannot be changed
+    ImmutableOwner,
     /// Padding extension used to make an account exactly Multisig::LEN, used for testing
     #[cfg(test)]
     AccountPaddingTest = u16::MAX - 1,
@@ -580,6 +587,7 @@ impl ExtensionType {
             ExtensionType::TransferFeeConfig => pod_get_packed_len::<TransferFeeConfig>(),
             ExtensionType::TransferFeeAmount => pod_get_packed_len::<TransferFeeAmount>(),
             ExtensionType::MintCloseAuthority => pod_get_packed_len::<MintCloseAuthority>(),
+            ExtensionType::ImmutableOwner => pod_get_packed_len::<ImmutableOwner>(),
             ExtensionType::ConfidentialTransferMint => {
                 pod_get_packed_len::<ConfidentialTransferMint>()
             }
@@ -631,9 +639,9 @@ impl ExtensionType {
             | ExtensionType::MintCloseAuthority
             | ExtensionType::ConfidentialTransferMint
             | ExtensionType::DefaultAccountState => AccountType::Mint,
-            ExtensionType::TransferFeeAmount | ExtensionType::ConfidentialTransferAccount => {
-                AccountType::Account
-            }
+            ExtensionType::ImmutableOwner
+            | ExtensionType::TransferFeeAmount
+            | ExtensionType::ConfidentialTransferAccount => AccountType::Account,
             #[cfg(test)]
             ExtensionType::AccountPaddingTest => AccountType::Account,
             #[cfg(test)]

token/program-2022/src/instruction.rs

@@ -487,6 +487,19 @@ pub enum TokenInstruction {
     /// See `extension::default_account_state::instruction::DefaultAccountStateInstruction` for
     /// further details about the extended instructions that share this instruction prefix
     DefaultAccountStateExtension,
+    /// Initialize the Immutable Owner extension for the given token account
+    ///
+    /// Fails if the account has already been initialized, so must be called before
+    /// `InitializeAccount`.
+    ///
+    /// Accounts expected by this instruction:
+    ///
+    ///   0. `[writable]`  The account to initialize.
+    //
+    /// Data expected by this instruction:
+    ///   None
+    ///
+    InitializeImmutableOwner,
 }
 impl TokenInstruction {
     /// Unpacks a byte buffer into a [TokenInstruction](enum.TokenInstruction.html).
@@ -597,6 +610,7 @@ impl TokenInstruction {
             }
             24 => Self::ConfidentialTransferExtension,
             25 => Self::DefaultAccountStateExtension,
+            26 => Self::InitializeImmutableOwner,
             _ => return Err(TokenError::InvalidInstruction.into()),
         })
     }
@@ -718,6 +732,9 @@ impl TokenInstruction {
             &Self::DefaultAccountStateExtension => {
                 buf.push(25);
             }
+            &Self::InitializeImmutableOwner => {
+                buf.push(26);
+            }
         };
         buf
     }
@@ -1456,6 +1473,19 @@ pub fn initialize_mint_close_authority(
     })
 }
 
+/// Create an `InitializeImmutableOwner` instruction
+pub fn initialize_immutable_owner(
+    token_program_id: &Pubkey,
+    token_account: &Pubkey,
+) -> Result<Instruction, ProgramError> {
+    check_program_account(token_program_id)?;
+    Ok(Instruction {
+        program_id: *token_program_id,
+        accounts: vec![AccountMeta::new(*token_account, false)],
+        data: TokenInstruction::InitializeImmutableOwner.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)