From 2da02992b79eec57095e4aa4010495310be86367 Mon Sep 17 00:00:00 2001
From: Denis K <100715772+dk-neon@users.noreply.github.com>
Date: Sat, 28 Jan 2023 14:03:31 +0500
Subject: [PATCH] Fixed bug with mishmash of traces from different programs in
 single log (#29605)

* Fixed bug with mishmash of traces from different programs in single log

* Transformed instructions trace log into the stack of (sub)programs' trace logs
---
 program-runtime/src/invoke_context.rs | 11 ++++++++---
 programs/sbf/tests/programs.rs        | 17 +++++++++++++++--
 rbpf-cli/src/main.rs                  |  8 +++++++-
 3 files changed, 30 insertions(+), 6 deletions(-)

diff --git a/program-runtime/src/invoke_context.rs b/program-runtime/src/invoke_context.rs
index 9dcadf57b1..11e9c6b093 100644
--- a/program-runtime/src/invoke_context.rs
+++ b/program-runtime/src/invoke_context.rs
@@ -61,7 +61,10 @@ impl std::fmt::Debug for BuiltinProgram {
 
 impl<'a> ContextObject for InvokeContext<'a> {
     fn trace(&mut self, state: [u64; 12]) {
-        self.trace_log.push(state);
+        self.trace_log_stack
+            .last_mut()
+            .expect("Inconsistent trace log stack")
+            .push(state);
     }
 
     fn consume(&mut self, amount: u64) {
@@ -104,7 +107,7 @@ pub struct InvokeContext<'a> {
     pre_accounts: Vec<PreAccount>,
     builtin_programs: &'a [BuiltinProgram],
     pub sysvar_cache: Cow<'a, SysvarCache>,
-    pub trace_log: Vec<[u64; 12]>,
+    pub trace_log_stack: Vec<Vec<[u64; 12]>>,
     log_collector: Option<Rc<RefCell<LogCollector>>>,
     compute_budget: ComputeBudget,
     current_compute_budget: ComputeBudget,
@@ -139,7 +142,7 @@ impl<'a> InvokeContext<'a> {
             pre_accounts: Vec::new(),
             builtin_programs,
             sysvar_cache,
-            trace_log: Vec::new(),
+            trace_log_stack: vec![Vec::new()],
             log_collector,
             current_compute_budget: compute_budget,
             compute_budget,
@@ -271,12 +274,14 @@ impl<'a> InvokeContext<'a> {
             }
         }
 
+        self.trace_log_stack.push(Vec::new());
         self.syscall_context.push(None);
         self.transaction_context.push()
     }
 
     /// Pop a stack frame from the invocation stack
     pub fn pop(&mut self) -> Result<(), InstructionError> {
+        self.trace_log_stack.pop();
         self.syscall_context.pop();
         self.transaction_context.pop()
     }
diff --git a/programs/sbf/tests/programs.rs b/programs/sbf/tests/programs.rs
index 1433f99bfc..223a211ecd 100644
--- a/programs/sbf/tests/programs.rs
+++ b/programs/sbf/tests/programs.rs
@@ -148,10 +148,23 @@ fn run_program(name: &str) -> u64 {
                 }
                 instruction_count = compute_units_consumed;
                 if i == 0 {
-                    trace_log = Some(vm.env.context_object_pointer.trace_log.clone());
+                    trace_log = Some(
+                        vm.env
+                            .context_object_pointer
+                            .trace_log_stack
+                            .last()
+                            .expect("Inconsistent trace log stack")
+                            .clone(),
+                    );
                 } else {
                     let interpreter = trace_log.as_ref().unwrap().as_slice();
-                    let mut jit = vm.env.context_object_pointer.trace_log.as_slice();
+                    let mut jit = vm
+                        .env
+                        .context_object_pointer
+                        .trace_log_stack
+                        .last()
+                        .expect("Inconsistent trace log stack")
+                        .as_slice();
                     if jit.len() > interpreter.len() {
                         jit = &jit[0..interpreter.len()];
                     }
diff --git a/rbpf-cli/src/main.rs b/rbpf-cli/src/main.rs
index d32f799a96..0ffb519a02 100644
--- a/rbpf-cli/src/main.rs
+++ b/rbpf-cli/src/main.rs
@@ -306,7 +306,13 @@ before execting it in the virtual machine.",
     let (instruction_count, result) = vm.execute_program(matches.value_of("use").unwrap() != "jit");
     let duration = Instant::now() - start_time;
     if matches.occurrences_of("trace") > 0 {
-        let trace_log = vm.env.context_object_pointer.trace_log.as_slice();
+        let trace_log = vm
+            .env
+            .context_object_pointer
+            .trace_log_stack
+            .last()
+            .expect("Inconsistent trace log stack")
+            .as_slice();
         if matches.value_of("trace").unwrap() == "stdout" {
             analysis
                 .analyze()