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

* token: Reassign and reallocate accounts on close

* Revert "Refactor unpack and make test more robust (#3417)"

This reverts commit 3cd0c9a82209eaf445662c00122188155594ad4a.

* Revert "check that unpack is tolerant of small sizes (#3416)"

This reverts commit c202d51918c994230fc71f9b8829a7f7300d5d12.

* Also revert 842d4bfad96329968121d0460403268fa73ffe4a

Author: joncinque

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);
-        }
-    }
 }

program/src/processor.rs

@@ -13,10 +13,11 @@ use solana_program::{
     msg,
     program::set_return_data,
     program_error::ProgramError,
-    program_memory::{sol_memcmp, sol_memset},
+    program_memory::sol_memcmp,
     program_option::COption,
     program_pack::{IsInitialized, Pack},
     pubkey::{Pubkey, PUBKEY_BYTES},
+    system_program,
     sysvar::{rent::Rent, Sysvar},
 };
 
@@ -696,8 +697,7 @@ impl Processor {
             .ok_or(TokenError::Overflow)?;
 
         **source_account_info.lamports.borrow_mut() = 0;
-
-        sol_memset(*source_account_info.data.borrow_mut(), 0, Account::LEN);
+        delete_account(source_account_info)?;
 
         Ok(())
     }
@@ -1007,6 +1007,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 super::*;

program/tests/close_account.rs

@@ -0,0 +1,202 @@
+#![cfg(feature = "test-bpf")]
+
+use {
+    solana_program_test::{processor, tokio, ProgramTest, ProgramTestContext},
+    solana_sdk::{
+        instruction::InstructionError,
+        program_pack::Pack,
+        pubkey::Pubkey,
+        signature::Signer,
+        signer::keypair::Keypair,
+        system_instruction,
+        transaction::{Transaction, TransactionError},
+    },
+    spl_token::{
+        instruction,
+        processor::Processor,
+        state::{Account, Mint},
+    },
+};
+
+async fn setup_mint_and_account(
+    context: &mut ProgramTestContext,
+    mint: &Keypair,
+    token_account: &Keypair,
+    owner: &Pubkey,
+    token_program_id: &Pubkey,
+) {
+    let rent = context.banks_client.get_rent().await.unwrap();
+    let mint_authority_pubkey = Pubkey::new_unique();
+
+    let space = Mint::LEN;
+    let tx = Transaction::new_signed_with_payer(
+        &[
+            system_instruction::create_account(
+                &context.payer.pubkey(),
+                &mint.pubkey(),
+                rent.minimum_balance(space),
+                space as u64,
+                &token_program_id,
+            ),
+            instruction::initialize_mint(
+                &token_program_id,
+                &mint.pubkey(),
+                &mint_authority_pubkey,
+                None,
+                9,
+            )
+            .unwrap(),
+        ],
+        Some(&context.payer.pubkey()),
+        &[&context.payer, mint],
+        context.last_blockhash,
+    );
+    context.banks_client.process_transaction(tx).await.unwrap();
+    let space = Account::LEN;
+    let tx = Transaction::new_signed_with_payer(
+        &[
+            system_instruction::create_account(
+                &context.payer.pubkey(),
+                &token_account.pubkey(),
+                rent.minimum_balance(space),
+                space as u64,
+                &token_program_id,
+            ),
+            instruction::initialize_account(
+                &token_program_id,
+                &token_account.pubkey(),
+                &mint.pubkey(),
+                &owner,
+            )
+            .unwrap(),
+        ],
+        Some(&context.payer.pubkey()),
+        &[&context.payer, &token_account],
+        context.last_blockhash,
+    );
+    context.banks_client.process_transaction(tx).await.unwrap();
+}
+
+#[tokio::test]
+async fn success_init_after_close_account() {
+    let program_test =
+        ProgramTest::new("spl_token", spl_token::id(), processor!(Processor::process));
+    let mut context = program_test.start_with_context().await;
+    let mint = Keypair::new();
+    let token_account = Keypair::new();
+    let owner = Keypair::new();
+    let token_program_id = spl_token::id();
+    setup_mint_and_account(
+        &mut context,
+        &mint,
+        &token_account,
+        &owner.pubkey(),
+        &token_program_id,
+    )
+    .await;
+
+    let destination = Pubkey::new_unique();
+    let tx = Transaction::new_signed_with_payer(
+        &[
+            instruction::close_account(
+                &token_program_id,
+                &token_account.pubkey(),
+                &destination,
+                &owner.pubkey(),
+                &[],
+            )
+            .unwrap(),
+            system_instruction::create_account(
+                &context.payer.pubkey(),
+                &token_account.pubkey(),
+                1_000_000_000,
+                Account::LEN as u64,
+                &token_program_id,
+            ),
+            instruction::initialize_account(
+                &token_program_id,
+                &token_account.pubkey(),
+                &mint.pubkey(),
+                &owner.pubkey(),
+            )
+            .unwrap(),
+        ],
+        Some(&context.payer.pubkey()),
+        &[&context.payer, &owner, &token_account],
+        context.last_blockhash,
+    );
+    context.banks_client.process_transaction(tx).await.unwrap();
+    let destination = context
+        .banks_client
+        .get_account(destination)
+        .await
+        .unwrap()
+        .unwrap();
+    assert!(destination.lamports > 0);
+}
+
+#[tokio::test]
+async fn fail_init_after_close_account() {
+    let program_test =
+        ProgramTest::new("spl_token", spl_token::id(), processor!(Processor::process));
+    let mut context = program_test.start_with_context().await;
+    let mint = Keypair::new();
+    let token_account = Keypair::new();
+    let owner = Keypair::new();
+    let token_program_id = spl_token::id();
+    setup_mint_and_account(
+        &mut context,
+        &mint,
+        &token_account,
+        &owner.pubkey(),
+        &token_program_id,
+    )
+    .await;
+
+    let destination = Pubkey::new_unique();
+    let tx = Transaction::new_signed_with_payer(
+        &[
+            instruction::close_account(
+                &token_program_id,
+                &token_account.pubkey(),
+                &destination,
+                &owner.pubkey(),
+                &[],
+            )
+            .unwrap(),
+            system_instruction::transfer(
+                &context.payer.pubkey(),
+                &token_account.pubkey(),
+                1_000_000_000,
+            ),
+            instruction::initialize_account(
+                &token_program_id,
+                &token_account.pubkey(),
+                &mint.pubkey(),
+                &owner.pubkey(),
+            )
+            .unwrap(),
+        ],
+        Some(&context.payer.pubkey()),
+        &[&context.payer, &owner],
+        context.last_blockhash,
+    );
+    #[allow(clippy::useless_conversion)]
+    let error: TransactionError = context
+        .banks_client
+        .process_transaction(tx)
+        .await
+        .unwrap_err()
+        .unwrap()
+        .into();
+    assert_eq!(
+        error,
+        TransactionError::InstructionError(2, InstructionError::InvalidAccountData)
+    );
+    assert!(context
+        .banks_client
+        .get_account(destination)
+        .await
+        .unwrap()
+        .is_none());
+}