From aadf4b9b63e2e754c8e2b96ea1b852c9d622d9bd Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Alexander=20Mei=C3=9Fner?= <AlexanderMeissner@gmx.net>
Date: Mon, 10 Jan 2022 18:26:51 +0100
Subject: [PATCH] Moves InvokeContext::return_data to TransactionContext.
 (#22411)

---
 program-runtime/src/invoke_context.rs |  5 ++---
 programs/bpf/tests/programs.rs        |  4 ++--
 programs/bpf_loader/src/lib.rs        |  3 ++-
 programs/bpf_loader/src/syscalls.rs   | 12 +++++++++---
 sdk/src/transaction_context.rs        |  8 ++++++--
 5 files changed, 21 insertions(+), 11 deletions(-)

diff --git a/program-runtime/src/invoke_context.rs b/program-runtime/src/invoke_context.rs
index f5af1eb3fe..d43f863efa 100644
--- a/program-runtime/src/invoke_context.rs
+++ b/program-runtime/src/invoke_context.rs
@@ -181,7 +181,6 @@ impl<'a> StackFrame<'a> {
 
 pub struct InvokeContext<'a> {
     pub transaction_context: &'a mut TransactionContext,
-    pub return_data: (Pubkey, Vec<u8>),
     invoke_stack: Vec<StackFrame<'a>>,
     rent: Rent,
     pre_accounts: Vec<PreAccount>,
@@ -218,7 +217,6 @@ impl<'a> InvokeContext<'a> {
     ) -> Self {
         Self {
             transaction_context,
-            return_data: (Pubkey::default(), Vec::new()),
             invoke_stack: Vec::with_capacity(compute_budget.max_invoke_depth),
             rent,
             pre_accounts: Vec::new(),
@@ -835,7 +833,8 @@ impl<'a> InvokeContext<'a> {
             .and_then(|_| {
                 let mut process_executable_chain_time =
                     Measure::start("process_executable_chain_time");
-                self.return_data = (program_id, Vec::new());
+                self.transaction_context
+                    .set_return_data(program_id, Vec::new())?;
                 let pre_remaining_units = self.compute_meter.borrow().get_remaining();
                 let execution_result = self.process_executable_chain(instruction_data);
                 let post_remaining_units = self.compute_meter.borrow().get_remaining();
diff --git a/programs/bpf/tests/programs.rs b/programs/bpf/tests/programs.rs
index 55f661570c..b1c1f49eae 100644
--- a/programs/bpf/tests/programs.rs
+++ b/programs/bpf/tests/programs.rs
@@ -226,13 +226,13 @@ fn run_program(name: &str) -> u64 {
         let mut instruction_count = 0;
         let mut tracer = None;
         for i in 0..2 {
-            invoke_context.return_data = (
+            invoke_context.transaction_context.set_return_data(
                 *invoke_context
                     .transaction_context
                     .get_program_key()
                     .unwrap(),
                 Vec::new(),
-            );
+            ).unwrap();
             let mut parameter_bytes = parameter_bytes.clone();
             {
                 let mut vm = create_vm(
diff --git a/programs/bpf_loader/src/lib.rs b/programs/bpf_loader/src/lib.rs
index bfb08b10d6..d14ad4264f 100644
--- a/programs/bpf_loader/src/lib.rs
+++ b/programs/bpf_loader/src/lib.rs
@@ -1099,7 +1099,8 @@ impl Executor for BpfExecutor {
                 trace!("BPF Program Instruction Trace:\n{}", trace_string);
             }
             drop(vm);
-            let (_returned_from_program_id, return_data) = &invoke_context.return_data;
+            let (_returned_from_program_id, return_data) =
+                invoke_context.transaction_context.get_return_data();
             if !return_data.is_empty() {
                 stable_log::program_return(&log_collector, &program_id, return_data);
             }
diff --git a/programs/bpf_loader/src/syscalls.rs b/programs/bpf_loader/src/syscalls.rs
index b4a53e5fb6..f2eaf36313 100644
--- a/programs/bpf_loader/src/syscalls.rs
+++ b/programs/bpf_loader/src/syscalls.rs
@@ -2749,14 +2749,20 @@ impl<'a, 'b> SyscallObject<BpfError> for SyscallSetReturnData<'a, 'b> {
             )
             .to_vec()
         };
-        let program_id = question_mark!(
+        let program_id = *question_mark!(
             invoke_context
                 .transaction_context
                 .get_program_key()
                 .map_err(SyscallError::InstructionError),
             result
         );
-        invoke_context.return_data = (*program_id, return_data);
+        question_mark!(
+            invoke_context
+                .transaction_context
+                .set_return_data(program_id, return_data)
+                .map_err(SyscallError::InstructionError),
+            result
+        );
 
         *result = Ok(0);
     }
@@ -2799,7 +2805,7 @@ impl<'a, 'b> SyscallObject<BpfError> for SyscallGetReturnData<'a, 'b> {
             result
         );
 
-        let (program_id, return_data) = &invoke_context.return_data;
+        let (program_id, return_data) = invoke_context.transaction_context.get_return_data();
         length = length.min(return_data.len() as u64);
         if length != 0 {
             question_mark!(
diff --git a/sdk/src/transaction_context.rs b/sdk/src/transaction_context.rs
index 2cc6ca8a9f..318c6e1fc7 100644
--- a/sdk/src/transaction_context.rs
+++ b/sdk/src/transaction_context.rs
@@ -180,8 +180,12 @@ impl TransactionContext {
     }
 
     /// Set the return data of the current InstructionContext
-    pub fn set_return_data(&mut self, data: Vec<u8>) -> Result<(), InstructionError> {
-        self.return_data = (*self.get_program_key()?, data);
+    pub fn set_return_data(
+        &mut self,
+        program_id: Pubkey,
+        data: Vec<u8>,
+    ) -> Result<(), InstructionError> {
+        self.return_data = (program_id, data);
         Ok(())
     }
 }