From 5e71f339c94f4ab2b2f20a6e7f5fc42c0661952f Mon Sep 17 00:00:00 2001 From: Tao Zhu <82401714+taozhu-chicago@users.noreply.github.com> Date: Thu, 25 Aug 2022 16:33:41 -0500 Subject: [PATCH] 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 --- rpc/src/transaction_status_service.rs | 4 +- runtime/src/message_processor.rs | 105 ++++++------ sdk/program/src/message/sanitized.rs | 153 +++++++++++++++--- sdk/program/src/message/versions/v0/loaded.rs | 83 ++++++++-- sdk/src/transaction/sanitized.rs | 19 ++- 5 files changed, 273 insertions(+), 91 deletions(-) diff --git a/rpc/src/transaction_status_service.rs b/rpc/src/transaction_status_service.rs index 4d393a0126..d0a3fde0cb 100644 --- a/rpc/src/transaction_status_service.rs +++ b/rpc/src/transaction_status_service.rs @@ -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, ) diff --git a/runtime/src/message_processor.rs b/runtime/src/message_processor.rs index c95bfc0b72..e93bdeb3ac 100644 --- a/runtime/src/message_processor.rs +++ b/runtime/src/message_processor.rs @@ -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, diff --git a/sdk/program/src/message/sanitized.rs b/sdk/program/src/message/sanitized.rs index 5fc64cd05f..52e75611db 100644 --- a/sdk/program/src/message/sanitized.rs +++ b/sdk/program/src/message/sanitized.rs @@ -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, +} + +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::>(); + 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 for SanitizeMessageError { } } -impl TryFrom for SanitizedMessage { +impl TryFrom for SanitizedMessage { type Error = SanitizeMessageError; - fn try_from(message: LegacyMessage) -> Result { + fn try_from(message: legacy::Message) -> Result { 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") + } + } + } } diff --git a/sdk/program/src/message/versions/v0/loaded.rs b/sdk/program/src/message/versions/v0/loaded.rs index c1a7adb260..dc3a12f661 100644 --- a/sdk/program/src/message/versions/v0/loaded.rs +++ b/sdk/program/src/message/versions/v0/loaded.rs @@ -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, } /// 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::>(); + 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| { + 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] diff --git a/sdk/src/transaction/sanitized.rs b/sdk/src/transaction/sanitized.rs index baedf8c08e..7d1b34a4ab 100644 --- a/sdk/src/transaction/sanitized.rs +++ b/sdk/src/transaction/sanitized.rs @@ -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 { match &self.message { - SanitizedMessage::Legacy(message) => message.serialize(), + SanitizedMessage::Legacy(legacy_message) => legacy_message.message.serialize(), SanitizedMessage::V0(loaded_msg) => loaded_msg.message.serialize(), } }