Add GetAccountDataSize implementation (#2736)

* Add helper to get required Account extensions

* Add GetAccountDataSize processor

* Add get_account_data_size instruction

* Add test harness to check return data

* Add test of basic mint/account get-len

* Move method inside impl ExtensionType

Cargo.lock

@@ -3809,6 +3809,7 @@ version = "0.1.0"
 dependencies = [
  "arrayref",
  "bytemuck",
+ "lazy_static",
  "num-derive",
  "num-traits",
  "num_enum",

token/program-2022/Cargo.toml

@@ -23,6 +23,7 @@ solana-zk-token-sdk = "0.1.0"
 thiserror = "1.0"
 
 [dev-dependencies]
+lazy_static = "1.4.0"
 solana-program-test = "1.9.2"
 solana-sdk = "1.9.2"
 

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

@@ -496,6 +496,26 @@ impl ExtensionType {
             ExtensionType::MintPaddingTest => AccountType::Mint,
         }
     }
+
+    /// Get the list of required AccountType::Account ExtensionTypes based on a set of
+    /// AccountType::Mint ExtensionTypes
+    pub fn get_account_extensions(mint_extension_types: &[Self]) -> Vec<Self> {
+        let mut account_extension_types = vec![];
+        for extension_type in mint_extension_types {
+            #[allow(clippy::single_match)]
+            match extension_type {
+                ExtensionType::TransferFeeConfig => {
+                    account_extension_types.push(ExtensionType::TransferFeeAmount);
+                }
+                #[cfg(test)]
+                ExtensionType::MintPaddingTest => {
+                    account_extension_types.push(ExtensionType::AccountPaddingTest);
+                }
+                _ => {}
+            }
+        }
+        account_extension_types
+    }
 }
 
 /// Get the required account data length for the given ExtensionTypes
@@ -1101,4 +1121,53 @@ mod test {
         expect.extend_from_slice(&(ExtensionType::Uninitialized as u16).to_le_bytes());
         assert_eq!(expect, buffer);
     }
+
+    #[test]
+    fn test_get_account_extensions() {
+        // Some mint extensions with no required account extensions
+        let mint_extensions = vec![
+            ExtensionType::MintCloseAuthority,
+            ExtensionType::Uninitialized,
+        ];
+        assert_eq!(
+            ExtensionType::get_account_extensions(&mint_extensions),
+            vec![]
+        );
+
+        // One mint extension with required account extension, one without
+        let mint_extensions = vec![
+            ExtensionType::TransferFeeConfig,
+            ExtensionType::MintCloseAuthority,
+        ];
+        assert_eq!(
+            ExtensionType::get_account_extensions(&mint_extensions),
+            vec![ExtensionType::TransferFeeAmount]
+        );
+
+        // Some mint extensions both with required account extensions
+        let mint_extensions = vec![
+            ExtensionType::TransferFeeConfig,
+            ExtensionType::MintPaddingTest,
+        ];
+        assert_eq!(
+            ExtensionType::get_account_extensions(&mint_extensions),
+            vec![
+                ExtensionType::TransferFeeAmount,
+                ExtensionType::AccountPaddingTest
+            ]
+        );
+
+        // Demonstrate that method does not dedupe inputs or outputs
+        let mint_extensions = vec![
+            ExtensionType::TransferFeeConfig,
+            ExtensionType::TransferFeeConfig,
+        ];
+        assert_eq!(
+            ExtensionType::get_account_extensions(&mint_extensions),
+            vec![
+                ExtensionType::TransferFeeAmount,
+                ExtensionType::TransferFeeAmount
+            ]
+        );
+    }
 }

token/program-2022/src/instruction.rs

@@ -1405,6 +1405,19 @@ 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_readonly(*mint_pubkey, false)],
+        data: TokenInstruction::GetAccountDataSize.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/processor.rs

