Register tracing export

Cargo.toml

@@ -29,6 +29,7 @@ bs58 = "0.5.1"
 chrono = "0.4.42"
 clap = "4.5.51"
 criterion = "0.7.0"
+hex = "0.4.3"
 ed25519-dalek = "=1.0.1"
 libsecp256k1 = "0.6.0"
 mollusk-svm = { path = "harness", version = "0.7.2" }
@@ -53,6 +54,7 @@ serde = "1.0.203"
 serde_json = "1.0.117"
 serde_yaml = "0.9.34"
 serial_test = "2.0"
+sha2 = "0.10.9"
 solana-account = "3.2.0"
 solana-account-info = "3.0"
 solana-bpf-loader-program = "3.1.0"

harness/Cargo.toml

@@ -27,6 +27,11 @@ fuzz-fd = [
 ]
 invocation-inspect-callback = []
 precompiles = ["dep:agave-precompiles"]
+register-tracing = [
+    "invocation-inspect-callback",
+    "dep:hex",
+    "dep:sha2"
+]
 serde = [
     "dep:serde",
     "mollusk-svm-result/serde",
@@ -37,13 +42,15 @@ agave-feature-set = { workspace = true, features = ["agave-unstable-api"] }
 agave-precompiles = { workspace = true, features = ["agave-unstable-api"], optional = true }
 agave-syscalls = { workspace = true }
 bincode = { workspace = true }
+hex = { workspace = true, optional = true }
 serde = { workspace = true, features = ["derive"], optional = true }
 mollusk-svm-error = { workspace = true }
 mollusk-svm-fuzz-fixture = { workspace = true, optional = true }
 mollusk-svm-fuzz-fixture-firedancer = { workspace = true, optional = true }
 mollusk-svm-fuzz-fs = { workspace = true, optional = true }
 mollusk-svm-keys = { workspace = true }
 mollusk-svm-result = { workspace = true }
+sha2 = { workspace = true, optional = true }
 solana-account = { workspace = true }
 solana-bpf-loader-program = { workspace = true, features = ["agave-unstable-api"] }
 solana-clock = { workspace = true }

harness/src/file.rs

@@ -27,7 +27,7 @@ use {
     },
 };
 
