Refactor instruction compilation and update message account key ordering (#23729)

* Refactor: Make instruction compilation usable for other message versions

* apply trents feedback

* Fix tests

* Fix bpf compatiblity
This commit is contained in:
Justin Starry
2022-03-21 20:53:32 +08:00
committed by GitHub
parent a1a29b0b86
commit 15357480ec
11 changed files with 1212 additions and 972 deletions

View File

@ -1,6 +1,10 @@
use {
crate::{message::v0::LoadedAddresses, pubkey::Pubkey},
std::ops::Index,
crate::{
instruction::{CompiledInstruction, Instruction},
message::v0::LoadedAddresses,
pubkey::Pubkey,
},
std::{collections::BTreeMap, ops::Index},
};
/// Collection of static and dynamically loaded keys used to load accounts
@ -75,6 +79,33 @@ impl<'a> AccountKeys<'a> {
pub fn iter(&self) -> impl Iterator<Item = &'a Pubkey> {
self.key_segment_iter().flatten()
}
/// Compile instructions using the order of account keys to determine
/// compiled instruction account indexes.
pub fn compile_instructions(&self, instructions: &[Instruction]) -> Vec<CompiledInstruction> {
let account_index_map: BTreeMap<&Pubkey, u8> = BTreeMap::from_iter(
self.iter()
.enumerate()
.map(|(index, key)| (key, index as u8)),
);
instructions
.iter()
.map(|ix| {
let accounts: Vec<u8> = ix
.accounts
.iter()
.map(|account_meta| *account_index_map.get(&account_meta.pubkey).unwrap())
.collect();
CompiledInstruction {
program_id_index: *account_index_map.get(&ix.program_id).unwrap(),
data: ix.data.clone(),
accounts,
}
})
.collect()
}
}
#[cfg(test)]

View File

