Keep track of data reallocations in BorrowedAccount

programs/bpf_loader/src/syscalls/cpi.rs

@@ -112,6 +112,7 @@ struct CallerAccount<'a, 'b> {
     // the pointer field and ref_to_len_in_vm points to the length field.
     vm_data_addr: u64,
     ref_to_len_in_vm: VmValue<'b, 'a, u64>,
+    realloc_counter: usize,
 }
 
 impl<'a, 'b> CallerAccount<'a, 'b> {
@@ -257,6 +258,7 @@ impl<'a, 'b> CallerAccount<'a, 'b> {
             serialized_data,
             vm_data_addr,
             ref_to_len_in_vm,
+            realloc_counter: 0,
         })
     }
 
@@ -366,6 +368,7 @@ impl<'a, 'b> CallerAccount<'a, 'b> {
             serialized_data,
             vm_data_addr: account_info.data_addr,
             ref_to_len_in_vm,
+            realloc_counter: 0,
         })
     }
 
@@ -934,7 +937,7 @@ where
                 return Err(Box::new(SyscallError::InvalidLength));
             }
             #[allow(clippy::indexing_slicing)]
-            let caller_account =
+            let mut caller_account =
                 do_translate(
                     invoke_context,
                     memory_mapping,
@@ -944,6 +947,7 @@ where
                     &account_infos[caller_account_index],
                     serialized_metadata,
                 )?;
+            caller_account.realloc_counter = callee_account.get_realloc_counter();
 
             // before initiating CPI, the caller may have modified the
             // account (caller_account). We need to update the corresponding
@@ -1362,8 +1366,12 @@ fn update_caller_account(
             // because of BorrowedAccount::make_data_mut or by a program that uses the
             // AccountSharedData API directly (deprecated).
             let callee_ptr = callee_account.get_data().as_ptr() as u64;
-            if region.host_addr.get() != callee_ptr {
+            let realloc_counter = callee_account.get_realloc_counter();
+            if realloc_counter != caller_account.realloc_counter
+                || region.host_addr.get() != callee_ptr
+            {
                 region.host_addr.set(callee_ptr);
+                caller_account.realloc_counter = realloc_counter;
                 zero_all_mapped_spare_capacity = true;
             }
         }
@@ -2708,6 +2716,7 @@ mod tests {
                 serialized_data: data,
                 vm_data_addr: self.vm_addr + mem::size_of::<u64>() as u64,
                 ref_to_len_in_vm: VmValue::Translated(&mut self.len),
+                realloc_counter: 0,
             }
         }
     }

transaction-context/src/lib.rs

@@ -266,17 +266,6 @@ impl TransactionContext {
             .ok_or(InstructionError::NotEnoughAccountKeys)
     }
 
