Ensure mapping of callee is updated with direct mapping (anza-xyz#1093)

Consider this scenario:

 - Program increases length of an account
 - Program start CPI and adds this account as a read-only account
 - In fn update_callee_account() we resize account, which may change
   the pointer
 - Once CPI finishes, the program continues and may read/write from
   the account. The mapping must be up-to-date else we use stale
   pointers.

Note that we always call callee_account.set_data_length(), which
may change the pointer. In testing I found that resizing a vector
from 10240 down to 127 sometimes changes its pointer. So, always
update the pointer.

programs/bpf_loader/src/syscalls/cpi.rs

@@ -925,7 +925,7 @@ where
             // account (caller_account). We need to update the corresponding
             // BorrowedAccount (callee_account) so the callee can see the
             // changes.
-            update_callee_account(
+            let update_caller = update_callee_account(
                 invoke_context,
                 memory_mapping,
                 is_loader_deprecated,
@@ -934,7 +934,7 @@ where
                 direct_mapping,
             )?;
 
-            let caller_account = if instruction_account.is_writable {
+            let caller_account = if instruction_account.is_writable || update_caller {
                 Some(caller_account)
             } else {
                 None
@@ -1173,14 +1173,19 @@ fn cpi_common<S: SyscallInvokeSigned>(
 //
 // This method updates callee_account so the CPI callee can see the caller's
 // changes.
+//
+// When true is returned, the caller account must be updated after CPI. This
+// is only set for direct mapping when the pointer may have changed.
 fn update_callee_account(
     invoke_context: &InvokeContext,
     memory_mapping: &MemoryMapping,
     is_loader_deprecated: bool,
     caller_account: &CallerAccount,
     mut callee_account: BorrowedAccount<'_>,
     direct_mapping: bool,
-) -> Result<(), Error> {
+) -> Result<bool, Error> {
+    let mut must_update_caller = false;
+
     if callee_account.get_lamports() != *caller_account.lamports {
         callee_account.set_lamports(*caller_account.lamports)?;
     }
@@ -1198,7 +1203,11 @@ fn update_callee_account(
                 if is_loader_deprecated && realloc_bytes_used > 0 {
                     return Err(InstructionError::InvalidRealloc.into());
                 }
-                callee_account.set_data_length(post_len)?;
+                if prev_len != post_len {
+                    callee_account.set_data_length(post_len)?;
+                    // pointer to data may have changed, so caller must be updated
+                    must_update_caller = true;
+                }
                 if realloc_bytes_used > 0 {
                     let serialized_data = translate_slice::<u8>(
                         memory_mapping,
@@ -1239,7 +1248,7 @@ fn update_callee_account(
         callee_account.set_owner(caller_account.owner.as_ref())?;
     }
 
-    Ok(())
+    Ok(must_update_caller)
 }
 
 fn update_caller_account_perms(

programs/sbf/rust/invoke/src/lib.rs

@@ -1350,6 +1350,49 @@ fn process_instruction<'a>(
             let byte_index = usize::from_le_bytes(instruction_data[2..10].try_into().unwrap());
             target_account.data.borrow_mut()[byte_index] = instruction_data[10];
         }
+        TEST_CALLEE_ACCOUNT_UPDATES => {
+            msg!("TEST_CALLEE_ACCOUNT_UPDATES");
+
+            if instruction_data.len() < 2 + 2 * std::mem::size_of::<usize>() {
+                return Ok(());
+            }
+
+            let writable = instruction_data[1] != 0;
+            let resize = usize::from_le_bytes(instruction_data[2..10].try_into().unwrap());
+            let write_offset = usize::from_le_bytes(instruction_data[10..18].try_into().unwrap());
+            let invoke_struction = &instruction_data[18..];
+
+            let account = &accounts[ARGUMENT_INDEX];
+
+            if resize != 0 {
+                account.realloc(resize, false).unwrap();
+            }
+
+            if !invoke_struction.is_empty() {
+                // Invoke another program. With direct mapping, before CPI the callee will update the accounts (incl resizing)
+                // so the pointer may change.
+                let invoked_program_id = accounts[INVOKED_PROGRAM_INDEX].key;
+
+                invoke(
+                    &create_instruction(
+                        *invoked_program_id,
+                        &[
+                            (accounts[MINT_INDEX].key, false, false),
+                            (accounts[ARGUMENT_INDEX].key, writable, false),
+                            (invoked_program_id, false, false),
+                        ],
+                        invoke_struction.to_vec(),
+                    ),
+                    accounts,
+                )
+                .unwrap();
+            }
+
+            if write_offset != 0 {
+                // Ensure we still have access to the correct account
+                account.data.borrow_mut()[write_offset] ^= 0xe5;
+            }
+        }
         _ => panic!("unexpected program data"),
     }
 

programs/sbf/rust/invoke_dep/src/lib.rs

@@ -39,6 +39,7 @@ pub const TEST_CPI_INVALID_LAMPORTS_POINTER: u8 = 36;
 pub const TEST_CPI_INVALID_DATA_POINTER: u8 = 37;
 pub const TEST_CPI_CHANGE_ACCOUNT_DATA_MEMORY_ALLOCATION: u8 = 38;
 pub const TEST_WRITE_ACCOUNT: u8 = 39;
+pub const TEST_CALLEE_ACCOUNT_UPDATES: u8 = 40;
 
 pub const MINT_INDEX: usize = 0;
 pub const ARGUMENT_INDEX: usize = 1;

programs/sbf/tests/programs.rs

@@ -85,7 +85,7 @@ use {
         system_program,
         transaction::{SanitizedTransaction, Transaction, TransactionError},
     },
-    std::{cell::RefCell, str::FromStr, sync::Arc, time::Duration},
+    std::{assert_eq, cell::RefCell, str::FromStr, sync::Arc, time::Duration},
 };
 
 #[cfg(feature = "sbf_rust")]
@@ -4463,3 +4463,216 @@ fn test_deny_executable_write() {
         );
     }
 }
+
+#[test]
+fn test_update_callee_account() {
+    // Test that fn update_callee_account() works and we are updating the callee account on CPI.
+    solana_logger::setup();
+
+    let GenesisConfigInfo {
+        genesis_config,
+        mint_keypair,
+        ..
+    } = create_genesis_config(100_123_456_789);
+
+    for direct_mapping in [false, true] {
+        let mut bank = Bank::new_for_tests(&genesis_config);
+        let feature_set = Arc::make_mut(&mut bank.feature_set);
+        // by default test banks have all features enabled, so we only need to
+        // disable when needed
+        if !direct_mapping {
+            feature_set.deactivate(&feature_set::bpf_account_data_direct_mapping::id());
+        }
+        let (bank, bank_forks) = bank.wrap_with_bank_forks_for_tests();
+        let mut bank_client = BankClient::new_shared(bank.clone());
+        let authority_keypair = Keypair::new();
+
+        let (bank, invoke_program_id) = load_upgradeable_program_and_advance_slot(
+            &mut bank_client,
+            bank_forks.as_ref(),
+            &mint_keypair,
+            &authority_keypair,
+            "solana_sbf_rust_invoke",
+        );
+
+        let account_keypair = Keypair::new();
+
+        let mint_pubkey = mint_keypair.pubkey();
+
+        let account_metas = vec![
+            AccountMeta::new(mint_pubkey, true),
+            AccountMeta::new(account_keypair.pubkey(), false),
+            AccountMeta::new_readonly(invoke_program_id, false),
+        ];
+
+        // I. do CPI with account in read only (separate code path with direct mapping)
+        let mut account = AccountSharedData::new(42, 10240, &invoke_program_id);
+        let data: Vec<u8> = (0..10240).map(|n| n as u8).collect();
+        account.set_data(data);
+
+        bank.store_account(&account_keypair.pubkey(), &account);
+
+        let mut instruction_data = vec![TEST_CALLEE_ACCOUNT_UPDATES, 0];
+        instruction_data.extend_from_slice(20480usize.to_le_bytes().as_ref());
+        instruction_data.extend_from_slice(16384usize.to_le_bytes().as_ref());
+        // instruction data for inner CPI (2x)
+        instruction_data.extend_from_slice(&[TEST_CALLEE_ACCOUNT_UPDATES, 0]);
+        instruction_data.extend_from_slice(0usize.to_le_bytes().as_ref());
+        instruction_data.extend_from_slice(0usize.to_le_bytes().as_ref());
+
+        instruction_data.extend_from_slice(&[TEST_CALLEE_ACCOUNT_UPDATES, 0]);
+        instruction_data.extend_from_slice(0usize.to_le_bytes().as_ref());
+        instruction_data.extend_from_slice(0usize.to_le_bytes().as_ref());
+
+        let instruction = Instruction::new_with_bytes(
+            invoke_program_id,
+            &instruction_data,
+            account_metas.clone(),
+        );
+        let result = bank_client.send_and_confirm_instruction(&mint_keypair, instruction);
+        assert!(result.is_ok());
+
+        let data = bank_client
+            .get_account_data(&account_keypair.pubkey())
+            .unwrap()
+            .unwrap();
+
+        assert_eq!(data.len(), 20480);
+
+        data.iter().enumerate().for_each(|(i, v)| {
+            let expected = match i {
+                ..=10240 => i as u8,
+                16384 => 0xe5,
+                _ => 0,
+            };
+
+            assert_eq!(*v, expected, "offset:{i} {v:#x} != {expected:#x}");
+        });
+
+        // II. do CPI with account with resize to smaller and write
+        let mut account = AccountSharedData::new(42, 10240, &invoke_program_id);
+        let data: Vec<u8> = (0..10240).map(|n| n as u8).collect();
+        account.set_data(data);
+        bank.store_account(&account_keypair.pubkey(), &account);
+
+        let mut instruction_data = vec![TEST_CALLEE_ACCOUNT_UPDATES, 1];
+        instruction_data.extend_from_slice(20480usize.to_le_bytes().as_ref());
+        instruction_data.extend_from_slice(16384usize.to_le_bytes().as_ref());
+        // instruction data for inner CPI
+        instruction_data.extend_from_slice(&[TEST_CALLEE_ACCOUNT_UPDATES, 0]);
+        instruction_data.extend_from_slice(19480usize.to_le_bytes().as_ref());
+        instruction_data.extend_from_slice(8129usize.to_le_bytes().as_ref());
+
+        let instruction = Instruction::new_with_bytes(
+            invoke_program_id,
+            &instruction_data,
+            account_metas.clone(),
+        );
+        let result = bank_client.send_and_confirm_instruction(&mint_keypair, instruction);
+        assert!(result.is_ok());
+
+        let data = bank_client
+            .get_account_data(&account_keypair.pubkey())
+            .unwrap()
+            .unwrap();
+
+        assert_eq!(data.len(), 19480);
+
+        data.iter().enumerate().for_each(|(i, v)| {
+            let expected = match i {
+                8129 => (i as u8) ^ 0xe5,
+                ..=10240 => i as u8,
+                16384 => 0xe5,
+                _ => 0,
+            };
+
+            assert_eq!(*v, expected, "offset:{i} {v:#x} != {expected:#x}");
+        });
+
+        // III. do CPI with account with resize to larger and write
+        let mut account = AccountSharedData::new(42, 10240, &invoke_program_id);
+        let data: Vec<u8> = (0..10240).map(|n| n as u8).collect();
+        account.set_data(data);
+        bank.store_account(&account_keypair.pubkey(), &account);
+
+        let mut instruction_data = vec![TEST_CALLEE_ACCOUNT_UPDATES, 1];
+        instruction_data.extend_from_slice(16384usize.to_le_bytes().as_ref());
+        instruction_data.extend_from_slice(16384usize.to_le_bytes().as_ref());
+        // instruction data for inner CPI
+        instruction_data.extend_from_slice(&[TEST_CALLEE_ACCOUNT_UPDATES, 0]);
+        instruction_data.extend_from_slice(20480usize.to_le_bytes().as_ref());
+        instruction_data.extend_from_slice(16385usize.to_le_bytes().as_ref());
+
+        let instruction = Instruction::new_with_bytes(
+            invoke_program_id,
+            &instruction_data,
+            account_metas.clone(),
+        );
+        let result = bank_client.send_and_confirm_instruction(&mint_keypair, instruction);
+        assert!(result.is_ok());
+
+        let data = bank_client
+            .get_account_data(&account_keypair.pubkey())
+            .unwrap()
+            .unwrap();
+
+        assert_eq!(data.len(), 20480);
+
+        data.iter().enumerate().for_each(|(i, v)| {
+            let expected = match i {
+                ..=10240 => i as u8,
+                16384 | 16385 => 0xe5,
+                _ => 0,
+            };
+
+            assert_eq!(*v, expected, "offset:{i} {v:#x} != {expected:#x}");
+        });
+
+        // IV. do CPI with account with resize to larger and write
+        let mut account = AccountSharedData::new(42, 10240, &invoke_program_id);
+        let data: Vec<u8> = (0..10240).map(|n| n as u8).collect();
+        account.set_data(data);
+        bank.store_account(&account_keypair.pubkey(), &account);
+
+        let mut instruction_data = vec![TEST_CALLEE_ACCOUNT_UPDATES, 1];
+        instruction_data.extend_from_slice(16384usize.to_le_bytes().as_ref());
+        instruction_data.extend_from_slice(16384usize.to_le_bytes().as_ref());
+        // instruction data for inner CPI (2x)
+        instruction_data.extend_from_slice(&[TEST_CALLEE_ACCOUNT_UPDATES, 1]);
+        instruction_data.extend_from_slice(0usize.to_le_bytes().as_ref());
+        instruction_data.extend_from_slice(0usize.to_le_bytes().as_ref());
+
+        instruction_data.extend_from_slice(&[TEST_CALLEE_ACCOUNT_UPDATES, 1]);
+        instruction_data.extend_from_slice(0usize.to_le_bytes().as_ref());
+        instruction_data.extend_from_slice(0usize.to_le_bytes().as_ref());
+        // instruction data for inner CPI
+        instruction_data.extend_from_slice(&[TEST_CALLEE_ACCOUNT_UPDATES, 0]);
+        instruction_data.extend_from_slice(20480usize.to_le_bytes().as_ref());
+        instruction_data.extend_from_slice(16385usize.to_le_bytes().as_ref());
+
+        let instruction = Instruction::new_with_bytes(
+            invoke_program_id,
+            &instruction_data,
+            account_metas.clone(),
+        );
+        let result = bank_client.send_and_confirm_instruction(&mint_keypair, instruction);
+        assert!(result.is_ok());
+
+        let data = bank_client
+            .get_account_data(&account_keypair.pubkey())
+            .unwrap()
+            .unwrap();
+
+        assert_eq!(data.len(), 20480);
+
+        data.iter().enumerate().for_each(|(i, v)| {
+            let expected = match i {
+                ..=10240 => i as u8,
+                16384 | 16385 => 0xe5,
+                _ => 0,
+            };
+
+            assert_eq!(*v, expected, "offset:{i} {v:#x} != {expected:#x}");
+        });
+    }
+}