@ -0,0 +1,231 @@
use {
crate::{instruction::Instruction, message::MessageHeader, pubkey::Pubkey},
std::collections::BTreeMap,
};
/// A helper struct to collect pubkeys compiled for a set of instructions
#[derive(Default, Debug, PartialEq, Eq)]
pub(crate) struct CompiledKeys {
writable_signer_keys: Vec<Pubkey>,
readonly_signer_keys: Vec<Pubkey>,
writable_non_signer_keys: Vec<Pubkey>,
readonly_non_signer_keys: Vec<Pubkey>,
}
#[derive(Default, Debug)]
struct CompiledKeyMeta {
is_signer: bool,
is_writable: bool,
}
impl CompiledKeys {
/// Compiles the pubkeys referenced by a list of instructions and organizes by
/// signer/non-signer and writable/readonly.
pub(crate) fn compile(instructions: &[Instruction], payer: Option<Pubkey>) -> Self {
let mut key_meta_map = BTreeMap::<&Pubkey, CompiledKeyMeta>::new();
for ix in instructions {
key_meta_map.entry(&ix.program_id).or_default();
for account_meta in &ix.accounts {
let meta = key_meta_map.entry(&account_meta.pubkey).or_default();
meta.is_signer |= account_meta.is_signer;
meta.is_writable |= account_meta.is_writable;
}
}
if let Some(payer) = &payer {
key_meta_map.remove(payer);
}
let writable_signer_keys: Vec<Pubkey> = payer
.into_iter()
.chain(
key_meta_map
.iter()
.filter_map(|(key, meta)| (meta.is_signer && meta.is_writable).then(|| **key)),
)
.collect();
let readonly_signer_keys = key_meta_map
.iter()
.filter_map(|(key, meta)| (meta.is_signer && !meta.is_writable).then(|| **key))
.collect();
let writable_non_signer_keys = key_meta_map
.iter()
.filter_map(|(key, meta)| (!meta.is_signer && meta.is_writable).then(|| **key))
.collect();
let readonly_non_signer_keys = key_meta_map
.iter()
.filter_map(|(key, meta)| (!meta.is_signer && !meta.is_writable).then(|| **key))
.collect();
CompiledKeys {
writable_signer_keys,
readonly_signer_keys,
writable_non_signer_keys,
readonly_non_signer_keys,
}
}
pub(crate) fn try_into_message_components(self) -> Option<(MessageHeader, Vec<Pubkey>)> {
let header = MessageHeader {
num_required_signatures: u8::try_from(
self.writable_signer_keys
.len()
.checked_add(self.readonly_signer_keys.len())?,
)
.ok()?,
num_readonly_signed_accounts: u8::try_from(self.readonly_signer_keys.len()).ok()?,
num_readonly_unsigned_accounts: u8::try_from(self.readonly_non_signer_keys.len())
.ok()?,
};
let static_account_keys = std::iter::empty()
.chain(self.writable_signer_keys)
.chain(self.readonly_signer_keys)
.chain(self.writable_non_signer_keys)
.chain(self.readonly_non_signer_keys)
.collect();
Some((header, static_account_keys))
}
}
#[cfg(test)]
mod tests {
use {super::*, crate::instruction::AccountMeta};
#[test]
fn test_compile_with_dups() {
let program_id = Pubkey::new_unique();
let id0 = Pubkey::new_unique();
let id1 = Pubkey::new_unique();
let id2 = Pubkey::new_unique();
let keys = CompiledKeys::compile(
&[Instruction::new_with_bincode(
program_id,
&0,
vec![
AccountMeta::new(id0, true),
AccountMeta::new_readonly(id1, true),
AccountMeta::new(id2, false),
// duplicate the account inputs
AccountMeta::new(id0, true),
AccountMeta::new_readonly(id1, true),
AccountMeta::new(id2, false),
],
)],
None,
);
assert_eq!(
keys,
CompiledKeys {
writable_signer_keys: vec![id0],
readonly_signer_keys: vec![id1],
writable_non_signer_keys: vec![id2],
readonly_non_signer_keys: vec![program_id],
}
);
}
#[test]
fn test_compile_with_dup_payer() {
let program_id = Pubkey::new_unique();
let payer = Pubkey::new_unique();
let keys = CompiledKeys::compile(
&[Instruction::new_with_bincode(
program_id,
&0,
vec![AccountMeta::new_readonly(payer, false)],
)],
Some(payer),
);
assert_eq!(
keys,
CompiledKeys {
writable_signer_keys: vec![payer],
readonly_non_signer_keys: vec![program_id],
..CompiledKeys::default()
}
);
}
#[test]
fn test_compile_with_dup_signer_mismatch() {
let program_id = Pubkey::new_unique();
let id0 = Pubkey::new_unique();
let keys = CompiledKeys::compile(
&[Instruction::new_with_bincode(
program_id,
&0,
vec![AccountMeta::new(id0, false), AccountMeta::new(id0, true)],
)],
None,
);
// Ensure the dup writable key is a signer
assert_eq!(
keys,
CompiledKeys {
writable_signer_keys: vec![id0],
readonly_non_signer_keys: vec![program_id],
..CompiledKeys::default()
}
);
}
#[test]
fn test_compile_with_dup_signer_writable_mismatch() {
let program_id = Pubkey::new_unique();
let id0 = Pubkey::new_unique();
let keys = CompiledKeys::compile(
&[Instruction::new_with_bincode(
program_id,
&0,
vec![
AccountMeta::new_readonly(id0, true),
AccountMeta::new(id0, true),
],
)],
None,
);
// Ensure the dup signer key is writable
assert_eq!(
keys,
CompiledKeys {
writable_signer_keys: vec![id0],
readonly_non_signer_keys: vec![program_id],
..CompiledKeys::default()
}
);
}
#[test]
fn test_compile_with_dup_nonsigner_writable_mismatch() {
let program_id = Pubkey::new_unique();
let id0 = Pubkey::new_unique();
let keys = CompiledKeys::compile(
&[
Instruction::new_with_bincode(
program_id,
&0,
vec![
AccountMeta::new_readonly(id0, false),
AccountMeta::new(id0, false),
],
),
Instruction::new_with_bincode(program_id, &0, vec![AccountMeta::new(id0, false)]),
],
None,
);
// Ensure the dup nonsigner key is writable
assert_eq!(
keys,
CompiledKeys {
writable_non_signer_keys: vec![id0],
readonly_non_signer_keys: vec![program_id],
..CompiledKeys::default()
}
);
}
}

