token: Support both token programs in associated-token-account (#2904)

* token: Support both token programs in associated-token-account

* Rename spl-token -> spl-token-2022, spl-token-legacy -> spl-token

* Use spl-token by default

* Check mint in get_account_data_len

* Add checks for either token program

* Rebase fix

* Update version comments in tests

program/Cargo.toml

@@ -21,6 +21,7 @@ solana-program = "1.9.5"
 thiserror = "1.0"
 
 [dev-dependencies]
+lazy_static = "1.4.0"
 solana-program-test = "1.9.5"
 solana-sdk = "1.9.5"
 

program/src/instruction.rs

@@ -409,6 +409,31 @@ pub enum TokenInstruction {
         /// The freeze authority/multisignature of the mint.
         freeze_authority: COption<Pubkey>,
     },
+    /// Gets the required size of an account for the given mint as a little-endian
+    /// `u64`.
+    ///
+    /// Return data can be fetched using `sol_get_return_data` and deserializing
+    /// the return data as a little-endian `u64`.
+    ///
+    /// Accounts expected by this instruction:
+    ///
+    ///   0. `[]` The mint to calculate for
+    GetAccountDataSize, // typically, there's also data, but this program ignores it
+    /// Initialize the Immutable Owner extension for the given token account
+    ///
+    /// Fails if the account has already been initialized, so must be called before
+    /// `InitializeAccount`.
+    ///
+    /// No-ops in this version of the program, but is included for compatibility
+    /// with the Associated Token Account program.
+    ///
+    /// 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).
@@ -529,6 +554,8 @@ impl TokenInstruction {
                     decimals,
                 }
             }
+            21 => Self::GetAccountDataSize,
+            22 => Self::InitializeImmutableOwner,
             _ => return Err(TokenError::InvalidInstruction.into()),
         })
     }
@@ -625,6 +652,12 @@ impl TokenInstruction {
                 buf.extend_from_slice(mint_authority.as_ref());
                 Self::pack_pubkey_option(freeze_authority, &mut buf);
             }
+            &Self::GetAccountDataSize => {
+                buf.push(21);
+            }
+            &Self::InitializeImmutableOwner => {
+                buf.push(22);
+            }
         };
         buf
     }
@@ -1298,6 +1331,33 @@ pub fn sync_native(
     })
 }
 
