cpi: add tests that check that CPI updates all accounts at once

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

@@ -828,7 +828,7 @@ fn process_instruction(
             assert_eq!(account.data_len(), rc_box_size);
 
             // update the serialized length so we don't error out early with AccountDataSizeChanged
-            unsafe { *serialized_len_ptr = account.data_len() as u64 };
+            unsafe { *serialized_len_ptr = rc_box_size as u64 };
 
             // hack to avoid dropping the RcBox, which is supposed to be on the
             // heap but we put it into account data. If we don't do this,
@@ -847,13 +847,17 @@ fn process_instruction(
             const INVOKE_PROGRAM_INDEX: usize = 3;
             const REALLOC_PROGRAM_INDEX: usize = 4;
             let account = &accounts[ARGUMENT_INDEX];
+            let target_account_index = instruction_data[1] as usize;
+            let target_account = &accounts[target_account_index];
             let realloc_program_id = accounts[REALLOC_PROGRAM_INDEX].key;
             let invoke_program_id = accounts[INVOKE_PROGRAM_INDEX].key;
             account.realloc(0, false).unwrap();
             account.assign(realloc_program_id);
+            target_account.realloc(0, false).unwrap();
+            target_account.assign(realloc_program_id);
 
             let rc_box_addr =
-                account.data.borrow_mut().as_mut_ptr() as *mut RcBox<RefCell<&mut [u8]>>;
+                target_account.data.borrow_mut().as_mut_ptr() as *mut RcBox<RefCell<&mut [u8]>>;
             let rc_box_size = mem::size_of::<RcBox<RefCell<&mut [u8]>>>();
             unsafe {
                 std::ptr::write(
@@ -876,6 +880,8 @@ fn process_instruction(
                 );
             }
 
+            let serialized_len_ptr =
+                unsafe { account.data.borrow_mut().as_mut_ptr().offset(-8) as *mut u64 };
             // Place a RcBox<RefCell<&mut [u8]>> in the account data. This
             // allows us to test having CallerAccount::ref_to_len_in_vm in an
             // account region.
@@ -897,6 +903,7 @@ fn process_instruction(
                     *realloc_program_id,
                     &[
                         (accounts[ARGUMENT_INDEX].key, true, false),
+                        (target_account.key, true, false),
                         (realloc_program_id, false, false),
                         (invoke_program_id, false, false),
                     ],
@@ -906,6 +913,7 @@ fn process_instruction(
             )
             .unwrap();
 
+            unsafe { *serialized_len_ptr = rc_box_size as u64 };
             // hack to avoid dropping the RcBox, which is supposed to be on the
             // heap but we put it into account data. If we don't do this,
             // dropping the Rc will cause

programs/sbf/tests/programs.rs

@@ -4066,34 +4066,57 @@ fn test_cpi_account_ownership_writability() {
             }
         );
 
-        // Similar to the test above where we try to make CPI write into account
-        // data. This variant is for when direct mapping is enabled.
-        let account = AccountSharedData::new(42, 0, &invoke_program_id);
-        bank.store_account(&account_keypair.pubkey(), &account);
-        let instruction_data = vec![TEST_FORBID_LEN_UPDATE_AFTER_OWNERSHIP_CHANGE, 42, 42, 42];
-        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);
-        if direct_mapping {
-            assert_eq!(
-                result.unwrap_err().unwrap(),
-                TransactionError::InstructionError(
-                    0,
-                    InstructionError::ExternalAccountDataModified
-                )
+        // We're going to try and make CPI write ref_to_len_in_vm into a 2nd
+        // account, so we add an extra one here.
+        let account2_keypair = Keypair::new();
+        let mut account_metas = account_metas.clone();
+        account_metas.push(AccountMeta::new(account2_keypair.pubkey(), false));
+
+        for target_account in [1, account_metas.len() as u8 - 1] {
+            // Similar to the test above where we try to make CPI write into account
+            // data. This variant is for when direct mapping is enabled.
+            let account = AccountSharedData::new(42, 0, &invoke_program_id);
+            bank.store_account(&account_keypair.pubkey(), &account);
+            let account = AccountSharedData::new(42, 0, &invoke_program_id);
+            bank.store_account(&account2_keypair.pubkey(), &account);
+            let instruction_data = vec![
+                TEST_FORBID_LEN_UPDATE_AFTER_OWNERSHIP_CHANGE,
+                target_account,
+                42,
+                42,
+            ];
+            let instruction = Instruction::new_with_bytes(
+                invoke_program_id,
+                &instruction_data,
+                account_metas.clone(),
             );
-        } else {
-            // we expect this to succeed as after updating `ref_to_len_in_vm`,
-            // CPI will sync the actual account data between the callee and the
-            // caller, _always_ writing over the location pointed by
-            // `ref_to_len_in_vm`. To verify this, we check that the account
-            // data is in fact all zeroes like it is in the callee.
-            result.unwrap();
-            let account = bank.get_account(&account_keypair.pubkey()).unwrap();
-            assert_eq!(account.data(), vec![0; 40]);
+            let message = Message::new(&[instruction], Some(&mint_pubkey));
+            let tx = Transaction::new(&[&mint_keypair], message.clone(), bank.last_blockhash());
+            let (result, _, logs) = process_transaction_and_record_inner(&bank, tx);
+            if direct_mapping {
+                assert_eq!(
+                    result.unwrap_err(),
+                    TransactionError::InstructionError(
+                        0,
+                        InstructionError::ProgramFailedToComplete
+                    )
+                );
+                // We haven't moved the data pointer, but ref_to_len_vm _is_ in
+                // the account data vm range and that's not allowed either.
+                assert!(
+                    logs.iter().any(|log| log.contains("Invalid pointer")),
+                    "{logs:?}"
+                );
+            } else {
+                // we expect this to succeed as after updating `ref_to_len_in_vm`,
+                // CPI will sync the actual account data between the callee and the
+                // caller, _always_ writing over the location pointed by
+                // `ref_to_len_in_vm`. To verify this, we check that the account
+                // data is in fact all zeroes like it is in the callee.
+                result.unwrap();
+                let account = bank.get_account(&account_keypair.pubkey()).unwrap();
+                assert_eq!(account.data(), vec![0; 40]);
+            }
         }
     }
 }