-    /// Searches for an account by its key
-    #[cfg(not(target_os = "solana"))]
-    pub fn get_account_at_index(
-        &self,
-        index_in_transaction: IndexOfAccount,
-    ) -> Result<&RefCell<AccountSharedData>, InstructionError> {
-        self.accounts
-            .get(index_in_transaction)
-            .ok_or(InstructionError::NotEnoughAccountKeys)
-    }
-
     /// Searches for an account by its key
     pub fn find_index_of_account(&self, pubkey: &Pubkey) -> Option<IndexOfAccount> {
         self.account_keys
@@ -413,7 +402,9 @@ impl TransactionContext {
                 .and_then(|instruction_context| {
                     // Verify all executable accounts have no outstanding refs
                     for account_index in instruction_context.program_accounts.iter() {
-                        self.get_account_at_index(*account_index)?
+                        self.accounts
+                            .get(*account_index)
+                            .ok_or(InstructionError::NotEnoughAccountKeys)?
                             .try_borrow_mut()
                             .map_err(|_| InstructionError::AccountBorrowOutstanding)?;
                     }
@@ -464,13 +455,16 @@ impl TransactionContext {
             }
             let index_in_transaction = instruction_context
                 .get_index_of_instruction_account_in_transaction(instruction_account_index)?;
-            instruction_accounts_lamport_sum = (self
-                .get_account_at_index(index_in_transaction)?
-                .try_borrow()
-                .map_err(|_| InstructionError::AccountBorrowOutstanding)?
-                .lamports() as u128)
-                .checked_add(instruction_accounts_lamport_sum)
-                .ok_or(InstructionError::ArithmeticOverflow)?;
+            instruction_accounts_lamport_sum = u128::from(
+                self.accounts
+                    .get(index_in_transaction)
+                    .ok_or(InstructionError::NotEnoughAccountKeys)?
+                    .try_borrow()
+                    .map_err(|_| InstructionError::AccountBorrowOutstanding)?
+                    .lamports(),
+            )
+            .checked_add(instruction_accounts_lamport_sum)
+            .ok_or(InstructionError::ArithmeticOverflow)?;
         }
         Ok(instruction_accounts_lamport_sum)
     }
@@ -661,6 +655,7 @@ impl InstructionContext {
             index_in_transaction,
             index_in_instruction,
             account,
+            reallocation_count: 0,
         })
     }
 
@@ -758,6 +753,8 @@ pub struct BorrowedAccount<'a> {
     index_in_transaction: IndexOfAccount,
     index_in_instruction: IndexOfAccount,
     account: RefMut<'a, AccountSharedData>,
+    /// Increased each time the data in account is reallocated to a different address
+    reallocation_count: usize,
 }
 
 impl BorrowedAccount<'_> {
@@ -786,6 +783,12 @@ impl BorrowedAccount<'_> {
         self.account.owner()
     }
 
+    /// Returns the counter which is increased every time the data is reallocated to a different address
+    #[inline]
+    pub fn get_realloc_counter(&self) -> usize {
+        self.reallocation_count
+    }
+
     /// Assignes the owner of this account (transaction wide)
     #[cfg(not(target_os = "solana"))]
     pub fn set_owner(&mut self, pubkey: &[u8]) -> Result<(), InstructionError> {
@@ -904,7 +907,11 @@ impl BorrowedAccount<'_> {
         self.touch()?;
 
         self.update_accounts_resize_delta(data.len())?;
+        let p = self.account.data().as_ptr();
         self.account.set_data(data);
+        if self.account.data().as_ptr() != p {
+            self.reallocation_count = self.reallocation_count.saturating_add(1);
+        }
         Ok(())
     }
 
@@ -922,7 +929,11 @@ impl BorrowedAccount<'_> {
         // allocate + memcpy the current data if self.account is shared. We don't need the memcpy
         // here tho because account.set_data_from_slice(data) is going to replace the content
         // anyway.
+        let p = self.account.data().as_ptr();
         self.account.set_data_from_slice(data);
+        if self.account.data().as_ptr() != p {
+            self.reallocation_count = self.reallocation_count.saturating_add(1);
+        }
 
         Ok(())
     }
@@ -940,7 +951,11 @@ impl BorrowedAccount<'_> {
         }
         self.touch()?;
         self.update_accounts_resize_delta(new_length)?;
+        let p = self.account.data().as_ptr();
         self.account.resize(new_length, 0);
+        if self.account.data().as_ptr() != p {
+            self.reallocation_count = self.reallocation_count.saturating_add(1);
+        }
         Ok(())
     }
 
@@ -1008,7 +1023,11 @@ impl BorrowedAccount<'_> {
         // NOTE: The account memory region CoW code in bpf_loader::create_vm() implements the same
         // logic and must be kept in sync.
         if self.account.is_shared() {
+            let p = self.account.data().as_ptr();
             self.account.reserve(MAX_PERMITTED_DATA_INCREASE);
+            if self.account.data().as_ptr() != p {
+                self.reallocation_count = self.reallocation_count.saturating_add(1);
+            }
         }
     }