Impl get_/set_return_data syscalls for ProgramTest (#22647)

* Remove &mut self from set_return_data

* Impl get_/set_return_data for program-test SyscallStubs

* Add return_data program-test

program-test/src/lib.rs

@@ -355,6 +355,23 @@ impl solana_sdk::program_stubs::SyscallStubs for SyscallStubs {
     fn sol_get_rent_sysvar(&self, var_addr: *mut u8) -> u64 {
         get_sysvar(get_invoke_context().get_sysvar_cache().get_rent(), var_addr)
     }
+
+    fn sol_get_return_data(&self) -> Option<(Pubkey, Vec<u8>)> {
+        let (program_id, data) = get_invoke_context().transaction_context.get_return_data();
+        Some((*program_id, data.to_vec()))
+    }
+
+    fn sol_set_return_data(&self, data: &[u8]) {
+        let invoke_context = get_invoke_context();
+        let caller = *invoke_context
+            .transaction_context
+            .get_program_key()
+            .unwrap();
+        invoke_context
+            .transaction_context
+            .set_return_data(caller, data.to_vec())
+            .unwrap();
+    }
 }
 
 pub fn find_file(filename: &str) -> Option<PathBuf> {

program-test/tests/return_data.rs

@@ -0,0 +1,87 @@
+use {
+    solana_program_test::{processor, ProgramTest},
+    solana_sdk::{
+        account_info::{next_account_info, AccountInfo},
+        entrypoint::ProgramResult,
+        instruction::{AccountMeta, Instruction},
+        msg,
+        program::{get_return_data, invoke, set_return_data},
+        pubkey::Pubkey,
+        signature::Signer,
+        transaction::Transaction,
+    },
+    std::str::from_utf8,
+};
+
+// Process instruction to get return data from another program
+fn get_return_data_process_instruction(
+    _program_id: &Pubkey,
+    accounts: &[AccountInfo],
+    input: &[u8],
+) -> ProgramResult {
+    msg!("Processing get_return_data instruction before CPI");
+    let account_info_iter = &mut accounts.iter();
+    let invoked_program_info = next_account_info(account_info_iter)?;
+    invoke(
+        &Instruction {
+            program_id: *invoked_program_info.key,
+            accounts: vec![],
+            data: input.to_vec(),
+        },
+        &[invoked_program_info.clone()],
+    )?;
+    let return_data = get_return_data().unwrap();
+    msg!("Processing get_return_data instruction after CPI");
+    msg!("{}", from_utf8(&return_data.1).unwrap());
+    assert_eq!(return_data.1, input.to_vec());
+    Ok(())
+}
+
+// Process instruction to echo input back to another program
+#[allow(clippy::unnecessary_wraps)]
+fn set_return_data_process_instruction(
+    _program_id: &Pubkey,
+    _accounts: &[AccountInfo],
+    input: &[u8],
+) -> ProgramResult {
+    msg!("Processing invoked instruction before set_return_data");
+    set_return_data(input);
+    msg!("Processing invoked instruction after set_return_data");
+    Ok(())
+}
+
+#[tokio::test]
+async fn return_data() {
+    let get_return_data_program_id = Pubkey::new_unique();
+    let mut program_test = ProgramTest::new(
+        "get_return_data",
+        get_return_data_program_id,
+        processor!(get_return_data_process_instruction),
+    );
+    let set_return_data_program_id = Pubkey::new_unique();
+    program_test.add_program(
+        "set_return_data",
+        set_return_data_program_id,
+        processor!(set_return_data_process_instruction),
+    );
+
+    let mut context = program_test.start_with_context().await;
+    let instructions = vec![Instruction {
+        program_id: get_return_data_program_id,
+        accounts: vec![AccountMeta::new_readonly(set_return_data_program_id, false)],
+        data: vec![240, 159, 166, 150],
+    }];
+
+    let transaction = Transaction::new_signed_with_payer(
+        &instructions,
+        Some(&context.payer.pubkey()),
+        &[&context.payer],
+        context.last_blockhash,
+    );
+
+    context
+        .banks_client
+        .process_transaction(transaction)
+        .await
+        .unwrap();
+}

sdk/program/src/program_stubs.rs

@@ -87,7 +87,7 @@ pub trait SyscallStubs: Sync + Send {
     fn sol_get_return_data(&self) -> Option<(Pubkey, Vec<u8>)> {
         None
     }
-    fn sol_set_return_data(&mut self, _data: &[u8]) {}
+    fn sol_set_return_data(&self, _data: &[u8]) {}
     fn sol_log_data(&self, fields: &[&[u8]]) {
         println!("data: {}", fields.iter().map(base64::encode).join(" "));
     }
@@ -170,7 +170,7 @@ pub(crate) fn sol_get_return_data() -> Option<(Pubkey, Vec<u8>)> {
 }
 
 pub(crate) fn sol_set_return_data(data: &[u8]) {
-    SYSCALL_STUBS.write().unwrap().sol_set_return_data(data)
+    SYSCALL_STUBS.read().unwrap().sol_set_return_data(data)
 }
 
 pub(crate) fn sol_log_data(data: &[&[u8]]) {