From 4fadd553e1c549afd1d62aeb5ffa7ef31d1999d1 Mon Sep 17 00:00:00 2001 From: Jon Cinque Date: Wed, 3 Aug 2022 23:49:36 +0200 Subject: [PATCH] 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 c618de374e7c15c5abb443b6e35219c6257f0355. * Revert "check that unpack is tolerant of small sizes (#3416)" This reverts commit 22faa05f184853a72791d239b1bde025f110493d. * Also revert d7f352bd7891390a9401082e5082f298bf4c1913 --- .../program-2022-test/tests/close_account.rs | 167 +++++++++++++++ token/program-2022/src/instruction.rs | 12 +- token/program-2022/src/processor.rs | 28 ++- token/program/src/instruction.rs | 72 ++++--- token/program/src/processor.rs | 25 ++- token/program/tests/close_account.rs | 202 ++++++++++++++++++ 6 files changed, 454 insertions(+), 52 deletions(-) create mode 100644 token/program-2022-test/tests/close_account.rs create mode 100644 token/program/tests/close_account.rs diff --git a/token/program-2022-test/tests/close_account.rs b/token/program-2022-test/tests/close_account.rs new file mode 100644 index 00000000000..1e19d50ee92 --- /dev/null +++ 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 index 8fae81c07f0..e948a9991f3 100644 --- a/token/program-2022/src/instruction.rs +++ b/token/program-2022/src/instruction.rs @@ -1787,7 +1787,7 @@ pub(crate) fn encode_instruction, 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::(), 0..255) - ) { - let _no_panic = TokenInstruction::unpack(&data); - } - } } diff --git a/token/program-2022/src/processor.rs b/token/program-2022/src/processor.rs index fd9f3d5c128..6e3907dcc73 100644 --- 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::::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 index d31c1189fd9..c69c7bbe2eb 100644 --- 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::(), 0..255) - ) { - let _no_panic = TokenInstruction::unpack(&data); - } - } } diff --git a/token/program/src/processor.rs b/token/program/src/processor.rs index e6177092448..de7152ad865 100644 --- a/token/program/src/processor.rs +++ b/token/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::*; diff --git a/token/program/tests/close_account.rs b/token/program/tests/close_account.rs new file mode 100644 index 00000000000..195a7aa9d0f --- /dev/null +++ b/token/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()); +}