token: Reassign and reallocate accounts on close (#3415)

joncinque · web-flow · commit 4fadd553e1c5 · 2022-08-03T23:49:36.000+02:00
* token: Reassign and reallocate accounts on close * Revert "Refactor unpack and make test more robust (#3417)" This reverts commit c618de3. * Revert "check that unpack is tolerant of small sizes (#3416)" This reverts commit 22faa05. * Also revert d7f352b
diff --git a/token/program-2022-test/tests/close_account.rs b/token/program-2022-test/tests/close_account.rs
@@ -0,0 +1,167 @@
+#![cfg(feature = "test-bpf")]
+
+mod program_test;
+use {
+    program_test::TestContext,
+    solana_program_test::tokio,
+    solana_sdk::{
+        instruction::InstructionError, program_pack::Pack, pubkey::Pubkey, signature::Signer,
+        signer::keypair::Keypair, system_instruction, transaction::TransactionError,
+        transport::TransportError,
+    },
+    spl_token_2022::{instruction, state::Account},
+    spl_token_client::token::{ExtensionInitializationParams, TokenError as TokenClientError},
+};
+
+#[tokio::test]
+async fn success_init_after_close_account() {
+    let mut context = TestContext::new().await;
+    let payer = Keypair::from_bytes(&context.context.lock().await.payer.to_bytes()).unwrap();
+    context.init_token_with_mint(vec![]).await.unwrap();
+    let token = context.token_context.take().unwrap().token;
+    let token_program_id = spl_token_2022::id();
+    let owner = Keypair::new();
+    let token_account_keypair = Keypair::new();
+    let token_account = token
+        .create_auxiliary_token_account(&token_account_keypair, &owner.pubkey())
+        .await
+        .unwrap();
+
+    let destination = Pubkey::new_unique();
+    token
+        .process_ixs(
+            &[
+                instruction::close_account(
+                    &token_program_id,
+                    &token_account,
+                    &destination,
+                    &owner.pubkey(),
+                    &[],
+                )
+                .unwrap(),
+                system_instruction::create_account(
+                    &payer.pubkey(),
+                    &token_account,
+                    1_000_000_000,
+                    Account::LEN as u64,
+                    &token_program_id,
+                ),
+                instruction::initialize_account(
+                    &token_program_id,
+                    &token_account,
+                    token.get_address(),
+                    &owner.pubkey(),
+                )
+                .unwrap(),
+            ],
+            &[&owner, &payer, &token_account_keypair],
+        )
+        .await
+        .unwrap();
+    let destination = token.get_account(&destination).await.unwrap();
+    assert!(destination.lamports > 0);
+}
+
+#[tokio::test]
+async fn fail_init_after_close_account() {
+    let mut context = TestContext::new().await;
+    let payer = Keypair::from_bytes(&context.context.lock().await.payer.to_bytes()).unwrap();
+    context.init_token_with_mint(vec![]).await.unwrap();
+    let token = context.token_context.take().unwrap().token;
+    let token_program_id = spl_token_2022::id();
+    let owner = Keypair::new();
+    let token_account = token
+        .create_auxiliary_token_account(&Keypair::new(), &owner.pubkey())
+        .await
+        .unwrap();
+
+    let destination = Pubkey::new_unique();
+    let error = token
+        .process_ixs(
+            &[
+                instruction::close_account(
+                    &token_program_id,
+                    &token_account,
+                    &destination,
+                    &owner.pubkey(),
+                    &[],
+                )
+                .unwrap(),
+                system_instruction::transfer(&payer.pubkey(), &token_account, 1_000_000_000),
+                instruction::initialize_account(
+                    &token_program_id,
+                    &token_account,
+                    token.get_address(),
+                    &owner.pubkey(),
+                )
+                .unwrap(),
+            ],
+            &[&owner, &payer],
+        )
+        .await
+        .unwrap_err();
+    assert_eq!(
+        error,
+        TokenClientError::Client(Box::new(TransportError::TransactionError(
+            TransactionError::InstructionError(2, InstructionError::InvalidAccountData,)
+        )))
+    );
+    let error = token.get_account(&destination).await.unwrap_err();
+    assert_eq!(error, TokenClientError::AccountNotFound);
+}
+
+#[tokio::test]
+async fn fail_init_after_close_mint() {
+    let close_authority = Keypair::new();
+    let mut context = TestContext::new().await;
+    let payer = Keypair::from_bytes(&context.context.lock().await.payer.to_bytes()).unwrap();
+    context
+        .init_token_with_mint(vec![ExtensionInitializationParams::MintCloseAuthority {
+            close_authority: Some(close_authority.pubkey()),
+        }])
+        .await
+        .unwrap();
+    let token = context.token_context.take().unwrap().token;
+    let token_program_id = spl_token_2022::id();
+
+    let destination = Pubkey::new_unique();
+    let error = token
+        .process_ixs(
+            &[
+                instruction::close_account(
+                    &token_program_id,
+                    token.get_address(),
+                    &destination,
+                    &close_authority.pubkey(),
+                    &[],
+                )
+                .unwrap(),
+                system_instruction::transfer(&payer.pubkey(), token.get_address(), 1_000_000_000),
+                instruction::initialize_mint_close_authority(
+                    &token_program_id,
+                    token.get_address(),
+                    None,
+                )
+                .unwrap(),
+                instruction::initialize_mint(
+                    &token_program_id,
+                    token.get_address(),
+                    &close_authority.pubkey(),
+                    None,
+                    0,
+                )
+                .unwrap(),
+            ],
+            &[&close_authority, &payer],
+        )
+        .await
+        .unwrap_err();
+    assert_eq!(
+        error,
+        TokenClientError::Client(Box::new(TransportError::TransactionError(
+            TransactionError::InstructionError(2, InstructionError::InvalidAccountData,)
+        )))
+    );
+    let error = token.get_account(&destination).await.unwrap_err();
+    assert_eq!(error, TokenClientError::AccountNotFound);
+}
diff --git a/token/program-2022/src/instruction.rs b/token/program-2022/src/instruction.rs
@@ -1787,7 +1787,7 @@ pub(crate) fn encode_instruction<T: Into<u8>, D: Pod>(
 
 #[cfg(test)]
 mod test {
-    use {super::*, proptest::prelude::*};
+    use super::*;
 
     #[test]
     fn test_instruction_packing() {
@@ -2284,14 +2284,4 @@ mod test {
             ui_amount,
         ));
     }
-
-    proptest! {
-        #![proptest_config(ProptestConfig::with_cases(1024))]
-        #[test]
-        fn test_instruction_unpack_panic(
-            data in prop::collection::vec(any::<u8>(), 0..255)
-        ) {
-            let _no_panic = TokenInstruction::unpack(&data);
-        }
-    }
 }
diff --git a/token/program-2022/src/processor.rs b/token/program-2022/src/processor.rs
@@ -27,11 +27,10 @@ use {
         msg,
         program::{invoke, invoke_signed, set_return_data},
         program_error::ProgramError,
-        program_memory::sol_memset,
         program_option::COption,
         program_pack::Pack,
         pubkey::Pubkey,
-        system_instruction,
+        system_instruction, system_program,
         sysvar::{rent::Rent, Sysvar},
     },
     std::convert::{TryFrom, TryInto},
@@ -875,7 +874,7 @@ impl Processor {
             return Err(ProgramError::InvalidAccountData);
         }
 
-        let mut source_account_data = source_account_info.data.borrow_mut();
+        let source_account_data = source_account_info.data.borrow();
         if let Ok(source_account) = StateWithExtensions::<Account>::unpack(&source_account_data) {
             if !source_account.base.is_native() && source_account.base.amount != 0 {
                 return Err(TokenError::NonNativeHasBalance.into());
@@ -935,8 +934,8 @@ impl Processor {
             .ok_or(TokenError::Overflow)?;
 
         **source_account_info.lamports.borrow_mut() = 0;
-        let data_len = source_account_data.len();
-        sol_memset(*source_account_data, 0, data_len);
+        drop(source_account_data);
+        delete_account(source_account_info)?;
 
         Ok(())
     }
@@ -1388,6 +1387,25 @@ impl Processor {
     }
 }
 
+/// Helper function to mostly delete an account in a test environment.  We could
+/// potentially muck around the bytes assuming that a vec is passed in, but that
+/// would be more trouble than it's worth.
+#[cfg(not(target_arch = "bpf"))]
+fn delete_account(account_info: &AccountInfo) -> Result<(), ProgramError> {
+    account_info.assign(&system_program::id());
+    let mut account_data = account_info.data.borrow_mut();
+    let data_len = account_data.len();
+    solana_program::program_memory::sol_memset(*account_data, 0, data_len);
+    Ok(())
+}
+
+/// Helper function to totally delete an account on-chain
+#[cfg(target_arch = "bpf")]
+fn delete_account(account_info: &AccountInfo) -> Result<(), ProgramError> {
+    account_info.assign(&system_program::id());
+    account_info.realloc(0, false)
+}
+
 #[cfg(test)]
 mod tests {
     use {
diff --git a/token/program/src/instruction.rs b/token/program/src/instruction.rs
@@ -15,8 +15,6 @@ use std::mem::size_of;
 pub const MIN_SIGNERS: usize = 1;
 /// Maximum number of multisignature signers (max N)
 pub const MAX_SIGNERS: usize = 11;
-/// Serialized length of a u64, for unpacking
-const U64_BYTES: usize = 8;
 
 /// Instructions supported by the token program.
 #[repr(C)]
@@ -521,19 +519,47 @@ impl<'a> TokenInstruction<'a> {
             10 => Self::FreezeAccount,
             11 => Self::ThawAccount,
             12 => {
-                let (amount, decimals, _rest) = Self::unpack_amount_decimals(rest)?;
+                let (amount, rest) = rest.split_at(8);
+                let amount = amount
+                    .try_into()
+                    .ok()
+                    .map(u64::from_le_bytes)
+                    .ok_or(InvalidInstruction)?;
+                let (&decimals, _rest) = rest.split_first().ok_or(InvalidInstruction)?;
+
                 Self::TransferChecked { amount, decimals }
             }
             13 => {
-                let (amount, decimals, _rest) = Self::unpack_amount_decimals(rest)?;
+                let (amount, rest) = rest.split_at(8);
+                let amount = amount
+                    .try_into()
+                    .ok()
+                    .map(u64::from_le_bytes)
+                    .ok_or(InvalidInstruction)?;
+                let (&decimals, _rest) = rest.split_first().ok_or(InvalidInstruction)?;
+
                 Self::ApproveChecked { amount, decimals }
             }
             14 => {
-                let (amount, decimals, _rest) = Self::unpack_amount_decimals(rest)?;
+                let (amount, rest) = rest.split_at(8);
+                let amount = amount
+                    .try_into()
+                    .ok()
+                    .map(u64::from_le_bytes)
+                    .ok_or(InvalidInstruction)?;
+                let (&decimals, _rest) = rest.split_first().ok_or(InvalidInstruction)?;
+
                 Self::MintToChecked { amount, decimals }
             }
             15 => {
-                let (amount, decimals, _rest) = Self::unpack_amount_decimals(rest)?;
+                let (amount, rest) = rest.split_at(8);
+                let amount = amount
+                    .try_into()
+                    .ok()
+                    .map(u64::from_le_bytes)
+                    .ok_or(InvalidInstruction)?;
+                let (&decimals, _rest) = rest.split_first().ok_or(InvalidInstruction)?;
+
                 Self::BurnChecked { amount, decimals }
             }
             16 => {
@@ -562,7 +588,12 @@ impl<'a> TokenInstruction<'a> {
             21 => Self::GetAccountDataSize,
             22 => Self::InitializeImmutableOwner,
             23 => {
-                let (amount, _rest) = Self::unpack_u64(rest)?;
+                let (amount, _rest) = rest.split_at(8);
+                let amount = amount
+                    .try_into()
+                    .ok()
+                    .map(u64::from_le_bytes)
+                    .ok_or(InvalidInstruction)?;
                 Self::AmountToUiAmount { amount }
             }
             24 => {
@@ -714,21 +745,6 @@ impl<'a> TokenInstruction<'a> {
             COption::None => buf.push(0),
         }
     }
-
-    fn unpack_u64(input: &[u8]) -> Result<(u64, &[u8]), ProgramError> {
-        let value = input
-            .get(..U64_BYTES)
-            .and_then(|slice| slice.try_into().ok())
-            .map(u64::from_le_bytes)
-            .ok_or(TokenError::InvalidInstruction)?;
-        Ok((value, &input[U64_BYTES..]))
-    }
-
-    fn unpack_amount_decimals(input: &[u8]) -> Result<(u64, u8, &[u8]), ProgramError> {
-        let (amount, rest) = Self::unpack_u64(input)?;
-        let (&decimals, rest) = rest.split_first().ok_or(TokenError::InvalidInstruction)?;
-        Ok((amount, decimals, rest))
-    }
 }
 
 /// Specifies the authority type for SetAuthority instructions
@@ -1431,7 +1447,7 @@ pub fn is_valid_signer_index(index: usize) -> bool {
 
 #[cfg(test)]
 mod test {
-    use {super::*, proptest::prelude::*};
+    use super::*;
 
     #[test]
     fn test_instruction_packing() {
@@ -1673,14 +1689,4 @@ mod test {
         let unpacked = TokenInstruction::unpack(&expect).unwrap();
         assert_eq!(unpacked, check);
     }
-
-    proptest! {
-        #![proptest_config(ProptestConfig::with_cases(1024))]
-        #[test]
-        fn test_instruction_unpack_panic(
-            data in prop::collection::vec(any::<u8>(), 0..255)
-        ) {
-            let _no_panic = TokenInstruction::unpack(&data);
-        }
-    }
 }
diff --git a/token/program/src/processor.rs b/token/program/src/processor.rs
diff --git a/token/program/tests/close_account.rs b/token/program/tests/close_account.rs
