Add a SanitizedMessage type that caches writable accounts indexes (#27317)

* Add a SanitizedMessage type that caches writable accounts indexes

* Add is_writable_account_cache to both SanitizedMessage variants, cache is initialized in constructors
This commit is contained in:
Tao Zhu 2022-08-25 16:33:41 -05:00 committed by GitHub
parent b7b03cbb05
commit 5e71f339c9
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
5 changed files with 273 additions and 91 deletions

View File

@ -228,7 +228,7 @@ pub(crate) mod tests {
clock::Slot,
hash::Hash,
instruction::CompiledInstruction,
message::{Message, MessageHeader, SanitizedMessage},
message::{LegacyMessage, Message, MessageHeader, SanitizedMessage},
nonce::{self, state::DurableNonce},
nonce_account,
pubkey::Pubkey,
@ -372,7 +372,7 @@ pub(crate) mod tests {
durable_nonce_fee: Some(DurableNonceFee::from(
&NonceFull::from_partial(
rollback_partial,
&SanitizedMessage::Legacy(message),
&SanitizedMessage::Legacy(LegacyMessage::new(message)),
&[(pubkey, nonce_account)],
&rent_debits,
)

View File

@ -179,7 +179,7 @@ mod tests {
solana_sdk::{
account::{AccountSharedData, ReadableAccount},
instruction::{AccountMeta, Instruction, InstructionError},
message::{AccountKeys, Message},
message::{AccountKeys, LegacyMessage, Message},
native_loader::{self, create_loadable_account_for_test},
pubkey::Pubkey,
secp256k1_instruction::new_secp256k1_instruction,
@ -270,20 +270,21 @@ mod tests {
AccountMeta::new_readonly(readonly_pubkey, false),
];
let message = SanitizedMessage::Legacy(Message::new_with_compiled_instructions(
1,
0,
2,
account_keys.clone(),
Hash::default(),
AccountKeys::new(&account_keys, None).compile_instructions(&[
Instruction::new_with_bincode(
mock_system_program_id,
&MockSystemInstruction::Correct,
account_metas.clone(),
),
]),
));
let message =
SanitizedMessage::Legacy(LegacyMessage::new(Message::new_with_compiled_instructions(
1,
0,
2,
account_keys.clone(),
Hash::default(),
AccountKeys::new(&account_keys, None).compile_instructions(&[
Instruction::new_with_bincode(
mock_system_program_id,
&MockSystemInstruction::Correct,
account_metas.clone(),
),
]),
)));
let sysvar_cache = SysvarCache::default();
let result = MessageProcessor::process_message(
builtin_programs,
@ -320,20 +321,21 @@ mod tests {
0
);
let message = SanitizedMessage::Legacy(Message::new_with_compiled_instructions(
1,
0,
2,
account_keys.clone(),
Hash::default(),
AccountKeys::new(&account_keys, None).compile_instructions(&[
Instruction::new_with_bincode(
mock_system_program_id,
&MockSystemInstruction::TransferLamports { lamports: 50 },
account_metas.clone(),
),
]),
));
let message =
SanitizedMessage::Legacy(LegacyMessage::new(Message::new_with_compiled_instructions(
1,
0,
2,
account_keys.clone(),
Hash::default(),
AccountKeys::new(&account_keys, None).compile_instructions(&[
Instruction::new_with_bincode(
mock_system_program_id,
&MockSystemInstruction::TransferLamports { lamports: 50 },
account_metas.clone(),
),
]),
)));
let result = MessageProcessor::process_message(
builtin_programs,
&message,
@ -359,20 +361,21 @@ mod tests {
))
);
let message = SanitizedMessage::Legacy(Message::new_with_compiled_instructions(
1,
0,
2,
account_keys.clone(),
Hash::default(),
AccountKeys::new(&account_keys, None).compile_instructions(&[
Instruction::new_with_bincode(
mock_system_program_id,
&MockSystemInstruction::ChangeData { data: 50 },
account_metas,
),
]),
));
let message =
SanitizedMessage::Legacy(LegacyMessage::new(Message::new_with_compiled_instructions(
1,
0,
2,
account_keys.clone(),
Hash::default(),
AccountKeys::new(&account_keys, None).compile_instructions(&[
Instruction::new_with_bincode(
mock_system_program_id,
&MockSystemInstruction::ChangeData { data: 50 },
account_metas,
),
]),
)));
let result = MessageProcessor::process_message(
builtin_programs,
&message,
@ -501,14 +504,14 @@ mod tests {
];
// Try to borrow mut the same account
let message = SanitizedMessage::Legacy(Message::new(
let message = SanitizedMessage::Legacy(LegacyMessage::new(Message::new(
&[Instruction::new_with_bincode(
mock_program_id,
&MockSystemInstruction::BorrowFail,
account_metas.clone(),
)],
Some(transaction_context.get_key_of_account_at_index(0).unwrap()),
));
)));
let sysvar_cache = SysvarCache::default();
let result = MessageProcessor::process_message(
builtin_programs,
@ -536,14 +539,14 @@ mod tests {
);
// Try to borrow mut the same account in a safe way
let message = SanitizedMessage::Legacy(Message::new(
let message = SanitizedMessage::Legacy(LegacyMessage::new(Message::new(
&[Instruction::new_with_bincode(
mock_program_id,
&MockSystemInstruction::MultiBorrowMut,
account_metas.clone(),
)],
Some(transaction_context.get_key_of_account_at_index(0).unwrap()),
));
)));
let result = MessageProcessor::process_message(
builtin_programs,
&message,
@ -564,7 +567,7 @@ mod tests {
assert!(result.is_ok());
// Do work on the same transaction account but at different instruction accounts
let message = SanitizedMessage::Legacy(Message::new(
let message = SanitizedMessage::Legacy(LegacyMessage::new(Message::new(
&[Instruction::new_with_bincode(
mock_program_id,
&MockSystemInstruction::DoWork {
@ -574,7 +577,7 @@ mod tests {
account_metas,
)],
Some(transaction_context.get_key_of_account_at_index(0).unwrap()),
));
)));
let result = MessageProcessor::process_message(
builtin_programs,
&message,
@ -644,7 +647,7 @@ mod tests {
let mut transaction_context =
TransactionContext::new(accounts, Some(Rent::default()), 1, 2);
let message = SanitizedMessage::Legacy(Message::new(
let message = SanitizedMessage::Legacy(LegacyMessage::new(Message::new(
&[
new_secp256k1_instruction(
&libsecp256k1::SecretKey::random(&mut rand::thread_rng()),
@ -653,7 +656,7 @@ mod tests {
Instruction::new_with_bytes(mock_program_id, &[], vec![]),
],
None,
));
)));
let sysvar_cache = SysvarCache::default();
let result = MessageProcessor::process_message(
builtin_programs,

View File

@ -3,7 +3,7 @@ use {
hash::Hash,
instruction::CompiledInstruction,
message::{
legacy::Message as LegacyMessage,
legacy,
v0::{self, LoadedAddresses},
AccountKeys, MessageHeader,
},
@ -14,15 +14,61 @@ use {
solana_program::{system_instruction::SystemInstruction, system_program},
sysvar::instructions::{BorrowedAccountMeta, BorrowedInstruction},
},
std::convert::TryFrom,
std::{borrow::Cow, convert::TryFrom},
thiserror::Error,
};
#[derive(Debug, Clone)]
pub struct LegacyMessage<'a> {
/// Legacy message
pub message: Cow<'a, legacy::Message>,
/// List of boolean with same length as account_keys(), each boolean value indicates if
/// corresponding account key is writable or not.
pub is_writable_account_cache: Vec<bool>,
}
impl<'a> LegacyMessage<'a> {
pub fn new(message: legacy::Message) -> Self {
let is_writable_account_cache = message
.account_keys
.iter()
.enumerate()
.map(|(i, _key)| message.is_writable(i))
.collect::<Vec<_>>();
Self {
message: Cow::Owned(message),
is_writable_account_cache,
}
}
pub fn has_duplicates(&self) -> bool {
self.message.has_duplicates()
}
pub fn is_key_called_as_program(&self, key_index: usize) -> bool {
self.message.is_key_called_as_program(key_index)
}
/// Inspect all message keys for the bpf upgradeable loader
pub fn is_upgradeable_loader_present(&self) -> bool {
self.message.is_upgradeable_loader_present()
}
/// Returns the full list of account keys.
pub fn account_keys(&self) -> AccountKeys {
AccountKeys::new(&self.message.account_keys, None)
}
pub fn is_writable(&self, index: usize) -> bool {
*self.is_writable_account_cache.get(index).unwrap_or(&false)
}
}
/// Sanitized message of a transaction.
#[derive(Debug, Clone)]
pub enum SanitizedMessage {
/// Sanitized legacy message
Legacy(LegacyMessage),
Legacy(LegacyMessage<'static>),
/// Sanitized version #0 message with dynamically loaded addresses
V0(v0::LoadedMessage<'static>),
}
@ -47,11 +93,11 @@ impl From<SanitizeError> for SanitizeMessageError {
}
}
impl TryFrom<LegacyMessage> for SanitizedMessage {
impl TryFrom<legacy::Message> for SanitizedMessage {
type Error = SanitizeMessageError;
fn try_from(message: LegacyMessage) -> Result<Self, Self::Error> {
fn try_from(message: legacy::Message) -> Result<Self, Self::Error> {
message.sanitize()?;
Ok(Self::Legacy(message))
Ok(Self::Legacy(LegacyMessage::new(message)))
}
}
@ -68,15 +114,15 @@ impl SanitizedMessage {
/// readonly accounts
pub fn header(&self) -> &MessageHeader {
match self {
Self::Legacy(message) => &message.header,
Self::Legacy(legacy_message) => &legacy_message.message.header,
Self::V0(loaded_msg) => &loaded_msg.message.header,
}
}
/// Returns a legacy message if this sanitized message wraps one
pub fn legacy_message(&self) -> Option<&LegacyMessage> {
if let Self::Legacy(message) = &self {
Some(message)
pub fn legacy_message(&self) -> Option<&legacy::Message> {
if let Self::Legacy(legacy_message) = &self {
Some(&legacy_message.message)
} else {
None
}
@ -92,7 +138,7 @@ impl SanitizedMessage {
/// The hash of a recent block, used for timing out a transaction
pub fn recent_blockhash(&self) -> &Hash {
match self {
Self::Legacy(message) => &message.recent_blockhash,
Self::Legacy(legacy_message) => &legacy_message.message.recent_blockhash,
Self::V0(loaded_msg) => &loaded_msg.message.recent_blockhash,
}
}
@ -101,7 +147,7 @@ impl SanitizedMessage {
/// one atomic transaction if all succeed.
pub fn instructions(&self) -> &[CompiledInstruction] {
match self {
Self::Legacy(message) => &message.instructions,
Self::Legacy(legacy_message) => &legacy_message.message.instructions,
Self::V0(loaded_msg) => &loaded_msg.message.instructions,
}
}
@ -124,7 +170,7 @@ impl SanitizedMessage {
/// Returns the list of account keys that are loaded for this message.
pub fn account_keys(&self) -> AccountKeys {
match self {
Self::Legacy(message) => AccountKeys::new(&message.account_keys, None),
Self::Legacy(message) => message.account_keys(),
Self::V0(message) => message.account_keys(),
}
}
@ -275,9 +321,9 @@ mod tests {
#[test]
fn test_try_from_message() {
let legacy_message_with_no_signers = LegacyMessage {
let legacy_message_with_no_signers = legacy::Message {
account_keys: vec![Pubkey::new_unique()],
..LegacyMessage::default()
..legacy::Message::default()
};
assert_eq!(
@ -296,7 +342,7 @@ mod tests {
CompiledInstruction::new(2, &(), vec![0, 1]),
];
let message = SanitizedMessage::try_from(LegacyMessage::new_with_compiled_instructions(
let message = SanitizedMessage::try_from(legacy::Message::new_with_compiled_instructions(
1,
0,
2,
@ -320,14 +366,14 @@ mod tests {
let key4 = Pubkey::new_unique();
let key5 = Pubkey::new_unique();
let legacy_message = SanitizedMessage::try_from(LegacyMessage {
let legacy_message = SanitizedMessage::try_from(legacy::Message {
header: MessageHeader {
num_required_signatures: 2,
num_readonly_signed_accounts: 1,
num_readonly_unsigned_accounts: 1,
},
account_keys: vec![key0, key1, key2, key3],
..LegacyMessage::default()
..legacy::Message::default()
})
.unwrap();
@ -364,7 +410,7 @@ mod tests {
CompiledInstruction::new(3, &(), vec![0, 0]),
];
let message = SanitizedMessage::try_from(LegacyMessage::new_with_compiled_instructions(
let message = SanitizedMessage::try_from(legacy::Message::new_with_compiled_instructions(
2,
1,
2,
@ -391,4 +437,73 @@ mod tests {
HashSet::default()
);
}
#[test]
fn test_is_writable_account_cache() {
let key0 = Pubkey::new_unique();
let key1 = Pubkey::new_unique();
let key2 = Pubkey::new_unique();
let key3 = Pubkey::new_unique();
let key4 = Pubkey::new_unique();
let key5 = Pubkey::new_unique();
let legacy_message = SanitizedMessage::try_from(legacy::Message {
header: MessageHeader {
num_required_signatures: 2,
num_readonly_signed_accounts: 1,
num_readonly_unsigned_accounts: 1,
},
account_keys: vec![key0, key1, key2, key3],
..legacy::Message::default()
})
.unwrap();
match legacy_message {
SanitizedMessage::Legacy(message) => {
assert_eq!(
message.is_writable_account_cache.len(),
message.account_keys().len()
);
assert!(message.is_writable_account_cache.get(0).unwrap());
assert!(!message.is_writable_account_cache.get(1).unwrap());
assert!(message.is_writable_account_cache.get(2).unwrap());
assert!(!message.is_writable_account_cache.get(3).unwrap());
}
_ => {
panic!("Expect to be SanitizedMessage::LegacyMessage")
}
}
let v0_message = SanitizedMessage::V0(v0::LoadedMessage::new(
v0::Message {
header: MessageHeader {
num_required_signatures: 2,
num_readonly_signed_accounts: 1,
num_readonly_unsigned_accounts: 1,
},
account_keys: vec![key0, key1, key2, key3],
..v0::Message::default()
},
LoadedAddresses {
writable: vec![key4],
readonly: vec![key5],
},
));
match v0_message {
SanitizedMessage::V0(message) => {
assert_eq!(
message.is_writable_account_cache.len(),
message.account_keys().len()
);
assert!(message.is_writable_account_cache.get(0).unwrap());
assert!(!message.is_writable_account_cache.get(1).unwrap());
assert!(message.is_writable_account_cache.get(2).unwrap());
assert!(!message.is_writable_account_cache.get(3).unwrap());
assert!(message.is_writable_account_cache.get(4).unwrap());
assert!(!message.is_writable_account_cache.get(5).unwrap());
}
_ => {
panic!("Expect to be SanitizedMessage::V0")
}
}
}
}

View File

@ -14,6 +14,9 @@ pub struct LoadedMessage<'a> {
pub message: Cow<'a, v0::Message>,
/// Addresses loaded with on-chain address lookup tables
pub loaded_addresses: Cow<'a, LoadedAddresses>,
/// List of boolean with same length as account_keys(), each boolean value indicates if
/// corresponding account key is writable or not.
pub is_writable_account_cache: Vec<bool>,
}
/// Collection of addresses loaded from on-chain lookup tables, split
@ -53,17 +56,36 @@ impl LoadedAddresses {
impl<'a> LoadedMessage<'a> {
pub fn new(message: v0::Message, loaded_addresses: LoadedAddresses) -> Self {
Self {
let mut loaded_message = Self {
message: Cow::Owned(message),
loaded_addresses: Cow::Owned(loaded_addresses),
}
is_writable_account_cache: Vec::default(),
};
loaded_message.set_is_writable_account_cache();
loaded_message
}
pub fn new_borrowed(message: &'a v0::Message, loaded_addresses: &'a LoadedAddresses) -> Self {
Self {
let mut loaded_message = Self {
message: Cow::Borrowed(message),
loaded_addresses: Cow::Borrowed(loaded_addresses),
}
is_writable_account_cache: Vec::default(),
};
loaded_message.set_is_writable_account_cache();
loaded_message
}
fn set_is_writable_account_cache(&mut self) {
let is_writable_account_cache = self
.account_keys()
.iter()
.enumerate()
.map(|(i, _key)| self.is_writable_internal(i))
.collect::<Vec<_>>();
let _ = std::mem::replace(
&mut self.is_writable_account_cache,
is_writable_account_cache,
);
}
/// Returns the full list of static and dynamic account keys that are loaded for this message.
@ -105,7 +127,7 @@ impl<'a> LoadedMessage<'a> {
}
/// Returns true if the account at the specified index was loaded as writable
pub fn is_writable(&self, key_index: usize) -> bool {
fn is_writable_internal(&self, key_index: usize) -> bool {
if self.is_writable_index(key_index) {
if let Some(key) = self.account_keys().get(key_index) {
return !(is_builtin_key_or_sysvar(key) || self.demote_program_id(key_index));
@ -114,6 +136,13 @@ impl<'a> LoadedMessage<'a> {
false
}
pub fn is_writable(&self, key_index: usize) -> bool {
*self
.is_writable_account_cache
.get(key_index)
.unwrap_or(&false)
}
pub fn is_signer(&self, i: usize) -> bool {
i < self.message.header.num_required_signatures as usize
}
@ -227,15 +256,45 @@ mod tests {
#[test]
fn test_is_writable() {
let mut message = check_test_loaded_message().0;
solana_logger::setup();
let create_message_with_keys = |keys: Vec<Pubkey>| {
LoadedMessage::new(
v0::Message {
header: MessageHeader {
num_required_signatures: 1,
num_readonly_signed_accounts: 0,
num_readonly_unsigned_accounts: 1,
},
account_keys: keys[..2].to_vec(),
..v0::Message::default()
},
LoadedAddresses {
writable: keys[2..=2].to_vec(),
readonly: keys[3..].to_vec(),
},
)
};
message.message.to_mut().account_keys[0] = sysvar::clock::id();
assert!(message.is_writable_index(0));
assert!(!message.is_writable(0));
let key0 = Pubkey::new_unique();
let key1 = Pubkey::new_unique();
let key2 = Pubkey::new_unique();
{
let message = create_message_with_keys(vec![sysvar::clock::id(), key0, key1, key2]);
assert!(message.is_writable_index(0));
assert!(!message.is_writable(0));
}
message.message.to_mut().account_keys[0] = system_program::id();
assert!(message.is_writable_index(0));
assert!(!message.is_writable(0));
{
let message = create_message_with_keys(vec![system_program::id(), key0, key1, key2]);
assert!(message.is_writable_index(0));
assert!(!message.is_writable(0));
}
{
let message = create_message_with_keys(vec![key0, key1, system_program::id(), key2]);
assert!(message.is_writable_index(2));
assert!(!message.is_writable(2));
}
}
#[test]

View File

@ -5,8 +5,9 @@ use {
crate::{
hash::Hash,
message::{
legacy,
v0::{self, LoadedAddresses, MessageAddressTableLookup},
SanitizedMessage, VersionedMessage,
LegacyMessage, SanitizedMessage, VersionedMessage,
},
precompiles::verify_if_precompile,
pubkey::Pubkey,
@ -87,7 +88,9 @@ impl SanitizedTransaction {
let signatures = tx.signatures;
let SanitizedVersionedMessage { message } = tx.message;
let message = match message {
VersionedMessage::Legacy(message) => SanitizedMessage::Legacy(message),
VersionedMessage::Legacy(message) => {
SanitizedMessage::Legacy(LegacyMessage::new(message))
}
VersionedMessage::V0(message) => {
let loaded_addresses =
address_loader.load_addresses(&message.address_table_lookups)?;
@ -122,7 +125,9 @@ impl SanitizedTransaction {
let signatures = tx.signatures;
let message = match tx.message {
VersionedMessage::Legacy(message) => SanitizedMessage::Legacy(message),
VersionedMessage::Legacy(message) => {
SanitizedMessage::Legacy(LegacyMessage::new(message))
}
VersionedMessage::V0(message) => {
let loaded_addresses =
address_loader.load_addresses(&message.address_table_lookups)?;
@ -149,7 +154,7 @@ impl SanitizedTransaction {
Ok(Self {
message_hash: tx.message.hash(),
message: SanitizedMessage::Legacy(tx.message),
message: SanitizedMessage::Legacy(LegacyMessage::new(tx.message)),
is_simple_vote_tx: false,
signatures: tx.signatures,
})
@ -200,9 +205,9 @@ impl SanitizedTransaction {
signatures,
message: VersionedMessage::V0(v0::Message::clone(&sanitized_msg.message)),
},
SanitizedMessage::Legacy(message) => VersionedTransaction {
SanitizedMessage::Legacy(legacy_message) => VersionedTransaction {
signatures,
message: VersionedMessage::Legacy(message.clone()),
message: VersionedMessage::Legacy(legacy::Message::clone(&legacy_message.message)),
},
}
}
@ -260,7 +265,7 @@ impl SanitizedTransaction {
/// Return the serialized message data to sign.
fn message_data(&self) -> Vec<u8> {
match &self.message {
SanitizedMessage::Legacy(message) => message.serialize(),
SanitizedMessage::Legacy(legacy_message) => legacy_message.message.serialize(),
SanitizedMessage::V0(loaded_msg) => loaded_msg.message.serialize(),
}
}