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
This commit is contained in:
parent
d026da4a1b
commit
0ce08274f9
|
@ -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,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<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,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
@ -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());
|
||||
}
|
||||
}
|
Loading…
Reference in New Issue