Allow programs to realloc their accounts within limits (#19475)

programs/bpf/rust/realloc/src/instructions.rs

@@ -0,0 +1,71 @@
+//! @brief Example Rust-based BPF realloc test program
+
+use solana_program::{
+    instruction::{AccountMeta, Instruction},
+    pubkey::Pubkey,
+};
+
+pub const REALLOC: u8 = 1;
+pub const REALLOC_EXTEND: u8 = 2;
+pub const REALLOC_EXTEND_AND_FILL: u8 = 3;
+pub const REALLOC_AND_ASSIGN: u8 = 4;
+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];
+    instruction_data.extend_from_slice(&size.to_le_bytes());
+
+    *bump += 1;
+
+    Instruction::new_with_bytes(
+        *program_id,
+        &instruction_data,
+        vec![AccountMeta::new(*address, false)],
+    )
+}
+
+pub fn realloc_extend(
+    program_id: &Pubkey,
+    address: &Pubkey,
+    size: usize,
+    bump: &mut u8,
+) -> Instruction {
+    let mut instruction_data = vec![REALLOC_EXTEND, *bump];
+    instruction_data.extend_from_slice(&size.to_le_bytes());
+
+    *bump += 1;
+
+    Instruction::new_with_bytes(
+        *program_id,
+        &instruction_data,
+        vec![AccountMeta::new(*address, false)],
+    )
+}
+
+pub fn realloc_extend_and_fill(
+    program_id: &Pubkey,
+    address: &Pubkey,
+    size: usize,
+    fill: u8,
+    bump: &mut u64,
+) -> Instruction {
+    let mut instruction_data = vec![
+        REALLOC_EXTEND_AND_FILL,
+        fill,
+        *bump as u8,
+        (*bump / 255) as u8,
+    ];
+    instruction_data.extend_from_slice(&size.to_le_bytes());
+
+    *bump += 1;
+
+    Instruction::new_with_bytes(
+        *program_id,
+        &instruction_data,
+        vec![AccountMeta::new(*address, false)],
+    )
+}

programs/bpf/rust/realloc/src/lib.rs

@@ -0,0 +1,134 @@
+//! @brief Example Rust-based BPF realloc test program
+
+pub mod instructions;
+
+extern crate solana_program;
+use crate::instructions::*;
+use solana_program::{
+    account_info::AccountInfo, entrypoint, entrypoint::ProgramResult,
+    entrypoint::MAX_PERMITTED_DATA_INCREASE, msg, program::invoke, pubkey::Pubkey,
+    system_instruction, system_program,
+};
+use std::convert::TryInto;
+
+entrypoint!(process_instruction);
+#[allow(clippy::unnecessary_wraps)]
+fn process_instruction(
+    program_id: &Pubkey,
+    accounts: &[AccountInfo],
+    instruction_data: &[u8],
+) -> ProgramResult {
+    let account = &accounts[0];
+
+    match instruction_data[0] {
+        REALLOC => {
+            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, 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, false)?;
+            assert_eq!(new_len, account.data_len());
+        }
+        REALLOC_EXTEND_AND_FILL => {
+            let pre_len = account.data_len();
+            let fill = instruction_data[2];
+            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, 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, 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, false)?;
+            assert_eq!(pre_len + MAX_PERMITTED_DATA_INCREASE, account.data_len());
+            invoke(
+                &system_instruction::assign(account.key, program_id),
+                accounts,
+            )?;
+            assert_eq!(account.owner, program_id);
+        }
+        ASSIGN_TO_SELF_VIA_SYSTEM_PROGRAM_AND_REALLOC => {
+            msg!("assign to self via system program and realloc");
+            let pre_len = account.data_len();
+            invoke(
+                &system_instruction::assign(account.key, program_id),
+                accounts,
+            )?;
+            assert_eq!(account.owner, program_id);
+            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, false)?;
+            assert_eq!(account.data_len(), 0);
+            account.assign(accounts[1].key);
+            assert_eq!(account.owner, accounts[1].key);
+        }
+        CHECK => {
+            msg!("check");
+            assert_eq!(100, account.data_len());
+            let data = account.try_borrow_mut_data()?;
+            for x in data[0..5].iter() {
+                assert_eq!(0, *x);
+            }
+            for x in data[5..].iter() {
+                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!(),
+    }
+
+    Ok(())
+}