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:
Jon Cinque 2021-01-29 14:23:59 +01:00 committed by GitHub
parent d026da4a1b
commit 0ce08274f9
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 248 additions and 45 deletions

View File

@ -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,
}
}
}

118
program-test/tests/fuzz.rs Normal file
View File

@ -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());
}
}