zero-init realloc memory

solana-labs · Sep 24, 2021 · fee60ca · fee60ca
1 parent 1be1516
commit fee60ca
Show file tree

Hide file tree

Showing 7 changed files with 95 additions and 22 deletions.
diff --git a/programs/bpf/benches/bpf_loader.rs b/programs/bpf/benches/bpf_loader.rs
@@ -113,6 +113,7 @@ fn bench_program_alu(bencher: &mut Bencher) {
         executable.as_ref(),
         &mut inner_iter,
         &mut invoke_context,
+        &[],
     )
     .unwrap();
 
@@ -220,7 +221,7 @@ fn bench_instruction_count_tuner(_bencher: &mut Bencher) {
 
     // Serialize account data
     let keyed_accounts = invoke_context.get_keyed_accounts().unwrap();
-    let (mut serialized, _account_lengths) = serialize_parameters(
+    let (mut serialized, account_lengths) = serialize_parameters(
         &bpf_loader::id(),
         &solana_sdk::pubkey::new_rand(),
         keyed_accounts,
@@ -243,6 +244,7 @@ fn bench_instruction_count_tuner(_bencher: &mut Bencher) {
         executable.as_ref(),
         serialized.as_slice_mut(),
         &mut invoke_context,
+        &account_lengths,
     )
     .unwrap();
 

diff --git a/programs/bpf/rust/realloc/src/instructions.rs b/programs/bpf/rust/realloc/src/instructions.rs
@@ -13,6 +13,7 @@ pub const REALLOC_AND_ASSIGN_TO_SELF_VIA_SYSTEM_PROGRAM: u8 = 5;
 pub const ASSIGN_TO_SELF_VIA_SYSTEM_PROGRAM_AND_REALLOC: u8 = 6;
 pub const DEALLOC_AND_ASSIGN_TO_CALLER: u8 = 7;
 pub const CHECK: u8 = 8;
+pub const ZERO_INIT: u8 = 9;
 
 pub fn realloc(program_id: &Pubkey, address: &Pubkey, size: usize, bump: &mut u8) -> Instruction {
     let mut instruction_data = vec![REALLOC, *bump];

diff --git a/programs/bpf/rust/realloc/src/lib.rs b/programs/bpf/rust/realloc/src/lib.rs
@@ -25,15 +25,15 @@ fn process_instruction(
             let (bytes, _) = instruction_data[2..].split_at(std::mem::size_of::<usize>());
             let new_len = usize::from_le_bytes(bytes.try_into().unwrap());
             msg!("realloc to {}", new_len);
-            account.realloc(new_len)?;
+            account.realloc(new_len, false)?;
             assert_eq!(new_len, account.data_len());
         }
         REALLOC_EXTEND => {
             let pre_len = account.data_len();
             let (bytes, _) = instruction_data[2..].split_at(std::mem::size_of::<usize>());
             let new_len = pre_len + usize::from_le_bytes(bytes.try_into().unwrap());
             msg!("realloc extend by {}", new_len);
-            account.realloc(new_len)?;
+            account.realloc(new_len, false)?;
             assert_eq!(new_len, account.data_len());
         }
         REALLOC_EXTEND_AND_FILL => {
@@ -42,21 +42,21 @@ fn process_instruction(
             let (bytes, _) = instruction_data[4..].split_at(std::mem::size_of::<usize>());
             let new_len = pre_len + usize::from_le_bytes(bytes.try_into().unwrap());
             msg!("realloc extend by {}", new_len);
-            account.realloc(new_len)?;
+            account.realloc(new_len, false)?;
             assert_eq!(new_len, account.data_len());
             account.try_borrow_mut_data()?[pre_len..].fill(fill);
         }
         REALLOC_AND_ASSIGN => {
             msg!("realloc and assign");
-            account.realloc(MAX_PERMITTED_DATA_INCREASE)?;
+            account.realloc(MAX_PERMITTED_DATA_INCREASE, false)?;
             assert_eq!(MAX_PERMITTED_DATA_INCREASE, account.data_len());
             account.assign(&system_program::id());
             assert_eq!(*account.owner, system_program::id());
         }
         REALLOC_AND_ASSIGN_TO_SELF_VIA_SYSTEM_PROGRAM => {
             msg!("realloc and assign to self via system program");
             let pre_len = account.data_len();
-            account.realloc(pre_len + MAX_PERMITTED_DATA_INCREASE)?;
+            account.realloc(pre_len + MAX_PERMITTED_DATA_INCREASE, false)?;
             assert_eq!(pre_len + MAX_PERMITTED_DATA_INCREASE, account.data_len());
             invoke(
                 &system_instruction::assign(account.key, program_id),
@@ -72,12 +72,12 @@ fn process_instruction(
                 accounts,
             )?;
             assert_eq!(account.owner, program_id);
-            account.realloc(pre_len + MAX_PERMITTED_DATA_INCREASE)?;
+            account.realloc(pre_len + MAX_PERMITTED_DATA_INCREASE, false)?;
             assert_eq!(account.data_len(), pre_len + MAX_PERMITTED_DATA_INCREASE);
         }
         DEALLOC_AND_ASSIGN_TO_CALLER => {
             msg!("dealloc and assign to caller");
-            account.realloc(0)?;
+            account.realloc(0, false)?;
             assert_eq!(account.data_len(), 0);
             account.assign(accounts[1].key);
             assert_eq!(account.owner, accounts[1].key);
@@ -93,6 +93,40 @@ fn process_instruction(
                 assert_eq!(2, *x);
             }
         }
+        ZERO_INIT => {
+            account.realloc(10, false)?;
+            {
+                let mut data = account.try_borrow_mut_data()?;
+                for i in 0..10 {
+                    assert_eq!(0, data[i]);
+                }
+                data.fill(1);
+                for i in 0..10 {
+                    assert_eq!(1, data[i]);
+                }
+            }
+
+            account.realloc(5, false)?;
+            account.realloc(10, false)?;
+            {
+                let data = account.try_borrow_data()?;
+                for i in 0..10 {
+                    assert_eq!(1, data[i]);
+                }
+            }
+
+            account.realloc(5, false)?;
+            account.realloc(10, true)?;
+            {
+                let data = account.try_borrow_data()?;
+                for i in 0..5 {
+                    assert_eq!(1, data[i]);
+                }
+                for i in 5..10 {
+                    assert_eq!(0, data[i]);
+                }
+            }
+        }
         _ => panic!(),
     }
 

diff --git a/programs/bpf/rust/realloc_invoke/src/lib.rs b/programs/bpf/rust/realloc_invoke/src/lib.rs
@@ -84,7 +84,7 @@ fn process_instruction(
             )?;
             let new_len = pre_len + MAX_PERMITTED_DATA_INCREASE;
             assert_eq!(new_len, account.data_len());
-            account.realloc(new_len + MAX_PERMITTED_DATA_INCREASE)?;
+            account.realloc(new_len + MAX_PERMITTED_DATA_INCREASE, false)?;
             assert_eq!(new_len + MAX_PERMITTED_DATA_INCREASE, account.data_len());
         }
         INVOKE_REALLOC_AND_ASSIGN => {
@@ -130,7 +130,7 @@ fn process_instruction(
         }
         INVOKE_REALLOC_INVOKE_CHECK => {
             msg!("realloc invoke check size");
-            account.realloc(100)?;
+            account.realloc(100, false)?;
             assert_eq!(100, account.data_len());
             account.try_borrow_mut_data()?[pre_len..].fill(2);
             invoke(
@@ -146,7 +146,7 @@ fn process_instruction(
             let (bytes, _) = instruction_data[2..].split_at(std::mem::size_of::<usize>());
             let new_len = usize::from_le_bytes(bytes.try_into().unwrap());
             msg!("realloc to {}", new_len);
-            account.realloc(new_len)?;
+            account.realloc(new_len, false)?;
             assert_eq!(new_len, account.data_len());
             if pre_len < new_len {
                 account.try_borrow_mut_data()?[pre_len..].fill(instruction_data[1]);
@@ -156,7 +156,7 @@ fn process_instruction(
             msg!("realloc invoke recursive");
             let (bytes, _) = instruction_data[2..].split_at(std::mem::size_of::<usize>());
             let new_len = usize::from_le_bytes(bytes.try_into().unwrap());
-            account.realloc(new_len)?;
+            account.realloc(new_len, false)?;
             assert_eq!(new_len, account.data_len());
             account.try_borrow_mut_data()?[pre_len..].fill(instruction_data[1]);
             let final_len: usize = 200;
@@ -197,7 +197,7 @@ fn process_instruction(
                 accounts,
             )?;
             assert_eq!(pre_len, accounts[1].data_len());
-            accounts[1].realloc(pre_len + 1)?;
+            accounts[1].realloc(pre_len + 1, false)?;
             assert_eq!(pre_len + 1, accounts[1].data_len());
             assert_eq!(accounts[1].owner, program_id);
             let final_len: usize = 200;
@@ -243,7 +243,7 @@ fn process_instruction(
             )?;
             assert_eq!(account.owner, program_id);
             assert_eq!(account.data_len(), 0);
-            account.realloc(new_len)?;
+            account.realloc(new_len, false)?;
             assert_eq!(account.data_len(), new_len);
             {
                 let data = account.try_borrow_mut_data()?;
@@ -255,7 +255,7 @@ fn process_instruction(
         INVOKE_REALLOC_MAX_INVOKE_MAX => {
             msg!("invoke realloc max invoke max");
             assert_eq!(0, account.data_len());
-            account.realloc(MAX_PERMITTED_DATA_INCREASE)?;
+            account.realloc(MAX_PERMITTED_DATA_INCREASE, false)?;
             assert_eq!(MAX_PERMITTED_DATA_INCREASE, account.data_len());
             account.assign(invoke_program_id);
             assert_eq!(account.owner, invoke_program_id);

diff --git a/programs/bpf/tests/programs.rs b/programs/bpf/tests/programs.rs
@@ -2781,9 +2781,23 @@ fn test_program_bpf_realloc() {
         .unwrap();
     let data = bank_client.get_account_data(&pubkey).unwrap().unwrap();
     assert_eq!(0, data.len());
+
+    // zero-init
+    bank_client
+        .send_and_confirm_message(
+            &[&mint_keypair, &keypair],
+            Message::new(
+                &[Instruction::new_with_bytes(
+                    program_id,
+                    &[ZERO_INIT],
+                    vec![AccountMeta::new(pubkey, true)],
+                )],
+                Some(&mint_pubkey),
+            ),
+        )
+        .unwrap();
 }
 
-// #[ignore]
 #[cfg(feature = "bpf_rust")]
 #[test]
 fn test_program_bpf_realloc_invoke() {

diff --git a/rbpf-cli/src/main.rs b/rbpf-cli/src/main.rs
@@ -153,7 +153,7 @@ native machine code before execting it in the virtual machine.",
     let mut account_refcells = Vec::new();
     let default_account = RefCell::new(AccountSharedData::default());
     let key = solana_sdk::pubkey::new_rand();
-    let (mut mem, account_lenghts) = match matches.value_of("input").unwrap().parse::<usize>() {
+    let (mut mem, account_lengths) = match matches.value_of("input").unwrap().parse::<usize>() {
         Ok(allocate) => {
             accounts.push(KeyedAccount::new(&key, false, &default_account));
             (vec![0u8; allocate], vec![])
@@ -170,9 +170,9 @@ native machine code before execting it in the virtual machine.",
             }
             let lid = bpf_loader::id();
             let pid = Pubkey::new(&[0u8; 32]);
-            let (mut bytes, account_lenghts) =
+            let (mut bytes, account_lengths) =
                 serialize_parameters(&lid, &pid, &accounts, &input.insndata).unwrap();
-            (Vec::from(bytes.as_slice_mut()), account_lenghts)
+            (Vec::from(bytes.as_slice_mut()), account_lengths)
         }
     };
     let mut invoke_context = MockInvokeContext::new(accounts);
@@ -233,7 +233,7 @@ native machine code before execting it in the virtual machine.",
         executable.as_ref(),
         &mut mem,
         &mut invoke_context,
-        &account_lenghts,
+        &account_lengths,
     )
     .unwrap();
     let start_time = Instant::now();

diff --git a/sdk/program/src/account_info.rs b/sdk/program/src/account_info.rs
@@ -1,4 +1,6 @@
-use crate::{clock::Epoch, program_error::ProgramError, pubkey::Pubkey};
+use crate::{
+    clock::Epoch, program_error::ProgramError, program_memory::sol_memset, pubkey::Pubkey,
+};
 use std::{
     cell::{Ref, RefCell, RefMut},
     cmp, fmt,
@@ -114,7 +116,21 @@ impl<'a> AccountInfo<'a> {
             .map_err(|_| ProgramError::AccountBorrowFailed)
     }
 
-    pub fn realloc(&self, new_len: usize) -> Result<(), ProgramError> {
+    /// Realloc the account's data and optionally zero-initialize the new
+    /// memory.
+    ///
+    /// Note:  Account data can be increased within a single call by up to
+    /// `solana_program::entrypoint::MAX_PERMITTED_DATA_INCREASE` bytes.
+    ///
+    /// Note: Memory used to grow is already zero-initialized upon program
+    /// entrypoint and re-zeroing it wastes compute units.  If within the same
+    /// call a program reallocs from larger to smaller and back to larger again
+    /// the new space could contain stale data.  Pass `true` for `zero_init` in
+    /// this case, otherwise compute units will be wasted re-zero-initializing.
+    pub fn realloc(&self, new_len: usize, zero_init: bool) -> Result<(), ProgramError> {
+        let orig_len = self.data_len();
+
+        // realloc
         unsafe {
             // First set new length in the serialized data
             let ptr = self.try_borrow_mut_data()?.as_mut_ptr().offset(-8) as *mut u64;
@@ -124,6 +140,12 @@ impl<'a> AccountInfo<'a> {
             let ptr = &mut *(((self.data.as_ptr() as *const u64).offset(1) as u64) as *mut u64);
             *ptr = new_len as u64;
         }
+
+        // zero-init if requested
+        if zero_init && new_len > orig_len {
+            sol_memset(&mut self.try_borrow_mut_data()?[orig_len..], 0, new_len);
+        }
+
         Ok(())
     }