Skip to content

Commit

Permalink
cpi: tweak tests
Browse files Browse the repository at this point in the history 
Remove some copy pasta and rename two tests to better describe what they're doing
  • Loading branch information
@alessandrod
alessandrod committed Aug 30, 2023
1 parent 65a0c1c commit 821cb14
Show file tree
Hide file tree
Showing 3 changed files with 38 additions and 91 deletions.
6 changes: 3 additions & 3 deletions programs/sbf/rust/invoke/src/instructions.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,9 +21,9 @@ pub const TEST_RETURN_DATA_TOO_LARGE: u8 = 18;
pub const TEST_DUPLICATE_PRIVILEGE_ESCALATION_SIGNER: u8 = 19;
pub const TEST_DUPLICATE_PRIVILEGE_ESCALATION_WRITABLE: u8 = 20;
pub const TEST_MAX_ACCOUNT_INFOS_EXCEEDED: u8 = 21;
pub const TEST_FORBID_WRITE_AFTER_OWNERSHIP_CHANGE_TO_SYSTEM_ACCOUNT: u8 = 22;
pub const TEST_FORBID_WRITE_AFTER_OWNERSHIP_CHANGE_TO_SYSTEM_ACCOUNT_NESTED: u8 = 23;
pub const TEST_FORBID_WRITE_AFTER_OWNERSHIP_CHANGE_TO_OTHER_PROGRAM: u8 = 24;
pub const TEST_FORBID_WRITE_AFTER_OWNERSHIP_CHANGE_IN_CALLEE: u8 = 22;
pub const TEST_FORBID_WRITE_AFTER_OWNERSHIP_CHANGE_IN_CALLEE_NESTED: u8 = 23;
pub const TEST_FORBID_WRITE_AFTER_OWNERSHIP_CHANGE_IN_CALLER: u8 = 24;
pub const TEST_FORBID_LEN_UPDATE_AFTER_OWNERSHIP_CHANGE_MOVING_DATA_POINTER: u8 = 25;
pub const TEST_FORBID_LEN_UPDATE_AFTER_OWNERSHIP_CHANGE: u8 = 26;
pub const TEST_ALLOW_WRITE_AFTER_OWNERSHIP_CHANGE_TO_CALLER: u8 = 27;
Expand Down
21 changes: 10 additions & 11 deletions programs/sbf/rust/invoke/src/processor.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -692,8 +692,8 @@ fn process_instruction(
invoked_instruction.accounts[1].is_writable = true;
invoke(&invoked_instruction, accounts)?;
}
TEST_FORBID_WRITE_AFTER_OWNERSHIP_CHANGE_TO_SYSTEM_ACCOUNT => {
msg!("TEST_FORBID_WRITE_AFTER_OWNERSHIP_CHANGE_TO_SYSTEM_ACCOUNT");
TEST_FORBID_WRITE_AFTER_OWNERSHIP_CHANGE_IN_CALLEE => {
msg!("TEST_FORBID_WRITE_AFTER_OWNERSHIP_CHANGE_IN_CALLEE");
invoke(
&create_instruction(
*program_id,
Expand All @@ -702,7 +702,7 @@ fn process_instruction(
(accounts[ARGUMENT_INDEX].key, true, false),
],
vec![
TEST_FORBID_WRITE_AFTER_OWNERSHIP_CHANGE_TO_SYSTEM_ACCOUNT_NESTED,
TEST_FORBID_WRITE_AFTER_OWNERSHIP_CHANGE_IN_CALLEE_NESTED,
42,
42,
42,
Expand All @@ -712,23 +712,22 @@ fn process_instruction(
)
.unwrap();
let account = &accounts[ARGUMENT_INDEX];
// this should cause the tx to fail with ReadonlyDataModified since
// the system program now owns the account
// this should cause the tx to fail since the callee changed ownership
unsafe {
*account
.data
.borrow_mut()
.get_unchecked_mut(instruction_data[1] as usize) = 42
};
}
TEST_FORBID_WRITE_AFTER_OWNERSHIP_CHANGE_TO_SYSTEM_ACCOUNT_NESTED => {
msg!("TEST_FORBID_WRITE_AFTER_OWNERSHIP_CHANGE_TO_SYSTEM_ACCOUNT_NESTED");
TEST_FORBID_WRITE_AFTER_OWNERSHIP_CHANGE_IN_CALLEE_NESTED => {
msg!("TEST_FORBID_WRITE_AFTER_OWNERSHIP_CHANGE_IN_CALLEE_NESTED");
let account = &accounts[ARGUMENT_INDEX];
account.data.borrow_mut().fill(0);
account.assign(&system_program::id());
}
TEST_FORBID_WRITE_AFTER_OWNERSHIP_CHANGE_TO_OTHER_PROGRAM => {
msg!("TEST_FORBID_WRITE_AFTER_OWNERSHIP_CHANGE_TO_OTHER_PROGRAM");
TEST_FORBID_WRITE_AFTER_OWNERSHIP_CHANGE_IN_CALLER => {
msg!("TEST_FORBID_WRITE_AFTER_OWNERSHIP_CHANGE_IN_CALLER");
let account = &accounts[ARGUMENT_INDEX];
let invoked_program_id = accounts[INVOKED_PROGRAM_INDEX].key;
account.data.borrow_mut().fill(0);
Expand All @@ -745,8 +744,8 @@ fn process_instruction(
accounts,
)
.unwrap();
// this should cause the tx to fail with ReadonlyDataModified since
// invoked_program_id now owns the account
// this should cause the tx to failsince invoked_program_id now owns
// the account
unsafe {
*account
.data
Expand Down
102 changes: 25 additions & 77 deletions programs/sbf/tests/programs.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4000,89 +4000,37 @@ fn test_cpi_account_ownership_writability() {
(2, 3), // first realloc byte
(2, 2 + MAX_PERMITTED_DATA_INCREASE as u8), // last realloc byte
] {
bank.register_recent_blockhash(&Hash::new_unique());
let account = AccountSharedData::new(42, account_size, &invoke_program_id);
bank.store_account(&account_keypair.pubkey(), &account);
for instruction_id in [
TEST_FORBID_WRITE_AFTER_OWNERSHIP_CHANGE_IN_CALLEE,
TEST_FORBID_WRITE_AFTER_OWNERSHIP_CHANGE_IN_CALLER,
] {
bank.register_recent_blockhash(&Hash::new_unique());
let account = AccountSharedData::new(42, account_size, &invoke_program_id);
bank.store_account(&account_keypair.pubkey(), &account);

let instruction = Instruction::new_with_bytes(
invoke_program_id,
&[
TEST_FORBID_WRITE_AFTER_OWNERSHIP_CHANGE_TO_SYSTEM_ACCOUNT,
byte_index,
42,
42,
],
account_metas.clone(),
);
let result = bank_client.send_and_confirm_instruction(&mint_keypair, instruction);
if (byte_index as usize) < account_size || direct_mapping {
assert_eq!(
result.unwrap_err().unwrap(),
TransactionError::InstructionError(
0,
InstructionError::ExternalAccountDataModified
)
);
} else {
// without direct mapping, changes to the realloc padding
// outside the account length are ignored
assert!(result.is_ok(), "{result:?}");
}

let account = AccountSharedData::new(42, account_size, &invoke_program_id);
bank.store_account(&account_keypair.pubkey(), &account);
let instruction = Instruction::new_with_bytes(
invoke_program_id,
&[
TEST_FORBID_WRITE_AFTER_OWNERSHIP_CHANGE_TO_OTHER_PROGRAM,
byte_index,
42,
42,
],
account_metas.clone(),
);
let result = bank_client.send_and_confirm_instruction(&mint_keypair, instruction);
if (byte_index as usize) < account_size || direct_mapping {
assert_eq!(
result.unwrap_err().unwrap(),
TransactionError::InstructionError(
0,
InstructionError::ExternalAccountDataModified
)
let instruction = Instruction::new_with_bytes(
invoke_program_id,
&[instruction_id, byte_index, 42, 42],
account_metas.clone(),
);
} else {
// without direct mapping, changes to the realloc padding
// outside the account length are ignored
assert!(result.is_ok(), "{result:?}");
}

// assign the account to invoked_program_id, which will then assign to invoke_program_id
// let mut account = bank.get_account(&account_keypair.pubkey()).unwrap();
// account.set_owner(invoked_program_id);
// bank.store_account(&account_keypair.pubkey(), &account);
let account = AccountSharedData::new(42, account_size, &invoked_program_id);
bank.store_account(&account_keypair.pubkey(), &account);
let result = bank_client.send_and_confirm_instruction(&mint_keypair, instruction);

let instruction = Instruction::new_with_bytes(
invoke_program_id,
&[
TEST_ALLOW_WRITE_AFTER_OWNERSHIP_CHANGE_TO_CALLER,
byte_index,
42,
42,
],
account_metas.clone(),
);
let result = bank_client.send_and_confirm_instruction(&mint_keypair, instruction);
assert!(result.is_ok(), "{result:?}");
let account = bank.get_account(&account_keypair.pubkey()).unwrap();
let mut data = vec![0; account.data().len()];
if (byte_index as usize) < data.len() {
data[byte_index as usize] = 42;
if (byte_index as usize) < account_size || direct_mapping {
assert_eq!(
result.unwrap_err().unwrap(),
TransactionError::InstructionError(
0,
InstructionError::ExternalAccountDataModified
)
);
} else {
// without direct mapping, changes to the realloc padding
// outside the account length are ignored
assert!(result.is_ok(), "{result:?}");
}
}
assert_eq!(account.data(), data);
}

// Test that the CPI code that updates `ref_to_len_in_vm` fails if we
// make it write to an invalid location. This is the first variant which
// correctly triggers ExternalAccountDataModified when direct mapping is
Expand Down

0 comments on commit 821cb14

Please sign in to comment.