View File

@ -15,14 +15,14 @@ use {
crate::{
bpf_loader, bpf_loader_deprecated, bpf_loader_upgradeable,
hash::Hash,
instruction::{AccountMeta, CompiledInstruction, Instruction},
message::MessageHeader,
instruction::{CompiledInstruction, Instruction},
message::{CompiledKeys, MessageHeader},
pubkey::Pubkey,
sanitize::{Sanitize, SanitizeError},
short_vec, system_instruction, system_program, sysvar, wasm_bindgen,
},
lazy_static::lazy_static,
std::{collections::BTreeSet, convert::TryFrom, str::FromStr},
std::{convert::TryFrom, str::FromStr},
};
lazy_static! {
@ -66,113 +66,6 @@ fn compile_instructions(ixs: &[Instruction], keys: &[Pubkey]) -> Vec<CompiledIns
ixs.iter().map(|ix| compile_instruction(ix, keys)).collect()
}
/// A helper struct to collect pubkeys referenced by a set of instructions and read-only counts
#[derive(Debug, PartialEq, Eq)]
struct InstructionKeys {
pub signed_keys: Vec<Pubkey>,
pub unsigned_keys: Vec<Pubkey>,
pub num_readonly_signed_accounts: u8,
pub num_readonly_unsigned_accounts: u8,
}
impl InstructionKeys {
fn new(
signed_keys: Vec<Pubkey>,
unsigned_keys: Vec<Pubkey>,
num_readonly_signed_accounts: u8,
num_readonly_unsigned_accounts: u8,
) -> Self {
Self {
signed_keys,
unsigned_keys,
num_readonly_signed_accounts,
num_readonly_unsigned_accounts,
}
}
}
/// Return pubkeys referenced by all instructions, with the ones needing signatures first. If the
/// payer key is provided, it is always placed first in the list of signed keys. Read-only signed
/// accounts are placed last in the set of signed accounts. Read-only unsigned accounts,
/// including program ids, are placed last in the set. No duplicates and order is preserved.
fn get_keys(instructions: &[Instruction], payer: Option<&Pubkey>) -> InstructionKeys {
let programs: Vec<_> = get_program_ids(instructions)
.iter()
.map(|program_id| AccountMeta {
pubkey: *program_id,
is_signer: false,
is_writable: false,
})
.collect();
let mut keys_and_signed: Vec<_> = instructions
.iter()
.flat_map(|ix| ix.accounts.iter())
.collect();
keys_and_signed.extend(&programs);
keys_and_signed.sort_by(|x, y| {
y.is_signer
.cmp(&x.is_signer)
.then(y.is_writable.cmp(&x.is_writable))
});
let payer_account_meta;
if let Some(payer) = payer {
payer_account_meta = AccountMeta {
pubkey: *payer,
is_signer: true,
is_writable: true,
};
keys_and_signed.insert(0, &payer_account_meta);
}
let mut unique_metas: Vec<AccountMeta> = vec![];
for account_meta in keys_and_signed {
// Promote to writable if a later AccountMeta requires it
if let Some(x) = unique_metas
.iter_mut()
.find(|x| x.pubkey == account_meta.pubkey)
{
x.is_writable |= account_meta.is_writable;
continue;
}
unique_metas.push(account_meta.clone());
}
let mut signed_keys = vec![];
let mut unsigned_keys = vec![];
let mut num_readonly_signed_accounts = 0;
let mut num_readonly_unsigned_accounts = 0;
for account_meta in unique_metas {
if account_meta.is_signer {
signed_keys.push(account_meta.pubkey);
if !account_meta.is_writable {
num_readonly_signed_accounts += 1;
}
} else {
unsigned_keys.push(account_meta.pubkey);
if !account_meta.is_writable {
num_readonly_unsigned_accounts += 1;
}
}
}
InstructionKeys::new(
signed_keys,
unsigned_keys,
num_readonly_signed_accounts,
num_readonly_unsigned_accounts,
)
}
/// Return program ids referenced by all instructions. No duplicates and order is preserved.
fn get_program_ids(instructions: &[Instruction]) -> Vec<Pubkey> {
let mut set = BTreeSet::new();
instructions
.iter()
.map(|ix| ix.program_id)
.filter(|&program_id| set.insert(program_id))
.collect()
}
/// A Solana transaction message (legacy).
///
/// See the [`message`] module documentation for further description.
@ -397,20 +290,16 @@ impl Message {
payer: Option<&Pubkey>,
blockhash: &Hash,
) -> Self {
let InstructionKeys {
mut signed_keys,
unsigned_keys,
num_readonly_signed_accounts,
num_readonly_unsigned_accounts,
} = get_keys(instructions, payer);
let num_required_signatures = signed_keys.len() as u8;
signed_keys.extend(&unsigned_keys);
let instructions = compile_instructions(instructions, &signed_keys);
let compiled_keys = CompiledKeys::compile(instructions, payer.cloned());
let (header, account_keys) = compiled_keys
.try_into_message_components()
.expect("overflow when compiling message keys");
let instructions = compile_instructions(instructions, &account_keys);
Self::new_with_compiled_instructions(
num_required_signatures,
num_readonly_signed_accounts,
num_readonly_unsigned_accounts,
signed_keys,
header.num_required_signatures,
header.num_readonly_signed_accounts,
header.num_readonly_unsigned_accounts,
account_keys,
*blockhash,
instructions,
)
@ -713,16 +602,6 @@ mod tests {
std::collections::HashSet,
};
#[test]
fn test_message_unique_program_ids() {
let program_id0 = Pubkey::default();
let program_ids = get_program_ids(&[
Instruction::new_with_bincode(program_id0, &0, vec![]),
Instruction::new_with_bincode(program_id0, &0, vec![]),
]);
assert_eq!(program_ids, vec![program_id0]);
}
#[test]
fn test_builtin_program_keys() {
let keys: HashSet<Pubkey> = BUILTIN_PROGRAMS_KEYS.iter().copied().collect();
@ -744,174 +623,6 @@ mod tests {
);
}
#[test]
fn test_message_unique_program_ids_not_adjacent() {
let program_id0 = Pubkey::default();
let program_id1 = Pubkey::new_unique();
let program_ids = get_program_ids(&[
Instruction::new_with_bincode(program_id0, &0, vec![]),
Instruction::new_with_bincode(program_id1, &0, vec![]),
Instruction::new_with_bincode(program_id0, &0, vec![]),
]);
assert_eq!(program_ids, vec![program_id0, program_id1]);
}
#[test]
fn test_message_unique_program_ids_order_preserved() {
let program_id0 = Pubkey::new_unique();
let program_id1 = Pubkey::default(); // Key less than program_id0
let program_ids = get_program_ids(&[
Instruction::new_with_bincode(program_id0, &0, vec![]),
Instruction::new_with_bincode(program_id1, &0, vec![]),
Instruction::new_with_bincode(program_id0, &0, vec![]),
]);
assert_eq!(program_ids, vec![program_id0, program_id1]);
}
#[test]
fn test_message_unique_keys_both_signed() {
let program_id = Pubkey::default();
let id0 = Pubkey::default();
let keys = get_keys(
&[
Instruction::new_with_bincode(program_id, &0, vec![AccountMeta::new(id0, true)]),
Instruction::new_with_bincode(program_id, &0, vec![AccountMeta::new(id0, true)]),
],
None,
);
assert_eq!(keys, InstructionKeys::new(vec![id0], vec![], 0, 0));
}
#[test]
fn test_message_unique_keys_signed_and_payer() {
let program_id = Pubkey::default();
let id0 = Pubkey::default();
let keys = get_keys(
&[Instruction::new_with_bincode(
program_id,
&0,
vec![AccountMeta::new(id0, true)],
)],
Some(&id0),
);
assert_eq!(keys, InstructionKeys::new(vec![id0], vec![], 0, 0));
}
#[test]
fn test_message_unique_keys_unsigned_and_payer() {
let program_id = Pubkey::default();
let id0 = Pubkey::default();
let keys = get_keys(
&[Instruction::new_with_bincode(
program_id,
&0,
vec![AccountMeta::new(id0, false)],
)],
Some(&id0),
);
assert_eq!(keys, InstructionKeys::new(vec![id0], vec![], 0, 0));
}
#[test]
fn test_message_unique_keys_one_signed() {
let program_id = Pubkey::default();
let id0 = Pubkey::default();
let keys = get_keys(
&[
Instruction::new_with_bincode(program_id, &0, vec![AccountMeta::new(id0, false)]),
Instruction::new_with_bincode(program_id, &0, vec![AccountMeta::new(id0, true)]),
],
None,
);
assert_eq!(keys, InstructionKeys::new(vec![id0], vec![], 0, 0));
}
#[test]
fn test_message_unique_keys_one_readonly_signed() {
let program_id = Pubkey::default();
let id0 = Pubkey::default();
let keys = get_keys(
&[
Instruction::new_with_bincode(
program_id,
&0,
vec![AccountMeta::new_readonly(id0, true)],
),
Instruction::new_with_bincode(program_id, &0, vec![AccountMeta::new(id0, true)]),
],
None,
);
// Ensure the key is no longer readonly
assert_eq!(keys, InstructionKeys::new(vec![id0], vec![], 0, 0));
}
#[test]
fn test_message_unique_keys_one_readonly_unsigned() {
let program_id = Pubkey::default();
let id0 = Pubkey::default();
let keys = get_keys(
&[
Instruction::new_with_bincode(
program_id,
&0,
vec![AccountMeta::new_readonly(id0, false)],
),
Instruction::new_with_bincode(program_id, &0, vec![AccountMeta::new(id0, false)]),
],
None,
);
// Ensure the key is no longer readonly
assert_eq!(keys, InstructionKeys::new(vec![], vec![id0], 0, 0));
}
#[test]
fn test_message_unique_keys_order_preserved() {
let program_id = Pubkey::default();
let id0 = Pubkey::new_unique();
let id1 = Pubkey::default(); // Key less than id0
let keys = get_keys(
&[
Instruction::new_with_bincode(program_id, &0, vec![AccountMeta::new(id0, false)]),
Instruction::new_with_bincode(program_id, &0, vec![AccountMeta::new(id1, false)]),
],
None,
);
assert_eq!(keys, InstructionKeys::new(vec![], vec![id0, id1], 0, 0));
}
#[test]
fn test_message_unique_keys_not_adjacent() {
let program_id = Pubkey::default();
let id0 = Pubkey::default();
let id1 = Pubkey::new_unique();
let keys = get_keys(
&[
Instruction::new_with_bincode(program_id, &0, vec![AccountMeta::new(id0, false)]),
Instruction::new_with_bincode(program_id, &0, vec![AccountMeta::new(id1, false)]),
Instruction::new_with_bincode(program_id, &0, vec![AccountMeta::new(id0, true)]),
],
None,
);
assert_eq!(keys, InstructionKeys::new(vec![id0], vec![id1], 0, 0));
}
#[test]
fn test_message_signed_keys_first() {
let program_id = Pubkey::default();
let id0 = Pubkey::default();
let id1 = Pubkey::new_unique();
let keys = get_keys(
&[
Instruction::new_with_bincode(program_id, &0, vec![AccountMeta::new(id0, false)]),
Instruction::new_with_bincode(program_id, &0, vec![AccountMeta::new(id1, true)]),
],
None,
);
assert_eq!(keys, InstructionKeys::new(vec![id1], vec![id0], 0, 0));
}
#[test]
// Ensure there's a way to calculate the number of required signatures.
fn test_message_signed_keys_len() {
@ -926,36 +637,6 @@ mod tests {
assert_eq!(message.header.num_required_signatures, 1);
}
#[test]
fn test_message_readonly_keys_last() {
let program_id = Pubkey::default();
let id0 = Pubkey::default(); // Identical key/program_id should be de-duped
let id1 = Pubkey::new_unique();
let id2 = Pubkey::new_unique();
let id3 = Pubkey::new_unique();
let keys = get_keys(
&[
Instruction::new_with_bincode(
program_id,
&0,
vec![AccountMeta::new_readonly(id0, false)],
),
Instruction::new_with_bincode(
program_id,
&0,
vec![AccountMeta::new_readonly(id1, true)],
),
Instruction::new_with_bincode(program_id, &0, vec![AccountMeta::new(id2, false)]),
Instruction::new_with_bincode(program_id, &0, vec![AccountMeta::new(id3, true)]),
],
None,
);
assert_eq!(
keys,
InstructionKeys::new(vec![id3, id1], vec![id2, id0], 1, 1)
);
}
#[test]
fn test_message_kitchen_sink() {
let program_id0 = Pubkey::new_unique();
@ -1007,32 +688,6 @@ mod tests {
assert_eq!(message.header.num_required_signatures, 2);
}
#[test]
fn test_message_program_last() {
let program_id = Pubkey::default();
let id0 = Pubkey::new_unique();
let id1 = Pubkey::new_unique();
let keys = get_keys(
&[
Instruction::new_with_bincode(
program_id,
&0,
vec![AccountMeta::new_readonly(id0, false)],
),
Instruction::new_with_bincode(
program_id,
&0,
vec![AccountMeta::new_readonly(id1, true)],
),
],
None,
);
assert_eq!(
keys,
InstructionKeys::new(vec![id1], vec![id0, program_id], 1, 2)
);
}
#[test]
fn test_program_position() {
let program_id0 = Pubkey::default();
@ -1103,7 +758,7 @@ mod tests {
);
assert_eq!(
message.get_account_keys_by_lock_type(),
(vec![&id1, &id0], vec![&id3, &id2, &program_id])
(vec![&id1, &id0], vec![&id3, &program_id, &id2])
);
}
@ -1199,7 +854,7 @@ mod tests {
let message = Message::new(&instructions, Some(&id1));
assert_eq!(
message.hash(),
Hash::from_str("CXRH7GHLieaQZRUjH1mpnNnUZQtU4V4RpJpAFgy77i3z").unwrap()
Hash::from_str("7VWCF4quo2CcWQFNUayZiorxpiR5ix8YzLebrXKf3fMF").unwrap()
)
}
}

View File

@ -37,6 +37,7 @@
//! types continue to be exposed to Solana programs, for backwards compatibility
//! reasons.
mod compiled_keys;
pub mod legacy;
#[cfg(not(target_arch = "bpf"))]
@ -49,6 +50,7 @@ mod non_bpf_modules {
pub use {account_keys::*, sanitized::*, versions::*};
}
use compiled_keys::*;
pub use legacy::Message;
#[cfg(not(target_arch = "bpf"))]
pub use non_bpf_modules::*;

View File

@ -102,6 +102,11 @@ impl TransactionContext {
.ok_or(InstructionError::NotEnoughAccountKeys)
}
/// Returns the keys for the accounts loaded in this Transaction
pub fn get_keys_of_accounts(&self) -> &[Pubkey] {
&self.account_keys
}
/// Searches for an account by its key
pub fn get_account_at_index(
&self,