Skip to content

Commit

Permalink
Ensure mapping of callee is updated with direct mapping (anza-xyz#1093)
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
seanyoung authored and neutrinoks committed Jul 17, 2024
1 parent 5e78916 commit 7b144af
Show file tree
Hide file tree
Showing 4 changed files with 272 additions and 6 deletions.
19 changes: 14 additions & 5 deletions programs/bpf_loader/src/syscalls/cpi.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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
Expand Down Expand Up @@ -1173,14 +1173,19 @@ fn cpi_common<S: SyscallInvokeSigned>(
//
// 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,
is_loader_deprecated: bool,
caller_account: &CallerAccount,
mut callee_account: BorrowedAccount<'_>,
direct_mapping: bool,
) -> Result<(), Error> {
) -> Result<bool, Error> {
let mut must_update_caller = false;

if callee_account.get_lamports() != *caller_account.lamports {
callee_account.set_lamports(*caller_account.lamports)?;
}
Expand All @@ -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::<u8>(
memory_mapping,
Expand Down Expand Up @@ -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(
Expand Down
43 changes: 43 additions & 0 deletions programs/sbf/rust/invoke/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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::<usize>() {
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"),
}

Expand Down
1 change: 1 addition & 0 deletions programs/sbf/rust/invoke_dep/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
215 changes: 214 additions & 1 deletion programs/sbf/tests/programs.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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")]
Expand Down Expand Up @@ -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<u8> = (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<u8> = (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<u8> = (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<u8> = (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}");
});
}
}

0 comments on commit 7b144af

Please sign in to comment.