Skip to content

Commit

Permalink
✨ Possibility to implement custom tx error handler
Browse files Browse the repository at this point in the history
  • Loading branch information
Ikrk authored and lukacan committed Mar 19, 2024
1 parent efb920a commit dfa6356
Show file tree
Hide file tree
Showing 5 changed files with 122 additions and 58 deletions.
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ incremented upon a breaking change and the patch version will be incremented for

## [Unreleased]
### Added
- feat/possibility to implement custom transaction error handling ([#145](https://github.com/Ackee-Blockchain/trident/pull/145))
- feat/support of automatically obtaining fully qualified paths of Data Accounts Custom types for `accounts_snapshots.rs` ([#141](https://github.com/Ackee-Blockchain/trident/pull/141))
- feat/allow direct accounts manipulation and storage ([#142](https://github.com/Ackee-Blockchain/trident/pull/142))
- feat/support of non-corresponding instruction and context names ([#130](https://github.com/Ackee-Blockchain/trident/pull/130))
Expand Down
31 changes: 20 additions & 11 deletions crates/client/derive/fuzz_test_executor/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -37,20 +37,29 @@ pub fn fuzz_test_executor(input: TokenStream) -> TokenStream {
let sig: Vec<&Keypair> = signers.iter().collect();
transaction.sign(&sig, client.get_last_blockhash());

client.process_transaction(transaction)
.map_err(|e| e.with_origin(Origin::Instruction(self.to_context_string())))?;
let tx_result = client.process_transaction(transaction)
.map_err(|e| e.with_origin(Origin::Instruction(self.to_context_string())));

snaphot.capture_after(client).unwrap();
let (acc_before, acc_after) = snaphot.get_snapshot()
.map_err(|e| e.with_origin(Origin::Instruction(self.to_context_string())))
.expect("Snapshot deserialization expect"); // we want to panic if we cannot unwrap to cause a crash
match tx_result {
Ok(_) => {
snaphot.capture_after(client).unwrap();
let (acc_before, acc_after) = snaphot.get_snapshot()
.map_err(|e| e.with_origin(Origin::Instruction(self.to_context_string())))
.expect("Snapshot deserialization expect"); // we want to panic if we cannot unwrap to cause a crash

if let Err(e) = ix.check(acc_before, acc_after, data).map_err(|e| e.with_origin(Origin::Instruction(self.to_context_string()))) {
eprintln!(
"CRASH DETECTED! Custom check after the {} instruction did not pass!",
self.to_context_string());
panic!("{}", e)
if let Err(e) = ix.check(acc_before, acc_after, data).map_err(|e| e.with_origin(Origin::Instruction(self.to_context_string()))) {
eprintln!(
"CRASH DETECTED! Custom check after the {} instruction did not pass!",
self.to_context_string());
panic!("{}", e)
}
},
Err(e) => {
let mut raw_accounts = snaphot.get_raw_pre_ix_accounts();
ix.tx_error_handler(e, data, &mut raw_accounts)?
}
}

}
}
});
Expand Down
41 changes: 39 additions & 2 deletions crates/client/src/fuzzer/data_builder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -80,7 +80,6 @@ where
#[cfg(fuzzing_debug)]
eprintln!("Currently processing: {}", fuzz_ix);

// fuzz_ix.run_fuzzer(program_id, &self.accounts, client)?;
if fuzz_ix
.run_fuzzer(program_id, &self.accounts, client)
.is_err()
Expand Down Expand Up @@ -125,22 +124,40 @@ pub trait FuzzDataBuilder<T: for<'a> Arbitrary<'a>> {
}
}

/// A trait providing methods to prepare data and accounts for the fuzzed instructions and allowing
/// users to implement custom invariants checks and transactions error handling.
pub trait IxOps<'info> {
/// The data to be passed as instruction data parameter
type IxData;
/// The accounts to be passed as instruction accounts
type IxAccounts;
/// The structure to which the instruction accounts will be deserialized
type IxSnapshot;

/// Provides instruction data for the fuzzed instruction.
/// It is assumed that the instruction data will be based on the fuzzer input stored in the `self.data` variable.
/// However it is on the developer to decide and it can be also for example a hardcoded constant.
/// You should only avoid any non-deterministic random values to preserve reproducibility of the tests.
fn get_data(
&self,
client: &mut impl FuzzClient,
fuzz_accounts: &mut Self::IxAccounts,
) -> Result<Self::IxData, FuzzingError>;

/// Provides accounts required for the fuzzed instruction. The method returns a tuple of signers and account metas.
fn get_accounts(
&self,
client: &mut impl FuzzClient,
fuzz_accounts: &mut Self::IxAccounts,
) -> Result<(Vec<Keypair>, Vec<AccountMeta>), FuzzingError>;

/// A method to implement custom invariants checks for a given instruction. This method is called after each
/// successfully executed instruction and by default does nothing. You can override this behavior by providing
/// your own implementation. You can access the snapshots of account states before and after the transaction for comparison.
///
/// If you want to detect a crash, you have to return a `FuzzingError` (or alternativelly panic).
///
/// If you want to perform checks also on a failed instruction execution, you can do so using the [`tx_error_handler`](trident_client::fuzzer::data_builder::IxOps::tx_error_handler) method.
#[allow(unused_variables)]
fn check(
&self,
Expand All @@ -151,11 +168,31 @@ pub trait IxOps<'info> {
Ok(())
}

/// A method to implement custom error handler for failed transactions.
///
/// The fuzzer might generate a sequence of one or more instructions that are executed sequentially.
/// By default, if the execution of one of the instructions fails, the remaining instructions are skipped
/// and are not executed. This can be overriden by implementing this method and returning `Ok(())`
/// instead of propagating the error.
///
/// You can also check the kind of the transaction error by inspecting the `e` parameter.
/// If you would like to detect a crash on a specific error, call `panic!()`.
///
/// If your accounts are malformed and the fuzzed program is unable to deserialize it, the transaction
/// execution will fail. In that case also the deserialization of accounts snapshot before executing
/// the instruction would fail. You are provided with the raw account infos snapshots and you are free
/// to deserialize the accounts by yourself and therefore also handling potential errors. To deserialize
/// the `pre_ix_acc_infos` raw accounts to a snapshot structure, you can call:
///
/// ```rust,ignore
/// self.deserialize_option(pre_ix_acc_infos)
/// ```
#[allow(unused_variables)]
fn tx_error_handler(
&self,
e: FuzzClientErrorWithOrigin,
ix_data: Self::IxData,
pre_ix_acc_infos: &'info [Option<AccountInfo<'info>>]
pre_ix_acc_infos: &'info mut [Option<AccountInfo<'info>>],
) -> Result<(), FuzzClientErrorWithOrigin> {
Err(e)
}
Expand Down
5 changes: 5 additions & 0 deletions crates/client/src/fuzzer/snapshot.rs
Original file line number Diff line number Diff line change
Expand Up @@ -93,6 +93,11 @@ where
}
}

pub fn get_raw_pre_ix_accounts(&'info mut self) -> Vec<Option<AccountInfo<'info>>> {
Self::set_missing_accounts_to_default(&mut self.before);
Self::calculate_account_info(&mut self.before, self.metas)
}

pub fn get_snapshot(&'info mut self) -> Result<(T::Ix, T::Ix), FuzzingErrorWithOrigin> {
// When user passes an account that is not initialized, the runtime will provide
// a default empty account to the program. If the uninitialized account is of type
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ impl FuzzTestExecutor<FuzzAccounts> for FuzzInstruction {
program_id: Pubkey,
accounts: &RefCell<FuzzAccounts>,
client: &mut impl FuzzClient,
) -> core::result::Result<(), Box<dyn std::error::Error + 'static>> {
) -> core::result::Result<(), FuzzClientErrorWithOrigin> {
match self {
FuzzInstruction::InitVesting(ix) => {
let (mut signers, metas) = ix
Expand Down Expand Up @@ -38,38 +38,44 @@ impl FuzzTestExecutor<FuzzAccounts> for FuzzInstruction {
signers.push(client.payer().clone());
let sig: Vec<&Keypair> = signers.iter().collect();
transaction.sign(&sig, client.get_last_blockhash());
let tx_res = client
let tx_result = client
.process_transaction(transaction)
.map_err(|e| {
e.with_origin(Origin::Instruction(self.to_context_string()))
});
if tx_res.is_ok() {
snaphot.capture_after(client).unwrap();
let (acc_before, acc_after) = snaphot
.get_snapshot()
.map_err(|e| {
e.with_origin(Origin::Instruction(self.to_context_string()))
})
.expect("Snapshot deserialization expect");
if let Err(e)
= ix
.check(acc_before, acc_after, data)
match tx_result {
Ok(_) => {
snaphot.capture_after(client).unwrap();
let (acc_before, acc_after) = snaphot
.get_snapshot()
.map_err(|e| {
e.with_origin(Origin::Instruction(self.to_context_string()))
})
{
.expect("Snapshot deserialization expect");
if let Err(e)
= ix
.check(acc_before, acc_after, data)
.map_err(|e| {
e.with_origin(Origin::Instruction(self.to_context_string()))
})
{
::std::io::_eprint(
format_args!(
"CRASH DETECTED! Custom check after the {0} instruction did not pass!\n",
self.to_context_string(),
),
);
};
{
::core::panicking::panic_display(&e);
{
::std::io::_eprint(
format_args!(
"CRASH DETECTED! Custom check after the {0} instruction did not pass!\n",
self.to_context_string(),
),
);
};
{
::core::panicking::panic_display(&e);
}
}
}
Err(e) => {
let mut raw_accounts = snaphot.get_raw_pre_ix_accounts();
ix.tx_error_handler(e, data, &mut raw_accounts)?
}
}
}
FuzzInstruction::WithdrawUnlocked(ix) => {
Expand Down Expand Up @@ -99,38 +105,44 @@ impl FuzzTestExecutor<FuzzAccounts> for FuzzInstruction {
signers.push(client.payer().clone());
let sig: Vec<&Keypair> = signers.iter().collect();
transaction.sign(&sig, client.get_last_blockhash());
let tx_res = client
let tx_result = client
.process_transaction(transaction)
.map_err(|e| {
e.with_origin(Origin::Instruction(self.to_context_string()))
});
if tx_res.is_ok() {
snaphot.capture_after(client).unwrap();
let (acc_before, acc_after) = snaphot
.get_snapshot()
.map_err(|e| {
e.with_origin(Origin::Instruction(self.to_context_string()))
})
.expect("Snapshot deserialization expect");
if let Err(e)
= ix
.check(acc_before, acc_after, data)
match tx_result {
Ok(_) => {
snaphot.capture_after(client).unwrap();
let (acc_before, acc_after) = snaphot
.get_snapshot()
.map_err(|e| {
e.with_origin(Origin::Instruction(self.to_context_string()))
})
{
.expect("Snapshot deserialization expect");
if let Err(e)
= ix
.check(acc_before, acc_after, data)
.map_err(|e| {
e.with_origin(Origin::Instruction(self.to_context_string()))
})
{
::std::io::_eprint(
format_args!(
"CRASH DETECTED! Custom check after the {0} instruction did not pass!\n",
self.to_context_string(),
),
);
};
{
::core::panicking::panic_display(&e);
{
::std::io::_eprint(
format_args!(
"CRASH DETECTED! Custom check after the {0} instruction did not pass!\n",
self.to_context_string(),
),
);
};
{
::core::panicking::panic_display(&e);
}
}
}
Err(e) => {
let mut raw_accounts = snaphot.get_raw_pre_ix_accounts();
ix.tx_error_handler(e, data, &mut raw_accounts)?
}
}
}
}
Expand Down

0 comments on commit dfa6356

Please sign in to comment.