+/// Creates a `GetAccountDataSize` instruction
+pub fn get_account_data_size(
+    token_program_id: &Pubkey,
+    mint_pubkey: &Pubkey,
+) -> Result<Instruction, ProgramError> {
+    check_program_account(token_program_id)?;
+
+    Ok(Instruction {
+        program_id: *token_program_id,
+        accounts: vec![AccountMeta::new(*mint_pubkey, false)],
+        data: TokenInstruction::GetAccountDataSize.pack(),
+    })
+}
+
+/// Creates a `InitializeImmutableOwner` instruction
+pub fn initialize_immutable_owner(
+    token_program_id: &Pubkey,
+    account_pubkey: &Pubkey,
+) -> Result<Instruction, ProgramError> {
+    check_program_account(token_program_id)?;
+    Ok(Instruction {
+        program_id: *token_program_id,
+        accounts: vec![AccountMeta::new(*account_pubkey, 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)
@@ -1518,5 +1578,19 @@ mod test {
         assert_eq!(packed, expect);
         let unpacked = TokenInstruction::unpack(&expect).unwrap();
         assert_eq!(unpacked, check);
+
+        let check = TokenInstruction::GetAccountDataSize;
+        let packed = check.pack();
+        let expect = vec![21u8];
+        assert_eq!(packed, expect);
+        let unpacked = TokenInstruction::unpack(&expect).unwrap();
+        assert_eq!(unpacked, check);
+
+        let check = TokenInstruction::InitializeImmutableOwner;
+        let packed = check.pack();
+        let expect = vec![22u8];
+        assert_eq!(packed, expect);
+        let unpacked = TokenInstruction::unpack(&expect).unwrap();
+        assert_eq!(unpacked, check);
     }
 }

program/src/processor.rs

@@ -11,6 +11,7 @@ use solana_program::{
     decode_error::DecodeError,
     entrypoint::ProgramResult,
     msg,
+    program::set_return_data,
     program_error::{PrintProgramError, ProgramError},
     program_memory::{sol_memcmp, sol_memset},
     program_option::COption,
@@ -757,6 +758,32 @@ impl Processor {
         Ok(())
     }
 
+    /// Processes a [GetAccountDataSize](enum.TokenInstruction.html) instruction
+    pub fn process_get_account_data_size(
+        program_id: &Pubkey,
+        accounts: &[AccountInfo],
+    ) -> ProgramResult {
+        let account_info_iter = &mut accounts.iter();
+        // make sure the mint is valid
+        let mint_info = next_account_info(account_info_iter)?;
+        Self::check_account_owner(program_id, mint_info)?;
+        let _ = Mint::unpack(&mint_info.data.borrow())?;
+        set_return_data(&Account::LEN.to_le_bytes());
+        Ok(())
+    }
+
+    /// Processes an [InitializeImmutableOwner](enum.TokenInstruction.html) instruction
+    pub fn process_initialize_immutable_owner(accounts: &[AccountInfo]) -> ProgramResult {
+        let account_info_iter = &mut accounts.iter();
+        let token_account_info = next_account_info(account_info_iter)?;
+        let account = Account::unpack_unchecked(&token_account_info.data.borrow())?;
+        if account.is_initialized() {
+            return Err(TokenError::AlreadyInUse.into());
+        }
+        msg!("Please upgrade to SPL Token 2022 for immutable owner support");
+        Ok(())
+    }
+
     /// Processes an [Instruction](enum.Instruction.html).
     pub fn process(program_id: &Pubkey, accounts: &[AccountInfo], input: &[u8]) -> ProgramResult {
         let instruction = TokenInstruction::unpack(input)?;
@@ -857,6 +884,14 @@ impl Processor {
                 msg!("Instruction: SyncNative");
                 Self::process_sync_native(program_id, accounts)
             }
+            TokenInstruction::GetAccountDataSize => {
+                msg!("Instruction: GetAccountDataSize");
+                Self::process_get_account_data_size(program_id, accounts)
+            }
+            TokenInstruction::InitializeImmutableOwner => {
+                msg!("Instruction: InitializeImmutableOwner");
+                Self::process_initialize_immutable_owner(accounts)
+            }
         }
     }
 
@@ -968,6 +1003,15 @@ mod tests {
     use solana_sdk::account::{
         create_account_for_test, create_is_signer_account_infos, Account as SolanaAccount,
     };
+    use std::sync::{Arc, RwLock};
+
+    lazy_static::lazy_static! {
+        static ref EXPECTED_DATA: Arc<RwLock<Vec<u8>>> = Arc::new(RwLock::new(Vec::new()));
+    }
+
+    fn set_expected_data(expected_data: Vec<u8>) {
+        *EXPECTED_DATA.write().unwrap() = expected_data;
+    }
 
     struct SyscallStubs {}
     impl solana_sdk::program_stubs::SyscallStubs for SyscallStubs {
@@ -1001,6 +1045,10 @@ mod tests {
             }
             solana_program::entrypoint::SUCCESS
         }
+
+        fn sol_set_return_data(&self, data: &[u8]) {
+            assert_eq!(&*EXPECTED_DATA.write().unwrap(), data)
+        }
     }
 
     fn do_process_instruction(
@@ -6355,4 +6403,88 @@ mod tests {
             )
         );
     }
+
+    #[test]
+    fn test_get_account_data_size() {
+        // see integration tests for return-data validity
+        let program_id = crate::id();
+        let owner_key = Pubkey::new_unique();
+        let mut rent_sysvar = rent_sysvar();
+        let mut mint_account =
+            SolanaAccount::new(mint_minimum_balance(), Mint::get_packed_len(), &program_id);
+        let mint_key = Pubkey::new_unique();
+        // fail if an invalid mint is passed in
+        assert_eq!(
+            Err(ProgramError::UninitializedAccount),
+            do_process_instruction(
+                get_account_data_size(&program_id, &mint_key).unwrap(),
+                vec![&mut mint_account],
+            )
+        );
+
+        do_process_instruction(
+            initialize_mint(&program_id, &mint_key, &owner_key, None, 2).unwrap(),
+            vec![&mut mint_account, &mut rent_sysvar],
+        )
+        .unwrap();
+
+        set_expected_data(Account::LEN.to_le_bytes().to_vec());
+        do_process_instruction(
+            get_account_data_size(&program_id, &mint_key).unwrap(),
+            vec![&mut mint_account],
+        )
+        .unwrap();
+    }
+
+    #[test]
+    fn test_initialize_immutable_owner() {
+        let program_id = crate::id();
+        let account_key = Pubkey::new_unique();
+        let mut account_account = SolanaAccount::new(
+            account_minimum_balance(),
+            Account::get_packed_len(),
+            &program_id,
+        );
+        let owner_key = Pubkey::new_unique();
+        let mut owner_account = SolanaAccount::default();
+        let mint_key = Pubkey::new_unique();
+        let mut mint_account =
+            SolanaAccount::new(mint_minimum_balance(), Mint::get_packed_len(), &program_id);
+        let mut rent_sysvar = rent_sysvar();
+
+        // create mint
+        do_process_instruction(
+            initialize_mint(&program_id, &mint_key, &owner_key, None, 2).unwrap(),
+            vec![&mut mint_account, &mut rent_sysvar],
+        )
+        .unwrap();
+
+        // success initialize immutable
+        do_process_instruction(
+            initialize_immutable_owner(&program_id, &account_key).unwrap(),
+            vec![&mut account_account],
+        )
+        .unwrap();
+
+        // create account
+        do_process_instruction(
+            initialize_account(&program_id, &account_key, &mint_key, &owner_key).unwrap(),
+            vec![
+                &mut account_account,
+                &mut mint_account,
+                &mut owner_account,
+                &mut rent_sysvar,
+            ],
+        )
+        .unwrap();
+
+        // fail post-init
+        assert_eq!(
+            Err(TokenError::AlreadyInUse.into()),
+            do_process_instruction(
+                initialize_immutable_owner(&program_id, &account_key).unwrap(),
+                vec![&mut account_account],
+            )
+        );
+    }
 }