@@ -1,10 +1,11 @@
 //! Program state processor
 
 use crate::{
+    check_program_account,
     error::TokenError,
     extension::{
         confidential_transfer::{self, ConfidentialTransferAccount},
-        transfer_fee, StateWithExtensionsMut,
+        get_account_len, transfer_fee, ExtensionType, StateWithExtensions, StateWithExtensionsMut,
     },
     instruction::{is_valid_signer_index, AuthorityType, TokenInstruction, MAX_SIGNERS},
     state::{Account, AccountState, Mint, Multisig},
@@ -15,6 +16,7 @@ use solana_program::{
     decode_error::DecodeError,
     entrypoint::ProgramResult,
     msg,
+    program::set_return_data,
     program_error::{PrintProgramError, ProgramError},
     program_option::COption,
     program_pack::{IsInitialized, Pack},
@@ -761,6 +763,24 @@ impl Processor {
         unimplemented!();
     }
 
+    /// Processes a [GetAccountDataSize](enum.TokenInstruction.html) instruction
+    pub fn process_get_account_data_size(accounts: &[AccountInfo]) -> ProgramResult {
+        let account_info_iter = &mut accounts.iter();
+        let mint_account_info = next_account_info(account_info_iter)?;
+
+        check_program_account(mint_account_info.owner)?;
+        let mint_data = mint_account_info.data.borrow();
+        let state = StateWithExtensions::<Mint>::unpack(&mint_data)?;
+        let mint_extensions: Vec<ExtensionType> = state.get_extension_types()?;
+
+        let account_extensions = ExtensionType::get_account_extensions(&mint_extensions);
+
+        let account_len = get_account_len(&account_extensions);
+        set_return_data(&account_len.to_le_bytes());
+
+        Ok(())
+    }
+
     /// Processes an [Instruction](enum.Instruction.html).
     pub fn process(program_id: &Pubkey, accounts: &[AccountInfo], input: &[u8]) -> ProgramResult {
         let instruction = TokenInstruction::unpack(input)?;
@@ -863,7 +883,8 @@ impl Processor {
                 Self::process_sync_native(program_id, accounts)
             }
             TokenInstruction::GetAccountDataSize => {
-                unimplemented!();
+                msg!("Instruction: GetAccountDataSize");
+                Self::process_get_account_data_size(accounts)
             }
             TokenInstruction::InitializeMintCloseAuthority { close_authority } => {
                 msg!("Instruction: InitializeMintCloseAuthority");
@@ -1007,6 +1028,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 {
@@ -1040,6 +1070,10 @@ mod tests {
             }
             solana_program::entrypoint::SUCCESS
         }
+
+        fn sol_set_return_data(&mut self, data: &[u8]) {
+            assert_eq!(&*EXPECTED_DATA.read().unwrap(), data)
+        }
     }
 
     fn do_process_instruction(
@@ -6251,4 +6285,83 @@ 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 owner_account = SolanaAccount::default();
+        let mut rent_sysvar = rent_sysvar();
+
+        // Base mint
+        let mut mint_account =
+            SolanaAccount::new(mint_minimum_balance(), Mint::get_packed_len(), &program_id);
+        let mint_key = Pubkey::new_unique();
+        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(get_account_len(&[]).to_le_bytes().to_vec());
+        do_process_instruction(
+            get_account_data_size(&program_id, &mint_key).unwrap(),
+            vec![&mut mint_account],
+        )
+        .unwrap();
+
+        // TODO: Extended mint
+
+        // Invalid mint
+        let mut invalid_mint_account = SolanaAccount::new(
+            account_minimum_balance(),
+            Account::get_packed_len(),
+            &program_id,
+        );
+        let invalid_mint_key = Pubkey::new_unique();
+        do_process_instruction(
+            initialize_account(&program_id, &invalid_mint_key, &mint_key, &owner_key).unwrap(),
+            vec![
+                &mut invalid_mint_account,
+                &mut mint_account,
+                &mut owner_account,
+                &mut rent_sysvar,
+            ],
+        )
+        .unwrap();
+
+        assert_eq!(
+            do_process_instruction(
+                get_account_data_size(&program_id, &invalid_mint_key).unwrap(),
+                vec![&mut invalid_mint_account],
+            ),
+            Err(ProgramError::InvalidAccountData)
+        );
+
+        // Invalid mint owner
+        let invalid_program_id = Pubkey::new_unique();
+        let mut invalid_mint_account = SolanaAccount::new(
+            mint_minimum_balance(),
+            Mint::get_packed_len(),
+            &invalid_program_id,
+        );
+        let invalid_mint_key = Pubkey::new_unique();
+        let mut instruction =
+            initialize_mint(&program_id, &invalid_mint_key, &owner_key, None, 2).unwrap();
+        instruction.program_id = invalid_program_id;
+        do_process_instruction(
+            instruction,
+            vec![&mut invalid_mint_account, &mut rent_sysvar],
+        )
+        .unwrap();
+
+        assert_eq!(
+            do_process_instruction(
+                get_account_data_size(&program_id, &invalid_mint_key).unwrap(),
+                vec![&mut invalid_mint_account],
+            ),
+            Err(ProgramError::IncorrectProgramId)
+        );
+    }
 }