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::{
|
solana_sdk::{
|
||||||
account::Account,
|
account::Account,
|
||||||
|
genesis_config::GenesisConfig,
|
||||||
keyed_account::KeyedAccount,
|
keyed_account::KeyedAccount,
|
||||||
process_instruction::{
|
process_instruction::{
|
||||||
stable_log, BpfComputeBudget, InvokeContext, ProcessInstructionWithContext,
|
stable_log, BpfComputeBudget, InvokeContext, ProcessInstructionWithContext,
|
||||||
|
@ -34,9 +35,13 @@ use {
|
||||||
mem::transmute,
|
mem::transmute,
|
||||||
path::{Path, PathBuf},
|
path::{Path, PathBuf},
|
||||||
rc::Rc,
|
rc::Rc,
|
||||||
sync::{Arc, RwLock},
|
sync::{
|
||||||
|
atomic::{AtomicBool, Ordering},
|
||||||
|
Arc, RwLock,
|
||||||
|
},
|
||||||
time::{Duration, Instant},
|
time::{Duration, Instant},
|
||||||
},
|
},
|
||||||
|
tokio::task::JoinHandle,
|
||||||
};
|
};
|
||||||
|
|
||||||
// Export types so test clients can limit their solana crate dependencies
|
// 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
|
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 {
|
pub struct ProgramTest {
|
||||||
accounts: Vec<(Pubkey, Account)>,
|
accounts: Vec<(Pubkey, Account)>,
|
||||||
builtins: Vec<Builtin>,
|
builtins: Vec<Builtin>,
|
||||||
|
@ -535,11 +579,7 @@ impl ProgramTest {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Start the test client
|
fn setup_bank(&self) -> (Arc<RwLock<BankForks>>, Keypair, Hash, GenesisConfig) {
|
||||||
///
|
|
||||||
/// 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) {
|
|
||||||
{
|
{
|
||||||
use std::sync::Once;
|
use std::sync::Once;
|
||||||
static ONCE: Once = Once::new();
|
static ONCE: Once = Once::new();
|
||||||
|
@ -564,7 +604,6 @@ impl ProgramTest {
|
||||||
let payer = gci.mint_keypair;
|
let payer = gci.mint_keypair;
|
||||||
debug!("Payer address: {}", payer.pubkey());
|
debug!("Payer address: {}", payer.pubkey());
|
||||||
debug!("Genesis config: {}", genesis_config);
|
debug!("Genesis config: {}", genesis_config);
|
||||||
let target_tick_duration = genesis_config.poh_config.target_tick_duration;
|
|
||||||
|
|
||||||
let mut bank = Bank::new(&genesis_config);
|
let mut bank = Bank::new(&genesis_config);
|
||||||
|
|
||||||
|
@ -581,7 +620,7 @@ impl ProgramTest {
|
||||||
}
|
}
|
||||||
|
|
||||||
// User-supplied additional builtins
|
// User-supplied additional builtins
|
||||||
for builtin in self.builtins {
|
for builtin in self.builtins.iter() {
|
||||||
bank.add_builtin(
|
bank.add_builtin(
|
||||||
&builtin.name,
|
&builtin.name,
|
||||||
builtin.id,
|
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() {
|
if bank.get_account(&address).is_some() {
|
||||||
info!("Overriding account at {}", address);
|
info!("Overriding account at {}", address);
|
||||||
}
|
}
|
||||||
|
@ -602,43 +641,15 @@ impl ProgramTest {
|
||||||
..BpfComputeBudget::default()
|
..BpfComputeBudget::default()
|
||||||
}));
|
}));
|
||||||
}
|
}
|
||||||
|
let bank = setup_fee_calculator(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();
|
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)));
|
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 transport = start_local_server(&bank_forks).await;
|
||||||
let banks_client = start_client(transport)
|
let banks_client = start_client(transport)
|
||||||
.await
|
.await
|
||||||
|
@ -654,12 +665,32 @@ impl ProgramTest {
|
||||||
.unwrap()
|
.unwrap()
|
||||||
.working_bank()
|
.working_bank()
|
||||||
.register_tick(&Hash::new_unique());
|
.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)
|
(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]
|
#[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