From 64235b638634f49febc7cc264d432991a7ad8bb6 Mon Sep 17 00:00:00 2001 From: Sean Young Date: Mon, 29 Apr 2024 10:38:49 +0000 Subject: [PATCH] Ensure mapping of callee is updated with direct mapping 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 | 19 ++- programs/sbf/rust/invoke/src/lib.rs | 43 +++++ programs/sbf/rust/invoke_dep/src/lib.rs | 1 + programs/sbf/tests/programs.rs | 215 +++++++++++++++++++++++- 4 files changed, 272 insertions(+), 6 deletions(-) diff --git a/programs/bpf_loader/src/syscalls/cpi.rs b/programs/bpf_loader/src/syscalls/cpi.rs index 71d5736f895e70..94046f5f741560 100644 --- a/programs/bpf_loader/src/syscalls/cpi.rs +++ b/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,6 +1173,9 @@ fn cpi_common( // // 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, @@ -1180,7 +1183,9 @@ fn update_callee_account( caller_account: &CallerAccount, mut callee_account: BorrowedAccount<'_>, direct_mapping: bool, -) -> Result<(), Error> { +) -> Result { + 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::( 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( diff --git a/programs/sbf/rust/invoke/src/lib.rs b/programs/sbf/rust/invoke/src/lib.rs index 72d23245ce7150..82461c2bc5595c 100644 --- a/programs/sbf/rust/invoke/src/lib.rs +++ b/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::() { + 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"), } diff --git a/programs/sbf/rust/invoke_dep/src/lib.rs b/programs/sbf/rust/invoke_dep/src/lib.rs index b335fb52f5b6b1..f33515f602bd8f 100644 --- a/programs/sbf/rust/invoke_dep/src/lib.rs +++ b/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; diff --git a/programs/sbf/tests/programs.rs b/programs/sbf/tests/programs.rs index 62f12bcef1d823..38ea3afffaa024 100644 --- a/programs/sbf/tests/programs.rs +++ b/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 = (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 = (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 = (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 = (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}"); + }); + } +}