-fn default_shared_object_dirs() -> Vec<PathBuf> {
+pub(crate) fn default_shared_object_dirs() -> Vec<PathBuf> {
     let mut search_path = vec![PathBuf::from("tests/fixtures")];
 
     if let Ok(bpf_out_dir) = std::env::var("BPF_OUT_DIR") {

harness/src/lib.rs

@@ -446,8 +446,12 @@ pub mod file;
 #[cfg(any(feature = "fuzz", feature = "fuzz-fd"))]
 pub mod fuzz;
 pub mod program;
+#[cfg(feature = "register-tracing")]
+pub mod register_tracing;
 pub mod sysvar;
 
+#[cfg(feature = "register-tracing")]
+use crate::register_tracing::DefaultRegisterTracingCallback;
 // Re-export result module from mollusk-svm-result crate
 pub use mollusk_svm_result as result;
 #[cfg(any(feature = "fuzz", feature = "fuzz-fd"))]
@@ -502,6 +506,9 @@ pub struct Mollusk {
     #[cfg(feature = "invocation-inspect-callback")]
     pub invocation_inspect_callback: Box<dyn InvocationInspectCallback>,
 
+    #[cfg(feature = "register-tracing")]
+    pub enable_register_tracing: bool,
+
     /// 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
@@ -557,7 +564,20 @@ impl Default for Mollusk {
         };
         #[cfg(not(feature = "fuzz"))]
         let feature_set = FeatureSet::all_enabled();
-        let program_cache = ProgramCache::new(&feature_set, &compute_budget);
+
+        let program_cache = ProgramCache::new(
+            &feature_set,
+            &compute_budget,
+            #[cfg(feature = "register-tracing")]
+            {
+                true
+            },
+            #[cfg(not(feature = "register-tracing"))]
+            {
+                false
+            },
+        );
+
         Self {
             config: Config::default(),
             compute_budget,
@@ -567,9 +587,20 @@ impl Default for Mollusk {
             program_cache,
             sysvars: Sysvars::default(),
 
-            #[cfg(feature = "invocation-inspect-callback")]
+            // Register tracing feature requires `invocation-inspect-callback`.
+            // Use tracing callback when both are active.
+            #[cfg(all(feature = "invocation-inspect-callback", feature = "register-tracing"))]
+            invocation_inspect_callback: Box::new(DefaultRegisterTracingCallback::default()),
+            // Use empty callback when only `invocation-inspect-callback` is active.
+            #[cfg(all(
+                feature = "invocation-inspect-callback",
+                not(feature = "register-tracing"),
+            ))]
             invocation_inspect_callback: Box::new(EmptyInvocationInspectCallback {}),
 
+            #[cfg(feature = "register-tracing")]
+            enable_register_tracing: true,
+
             #[cfg(feature = "fuzz-fd")]
             slot: 0,
         }
@@ -723,21 +754,26 @@ impl Mollusk {
             let runtime_features = self.feature_set.runtime_features();
             let sysvar_cache = self.sysvars.setup_sysvar_cache(accounts);
 
-            let program_runtime_environments = ProgramRuntimeEnvironments {
-                program_runtime_v1: Arc::new(
-                    create_program_runtime_environment_v1(
-                        &runtime_features,
+            let _enable_register_tracing = false;
+            #[cfg(feature = "register-tracing")]
+            let _enable_register_tracing = self.enable_register_tracing;
+
+            let program_runtime_environments: ProgramRuntimeEnvironments =
+                ProgramRuntimeEnvironments {
+                    program_runtime_v1: Arc::new(
+                        create_program_runtime_environment_v1(
+                            &runtime_features,
+                            &execution_budget,
+                            /* reject_deployment_of_broken_elfs */ false,
+                            /* debugging_features */ _enable_register_tracing,
+                        )
+                        .unwrap(),
+                    ),
+                    program_runtime_v2: Arc::new(create_program_runtime_environment_v2(
                         &execution_budget,
-                        /* reject_deployment_of_broken_elfs */ false,
-                        /* debugging_features */ false,
-                    )
-                    .unwrap(),
-                ),
-                program_runtime_v2: Arc::new(create_program_runtime_environment_v2(
-                    &execution_budget,
-                    /* debugging_features */ false,
-                )),
-            };
+                        /* debugging_features */ _enable_register_tracing,
+                    )),
+                };
 
             let mut invoke_context = InvokeContext::new(
                 &mut transaction_context,

harness/src/program.rs

@@ -74,15 +74,19 @@ pub struct ProgramCache {
 }
 
 impl ProgramCache {
-    pub fn new(feature_set: &FeatureSet, compute_budget: &ComputeBudget) -> Self {
+    pub fn new(
+        feature_set: &FeatureSet,
+        compute_budget: &ComputeBudget,
+        enable_register_tracing: bool,
+    ) -> Self {
         let me = Self {
             cache: Rc::new(RefCell::new(ProgramCacheForTxBatch::default())),
             entries_cache: Rc::new(RefCell::new(HashMap::new())),
             program_runtime_environment: create_program_runtime_environment_v1(
                 &feature_set.runtime_features(),
                 &compute_budget.to_budget(),
                 /* reject_deployment_of_broken_elfs */ false,
-                /* debugging_features */ false,
+                /* debugging_features */ enable_register_tracing,
             )
             .unwrap(),
         };

harness/src/register_tracing.rs

@@ -0,0 +1,142 @@
+use {
+    crate::{
+        file::{default_shared_object_dirs, read_file},
+        InvocationInspectCallback,
+    },
+    sha2::{Digest, Sha256},
+    solana_program_runtime::invoke_context::{Executable, InvokeContext, RegisterTrace},
+    solana_pubkey::Pubkey,
+    solana_transaction_context::{InstructionAccount, InstructionContext},
+    std::{fs::File, io::Write, path::PathBuf},
+};
+
+const DEFAULT_PATH: &str = "target/sbf/trace";
+
+pub struct DefaultRegisterTracingCallback {
+    pub sbf_trace_dir: String,
+}
+
+impl Default for DefaultRegisterTracingCallback {
+    fn default() -> Self {
+        Self {
+            // User can override default path with `SBF_TRACE_DIR` environment variable.
+            sbf_trace_dir: std::env::var("SBF_TRACE_DIR").unwrap_or(DEFAULT_PATH.to_string()),
+        }
+    }
+}
+
+impl DefaultRegisterTracingCallback {
+    pub fn handler(
+        &self,
+        instruction_context: InstructionContext,
+        executable: &Executable,
+        register_trace: RegisterTrace,
+    ) -> Result<(), Box<dyn std::error::Error>> {
+        if register_trace.is_empty() {
+            // Can't do much with an empty trace.
+            return Ok(());
+        }
+
+        let current_dir = std::env::current_dir()?;
+        let sbf_trace_dir = current_dir.join(&self.sbf_trace_dir);
+        std::fs::create_dir_all(&sbf_trace_dir)?;
+
+        let trace_digest = compute_hash(as_bytes(register_trace));
+        let base_fname = sbf_trace_dir.join(&trace_digest[..16]);
+        let mut regs_file = File::create(base_fname.with_extension("regs"))?;
+        let mut insns_file = File::create(base_fname.with_extension("insns"))?;
+        let mut so_hash_file = File::create(base_fname.with_extension("exec.sha256"))?;
+
+        // Get program_id.
+        let program_id = instruction_context.get_program_key()?;
+
+        // Persist the preload hash of the executable.
+        let _ = so_hash_file.write(
+            find_executable_pre_load_hash(executable)
+                .ok_or(format!(
+                    "Can't find shared object for executable with program_id: {program_id}"
+                ))?
+                .as_bytes(),
+        );
+
+        // Get the relocated executable.
+        let (_, program) = executable.get_text_bytes();
+        for regs in register_trace.iter() {
+            // The program counter is stored in r11.
+            let pc = regs[11];
+            // From the executable fetch the instruction this program counter points to.
+            let insn =
+                solana_program_runtime::solana_sbpf::ebpf::get_insn_unchecked(program, pc as usize)
+                    .to_array();
+
+            // Persist them in files.
+            let _ = regs_file.write(as_bytes(regs.as_slice()))?;
+            let _ = insns_file.write(insn.as_slice())?;
+        }
+
+        Ok(())
+    }
+}
+
+impl InvocationInspectCallback for DefaultRegisterTracingCallback {
+    fn before_invocation(&self, _: &Pubkey, _: &[u8], _: &[InstructionAccount], _: &InvokeContext) {
+    }
+
+    fn after_invocation(&self, invoke_context: &InvokeContext) {
+        invoke_context.iterate_vm_traces(
+            &|instruction_context: InstructionContext,
+              executable: &Executable,
+              register_trace: RegisterTrace| {
+                if let Err(e) = self.handler(instruction_context, executable, register_trace) {
+                    eprintln!("Error collecting the register tracing: {}", e);
+                }
+            },
+        );
+    }
+}
+
+pub(crate) fn as_bytes<T>(slice: &[T]) -> &[u8] {
+    unsafe { std::slice::from_raw_parts(slice.as_ptr() as *const u8, std::mem::size_of_val(slice)) }
+}
+
+fn find_so_files(dirs: &[PathBuf]) -> Vec<PathBuf> {
+    let mut so_files = Vec::new();
+
+    for dir in dirs {
+        if dir.is_dir() {
+            if let Ok(entries) = std::fs::read_dir(dir) {
+                for entry in entries.flatten() {
+                    let path = entry.path();
+                    if path.is_file() && path.extension().is_some_and(|ext| ext == "so") {
+                        so_files.push(path);
+                    }
+                }
+            }
+        }
+    }
+
+    so_files
+}
+
+fn find_executable_pre_load_hash(executable: &Executable) -> Option<String> {
+    find_so_files(&default_shared_object_dirs())
+        .iter()
+        .filter_map(|file| {
+            let so = read_file(file);
+            // Reconstruct a loaded Exectuable just to compare its relocated
+            // text bytes with the passed executable ones.
+            // If there's a match return the preload hash of the corresponding shared
+            // object.
+            Executable::load(&so, executable.get_loader().clone())
+                .ok()
+                .map(|e| Some((so, e)))
+                .unwrap_or(None)
+        })
+        .filter(|(_, e)| executable.get_text_bytes().1 == e.get_text_bytes().1)
+        .map(|(so, _)| compute_hash(&so))
+        .next_back()
+}
+
+fn compute_hash(slice: &[u8]) -> String {
+    hex::encode(Sha256::digest(slice).as_slice())
+}