feat: add invoke context inspection callback (#143)

Co-authored-by: Joe C <[email protected]>

Author: ssmike

harness/Cargo.toml

@@ -27,6 +27,7 @@ serde = [
     "dep:serde",
     "mollusk-svm-result/serde",
 ]
+invocation-inspect-callback = []
 
 [dependencies]
 agave-feature-set = { workspace = true }

harness/src/lib.rs

@@ -452,6 +452,8 @@ pub mod sysvar;
 pub use mollusk_svm_result as result;
 #[cfg(any(feature = "fuzz", feature = "fuzz-fd"))]
 use mollusk_svm_result::Compare;
+#[cfg(feature = "invocation-inspect-callback")]
+use solana_transaction_context::InstructionAccount;
 use {
     crate::{
         account_store::AccountStore, compile_accounts::CompiledAccounts, epoch_stake::EpochStake,
@@ -489,6 +491,12 @@ pub struct Mollusk {
     pub program_cache: ProgramCache,
     pub sysvars: Sysvars,
 
+    /// The callback which can be used to inspect invoke_context
+    /// and extract low-level information such as bpf traces, transaction context,
+    /// detailed timings, etc.
+    #[cfg(feature = "invocation-inspect-callback")]
+    pub invocation_inspect_callback: Box<dyn InvocationInspectCallback>,
+
     /// This field stores the slot only to be able to convert to and from FD
     /// fixtures and a Mollusk instance, since FD fixtures have a
     /// "slot context". However, this field is functionally irrelevant for
@@ -498,6 +506,36 @@ pub struct Mollusk {
     pub slot: u64,
 }
 
+#[cfg(feature = "invocation-inspect-callback")]
+pub trait InvocationInspectCallback {
+    fn before_invocation(
+        &self,
+        program_id: &Pubkey,
+        instruction_data: &[u8],
+        instruction_accounts: &[InstructionAccount],
+        invoke_context: &InvokeContext,
+    );
+
+    fn after_invocation(&self, invoke_context: &InvokeContext);
+}
+
+#[cfg(feature = "invocation-inspect-callback")]
+pub struct EmptyInvocationInspectCallback;
+
+#[cfg(feature = "invocation-inspect-callback")]
+impl InvocationInspectCallback for EmptyInvocationInspectCallback {
+    fn before_invocation(
+        &self,
+        _: &Pubkey,
+        _: &[u8],
+        _: &[InstructionAccount],
+        _: &InvokeContext,
+    ) {
+    }
+
+    fn after_invocation(&self, _: &InvokeContext) {}
+}
+
 impl Default for Mollusk {
     fn default() -> Self {
         #[rustfmt::skip]
@@ -528,6 +566,10 @@ impl Default for Mollusk {
             logger: None,
             program_cache,
             sysvars: Sysvars::default(),
+
+            #[cfg(feature = "invocation-inspect-callback")]
+            invocation_inspect_callback: Box::new(EmptyInvocationInspectCallback {}),
+
             #[cfg(feature = "fuzz-fd")]
             slot: 0,
         }
@@ -678,7 +720,16 @@ impl Mollusk {
                 self.compute_budget.to_budget(),
                 self.compute_budget.to_cost(),
             );
-            if invoke_context.is_precompile(&instruction.program_id) {
+
+            #[cfg(feature = "invocation-inspect-callback")]
+            self.invocation_inspect_callback.before_invocation(
+                &instruction.program_id,
+                &instruction.data,
+                &instruction_accounts,
+                &invoke_context,
+            );
+
+            let result = if invoke_context.is_precompile(&instruction.program_id) {
                 invoke_context.process_precompile(
                     &instruction.program_id,
                     &instruction.data,
@@ -694,7 +745,13 @@ impl Mollusk {
                     &mut compute_units_consumed,
                     &mut timings,
                 )
-            }
+            };
+
+            #[cfg(feature = "invocation-inspect-callback")]
+            self.invocation_inspect_callback
+                .after_invocation(&invoke_context);
+
+            result
         };
 
         let return_data = transaction_context.get_return_data().1.to_vec();