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

@@ -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<P: AsRef<Path>>(path: P) -> Vec<u8> {
     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<Builtin>,
@@ -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<RwLock<BankForks>>, 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,15 +620,15 @@ impl ProgramTest {
         }
 
         // User-supplied additional builtins
-        for builtin in self.builtins {
+        for builtin in self.builtins.iter() {
             bank.add_builtin(
                 &builtin.name,
                 builtin.id,
                 builtin.process_instruction_with_context,
             );
         }
 
-        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<T>(Arc<AtomicBool>, JoinHandle<T>);
+
+impl<T> Drop for DroppableTask<T> {
+    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<RwLock<BankForks>>,
+        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,
+        }
+    }
+}

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::<Vec<&Keypair>>();
+    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());
+    }
+}