From 0ce08274f95a93ea10fd0e5ceaf626e63ed392a9 Mon Sep 17 00:00:00 2001 From: Jon Cinque Date: Fri, 29 Jan 2021 14:23:59 +0100 Subject: [PATCH] program-test: Expose bank task to fix fuzzing (#14908) * program-test: Expose bank task to fix fuzzing * Run cargo fmt and clippy * Remove unnecessary print in test * Review feedback * Transition to AtomicBool --- program-test/src/lib.rs | 175 +++++++++++++++++++++++++++---------- program-test/tests/fuzz.rs | 118 +++++++++++++++++++++++++ 2 files changed, 248 insertions(+), 45 deletions(-) create mode 100644 program-test/tests/fuzz.rs diff --git a/program-test/src/lib.rs b/program-test/src/lib.rs index 9bb49c518..6676ab14d 100644 --- a/program-test/src/lib.rs +++ b/program-test/src/lib.rs @@ -19,6 +19,7 @@ use { }, solana_sdk::{ account::Account, + genesis_config::GenesisConfig, keyed_account::KeyedAccount, process_instruction::{ stable_log, BpfComputeBudget, InvokeContext, ProcessInstructionWithContext, @@ -34,9 +35,13 @@ use { mem::transmute, path::{Path, PathBuf}, rc::Rc, - sync::{Arc, RwLock}, + sync::{ + atomic::{AtomicBool, Ordering}, + Arc, RwLock, + }, time::{Duration, Instant}, }, + tokio::task::JoinHandle, }; // Export types so test clients can limit their solana crate dependencies @@ -351,6 +356,45 @@ pub fn read_file>(path: P) -> Vec { file_data } +fn setup_fee_calculator(bank: Bank) -> Bank { + // Realistic fee_calculator part 1: Fake a single signature by calling + // `bank.commit_transactions()` so that the fee calculator in the child bank will be + // initialized with a non-zero fee. + assert_eq!(bank.signature_count(), 0); + bank.commit_transactions( + &[], + None, + &mut [], + &[], + 0, + 1, + &mut ExecuteTimings::default(), + ); + assert_eq!(bank.signature_count(), 1); + + // Advance beyond slot 0 for a slightly more realistic test environment + let bank = Arc::new(bank); + let bank = Bank::new_from_parent(&bank, bank.collector_id(), bank.slot() + 1); + debug!("Bank slot: {}", bank.slot()); + + // Realistic fee_calculator part 2: Tick until a new blockhash is produced to pick up the + // non-zero fee calculator + let last_blockhash = bank.last_blockhash(); + while last_blockhash == bank.last_blockhash() { + bank.register_tick(&Hash::new_unique()); + } + let last_blockhash = bank.last_blockhash(); + // Make sure the new last_blockhash now requires a fee + assert_ne!( + bank.get_fee_calculator(&last_blockhash) + .expect("fee_calculator") + .lamports_per_signature, + 0 + ); + + bank +} + pub struct ProgramTest { accounts: Vec<(Pubkey, Account)>, builtins: Vec, @@ -535,11 +579,7 @@ impl ProgramTest { } } - /// Start the test client - /// - /// Returns a `BanksClient` interface into the test environment as well as a payer `Keypair` - /// with SOL for sending transactions - pub async fn start(self) -> (BanksClient, Keypair, Hash) { + fn setup_bank(&self) -> (Arc>, Keypair, Hash, GenesisConfig) { { use std::sync::Once; static ONCE: Once = Once::new(); @@ -564,7 +604,6 @@ impl ProgramTest { let payer = gci.mint_keypair; debug!("Payer address: {}", payer.pubkey()); debug!("Genesis config: {}", genesis_config); - let target_tick_duration = genesis_config.poh_config.target_tick_duration; let mut bank = Bank::new(&genesis_config); @@ -581,7 +620,7 @@ impl ProgramTest { } // User-supplied additional builtins - for builtin in self.builtins { + for builtin in self.builtins.iter() { bank.add_builtin( &builtin.name, builtin.id, @@ -589,7 +628,7 @@ impl ProgramTest { ); } - for (address, account) in self.accounts { + for (address, account) in self.accounts.iter() { if bank.get_account(&address).is_some() { info!("Overriding account at {}", address); } @@ -602,43 +641,15 @@ impl ProgramTest { ..BpfComputeBudget::default() })); } - - // Realistic fee_calculator part 1: Fake a single signature by calling - // `bank.commit_transactions()` so that the fee calculator in the child bank will be - // initialized with a non-zero fee. - assert_eq!(bank.signature_count(), 0); - bank.commit_transactions( - &[], - None, - &mut [], - &[], - 0, - 1, - &mut ExecuteTimings::default(), - ); - assert_eq!(bank.signature_count(), 1); - - // Advance beyond slot 0 for a slightly more realistic test environment - let bank = Arc::new(bank); - let bank = Bank::new_from_parent(&bank, bank.collector_id(), bank.slot() + 1); - debug!("Bank slot: {}", bank.slot()); - - // Realistic fee_calculator part 2: Tick until a new blockhash is produced to pick up the - // non-zero fee calculator + let bank = setup_fee_calculator(bank); let last_blockhash = bank.last_blockhash(); - while last_blockhash == bank.last_blockhash() { - bank.register_tick(&Hash::new_unique()); - } - let last_blockhash = bank.last_blockhash(); - // Make sure the new last_blockhash now requires a fee - assert_ne!( - bank.get_fee_calculator(&last_blockhash) - .expect("fee_calculator") - .lamports_per_signature, - 0 - ); - let bank_forks = Arc::new(RwLock::new(BankForks::new(bank))); + + (bank_forks, payer, last_blockhash, genesis_config) + } + + pub async fn start(self) -> (BanksClient, Keypair, Hash) { + let (bank_forks, payer, last_blockhash, genesis_config) = self.setup_bank(); let transport = start_local_server(&bank_forks).await; let banks_client = start_client(transport) .await @@ -654,12 +665,32 @@ impl ProgramTest { .unwrap() .working_bank() .register_tick(&Hash::new_unique()); - tokio::time::sleep(target_tick_duration).await; + tokio::time::sleep(genesis_config.poh_config.target_tick_duration).await; } }); (banks_client, payer, last_blockhash) } + + /// Start the test client + /// + /// Returns a `BanksClient` interface into the test environment as well as a payer `Keypair` + /// with SOL for sending transactions + pub async fn start_with_context(self) -> ProgramTestContext { + let (bank_forks, payer, last_blockhash, genesis_config) = self.setup_bank(); + let transport = start_local_server(&bank_forks).await; + let banks_client = start_client(transport) + .await + .unwrap_or_else(|err| panic!("Failed to start banks client: {}", err)); + + ProgramTestContext::new( + bank_forks, + banks_client, + payer, + last_blockhash, + genesis_config, + ) + } } #[async_trait] @@ -698,3 +729,57 @@ impl ProgramTestBanksClientExt for BanksClient { )) } } + +struct DroppableTask(Arc, JoinHandle); + +impl Drop for DroppableTask { + fn drop(&mut self) { + self.0.store(true, Ordering::Relaxed); + } +} + +pub struct ProgramTestContext { + pub banks_client: BanksClient, + pub payer: Keypair, + pub last_blockhash: Hash, + _bank_task: DroppableTask<()>, +} + +impl ProgramTestContext { + fn new( + bank_forks: Arc>, + banks_client: BanksClient, + payer: Keypair, + last_blockhash: Hash, + genesis_config: GenesisConfig, + ) -> Self { + // Run a simulated PohService to provide the client with new blockhashes. New blockhashes + // are required when sending multiple otherwise identical transactions in series from a + // test + let target_tick_duration = genesis_config.poh_config.target_tick_duration; + let exit = Arc::new(AtomicBool::new(false)); + let bank_task = DroppableTask( + exit.clone(), + tokio::spawn(async move { + loop { + if exit.load(Ordering::Relaxed) { + break; + } + bank_forks + .read() + .unwrap() + .working_bank() + .register_tick(&Hash::new_unique()); + tokio::time::sleep(target_tick_duration).await; + } + }), + ); + + Self { + banks_client, + payer, + last_blockhash, + _bank_task: bank_task, + } + } +} diff --git a/program-test/tests/fuzz.rs b/program-test/tests/fuzz.rs new file mode 100644 index 000000000..74c395b53 --- /dev/null +++ b/program-test/tests/fuzz.rs @@ -0,0 +1,118 @@ +use { + solana_banks_client::BanksClient, + solana_program::{ + account_info::AccountInfo, entrypoint::ProgramResult, hash::Hash, pubkey::Pubkey, + rent::Rent, + }, + solana_program_test::{processor, ProgramTest}, + solana_sdk::{ + signature::Keypair, signature::Signer, system_instruction, transaction::Transaction, + }, +}; + +// Dummy process instruction required to instantiate ProgramTest +#[allow(clippy::unnecessary_wraps)] +fn process_instruction( + _program_id: &Pubkey, + _accounts: &[AccountInfo], + _input: &[u8], +) -> ProgramResult { + Ok(()) +} + +#[test] +fn simulate_fuzz() { + let rt = tokio::runtime::Runtime::new().unwrap(); + let program_id = Pubkey::new_unique(); + // Initialize and start the test network + let program_test = ProgramTest::new( + "program-test-fuzz", + program_id, + processor!(process_instruction), + ); + + let (mut banks_client, payer, last_blockhash) = + rt.block_on(async { program_test.start().await }); + + // the honggfuzz `fuzz!` macro does not allow for async closures, + // so we have to use the runtime directly to run async functions + rt.block_on(async { + run_fuzz_instructions( + &[1, 2, 3, 4, 5], + &mut banks_client, + &payer, + last_blockhash, + &program_id, + ) + .await + }); +} + +#[test] +fn simulate_fuzz_with_context() { + let rt = tokio::runtime::Runtime::new().unwrap(); + let program_id = Pubkey::new_unique(); + // Initialize and start the test network + let program_test = ProgramTest::new( + "program-test-fuzz", + program_id, + processor!(process_instruction), + ); + + let mut test_state = rt.block_on(async { program_test.start_with_context().await }); + + // the honggfuzz `fuzz!` macro does not allow for async closures, + // so we have to use the runtime directly to run async functions + rt.block_on(async { + run_fuzz_instructions( + &[1, 2, 3, 4, 5], + &mut test_state.banks_client, + &test_state.payer, + test_state.last_blockhash, + &program_id, + ) + .await + }); +} + +async fn run_fuzz_instructions( + fuzz_instruction: &[u8], + banks_client: &mut BanksClient, + payer: &Keypair, + last_blockhash: Hash, + program_id: &Pubkey, +) { + let mut instructions = vec![]; + let mut signer_keypairs = vec![]; + for &i in fuzz_instruction { + let keypair = Keypair::new(); + let instruction = system_instruction::create_account( + &payer.pubkey(), + &keypair.pubkey(), + Rent::default().minimum_balance(i as usize), + i as u64, + program_id, + ); + instructions.push(instruction); + signer_keypairs.push(keypair); + } + // Process transaction on test network + let mut transaction = Transaction::new_with_payer(&instructions, Some(&payer.pubkey())); + let signers = [payer] + .iter() + .copied() + .chain(signer_keypairs.iter()) + .collect::>(); + transaction.partial_sign(&signers, last_blockhash); + + banks_client.process_transaction(transaction).await.unwrap(); + for keypair in signer_keypairs { + let account = banks_client + .get_account(keypair.pubkey()) + .await + .expect("account exists") + .unwrap(); + assert!(account.lamports > 0); + assert!(!account.data.is_empty()); + } +}