From b588d63eaf0e51a2877f969bc52f7f109a6d1cce Mon Sep 17 00:00:00 2001 From: Vladimir Guguiev <1524432+vovacodes@users.noreply.github.com> Date: Wed, 18 Oct 2023 13:58:27 +0200 Subject: [PATCH] feat(sdk-rs): add VaultTransactionMessageExt --- sdk/rs/src/client.rs | 21 ++- sdk/rs/src/lib.rs | 1 + .../compiled_keys.rs | 176 ++++++++++++++++++ sdk/rs/src/vault_transaction_message/mod.rs | 72 +++++++ 4 files changed, 261 insertions(+), 9 deletions(-) create mode 100644 sdk/rs/src/vault_transaction_message/compiled_keys.rs create mode 100644 sdk/rs/src/vault_transaction_message/mod.rs diff --git a/sdk/rs/src/client.rs b/sdk/rs/src/client.rs index bb8eb8c4..6194ae7e 100644 --- a/sdk/rs/src/client.rs +++ b/sdk/rs/src/client.rs @@ -313,22 +313,25 @@ pub fn spending_limit_use( /// ``` /// use squads_multisig::anchor_lang::AnchorSerialize; /// use squads_multisig::solana_program::pubkey::Pubkey; -/// use squads_multisig::solana_program::system_program; +/// use squads_multisig::solana_program::{system_instruction, system_program}; /// use squads_multisig::client::{ /// VaultTransactionCreateAccounts, /// VaultTransactionCreateArgs, /// vault_transaction_create, /// }; +/// use squads_multisig::pda::get_vault_pda; +/// use squads_multisig::vault_transaction_message::VaultTransactionMessageExt; /// use squads_multisig_program::TransactionMessage; /// -/// let message = TransactionMessage { -/// num_signers: 1, -/// num_writable_signers: 1, -/// num_writable_non_signers: 2, -/// account_keys: vec![].into(), -/// instructions: vec![].into(), -/// address_table_lookups: vec![].into(), -/// }.try_to_vec().unwrap(); +/// // Default vault (index 0). +/// let vault_pda = get_vault_pda(&Pubkey::new_unique(), 0, None).0; +/// +/// // Create a vault transaction message that includes 1 instruction - SOL transfer from the default vault. +/// let message = TransactionMessage::try_compile( +/// &vault_pda, +/// &[system_instruction::transfer(&vault_pda, &Pubkey::new_unique(), 1_000_000)], +/// &[] +/// ).unwrap().try_to_vec().unwrap(); /// /// let ix = vault_transaction_create( /// VaultTransactionCreateAccounts { diff --git a/sdk/rs/src/lib.rs b/sdk/rs/src/lib.rs index b27e0675..13cc58f0 100644 --- a/sdk/rs/src/lib.rs +++ b/sdk/rs/src/lib.rs @@ -7,6 +7,7 @@ pub use squads_multisig_program::anchor_lang::solana_program; pub mod client; pub mod pda; +pub mod vault_transaction_message; pub mod error { use thiserror::Error; diff --git a/sdk/rs/src/vault_transaction_message/compiled_keys.rs b/sdk/rs/src/vault_transaction_message/compiled_keys.rs new file mode 100644 index 00000000..0af03306 --- /dev/null +++ b/sdk/rs/src/vault_transaction_message/compiled_keys.rs @@ -0,0 +1,176 @@ +use std::collections::BTreeMap; + +use crate::solana_program::address_lookup_table_account::AddressLookupTableAccount; +use crate::solana_program::instruction::Instruction; +use crate::solana_program::message::v0::{LoadedAddresses, MessageAddressTableLookup}; +use crate::solana_program::message::{CompileError, MessageHeader}; + +use crate::solana_program::pubkey::Pubkey; + +/// A helper struct to collect pubkeys compiled for a set of instructions +/// +/// NOTE: The only difference between this and the original implementation from `solana_program` is that we don't mark the instruction programIds as invoked. +// /// It makes sense to do because the instructions will be called via CPI, so the programIds can come from Address Lookup Tables. +// /// This allows to compress the message size and avoid hitting the tx size limit during `vault_transaction_create` instruction calls. +#[derive(Default, Debug, Clone, PartialEq, Eq)] +pub(crate) struct CompiledKeys { + payer: Option, + key_meta_map: BTreeMap, +} + +#[derive(Default, Debug, Clone, PartialEq, Eq)] +struct CompiledKeyMeta { + is_signer: bool, + is_writable: bool, + is_invoked: 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) -> Self { + let mut key_meta_map = BTreeMap::::new(); + for ix in instructions { + let mut meta = key_meta_map.entry(ix.program_id).or_default(); + // NOTE: This is the only difference from the original. + // meta.is_invoked = true; + meta.is_invoked = false; + 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 { + let mut meta = key_meta_map.entry(*payer).or_default(); + meta.is_signer = true; + meta.is_writable = true; + } + Self { + payer, + key_meta_map, + } + } + + pub(crate) fn try_into_message_components( + self, + ) -> Result<(MessageHeader, Vec), CompileError> { + let try_into_u8 = |num: usize| -> Result { + u8::try_from(num).map_err(|_| CompileError::AccountIndexOverflow) + }; + + let Self { + payer, + mut key_meta_map, + } = self; + + if let Some(payer) = &payer { + key_meta_map.remove_entry(payer); + } + + let writable_signer_keys: Vec = 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: Vec = key_meta_map + .iter() + .filter_map(|(key, meta)| (meta.is_signer && !meta.is_writable).then(|| *key)) + .collect(); + let writable_non_signer_keys: Vec = key_meta_map + .iter() + .filter_map(|(key, meta)| (!meta.is_signer && meta.is_writable).then(|| *key)) + .collect(); + let readonly_non_signer_keys: Vec = key_meta_map + .iter() + .filter_map(|(key, meta)| (!meta.is_signer && !meta.is_writable).then(|| *key)) + .collect(); + + let signers_len = writable_signer_keys + .len() + .saturating_add(readonly_signer_keys.len()); + + let header = MessageHeader { + num_required_signatures: try_into_u8(signers_len)?, + num_readonly_signed_accounts: try_into_u8(readonly_signer_keys.len())?, + num_readonly_unsigned_accounts: try_into_u8(readonly_non_signer_keys.len())?, + }; + + let static_account_keys = std::iter::empty() + .chain(writable_signer_keys) + .chain(readonly_signer_keys) + .chain(writable_non_signer_keys) + .chain(readonly_non_signer_keys) + .collect(); + + Ok((header, static_account_keys)) + } + + #[cfg(not(target_os = "solana"))] + pub(crate) fn try_extract_table_lookup( + &mut self, + lookup_table_account: &AddressLookupTableAccount, + ) -> Result, CompileError> { + let (writable_indexes, drained_writable_keys) = self + .try_drain_keys_found_in_lookup_table(&lookup_table_account.addresses, |meta| { + !meta.is_signer && !meta.is_invoked && meta.is_writable + })?; + let (readonly_indexes, drained_readonly_keys) = self + .try_drain_keys_found_in_lookup_table(&lookup_table_account.addresses, |meta| { + !meta.is_signer && !meta.is_invoked && !meta.is_writable + })?; + + // Don't extract lookup if no keys were found + if writable_indexes.is_empty() && readonly_indexes.is_empty() { + return Ok(None); + } + + Ok(Some(( + MessageAddressTableLookup { + account_key: lookup_table_account.key, + writable_indexes, + readonly_indexes, + }, + LoadedAddresses { + writable: drained_writable_keys, + readonly: drained_readonly_keys, + }, + ))) + } + + #[cfg(not(target_os = "solana"))] + fn try_drain_keys_found_in_lookup_table( + &mut self, + lookup_table_addresses: &[Pubkey], + key_meta_filter: impl Fn(&CompiledKeyMeta) -> bool, + ) -> Result<(Vec, Vec), CompileError> { + let mut lookup_table_indexes = Vec::new(); + let mut drained_keys = Vec::new(); + + for search_key in self + .key_meta_map + .iter() + .filter_map(|(key, meta)| key_meta_filter(meta).then(|| key)) + { + for (key_index, key) in lookup_table_addresses.iter().enumerate() { + if key == search_key { + let lookup_table_index = u8::try_from(key_index) + .map_err(|_| CompileError::AddressTableLookupIndexOverflow)?; + + lookup_table_indexes.push(lookup_table_index); + drained_keys.push(*search_key); + break; + } + } + } + + for key in &drained_keys { + self.key_meta_map.remove_entry(key); + } + + Ok((lookup_table_indexes, drained_keys)) + } +} diff --git a/sdk/rs/src/vault_transaction_message/mod.rs b/sdk/rs/src/vault_transaction_message/mod.rs new file mode 100644 index 00000000..39f8a2f0 --- /dev/null +++ b/sdk/rs/src/vault_transaction_message/mod.rs @@ -0,0 +1,72 @@ +use squads_multisig_program::{CompiledInstruction, MessageAddressTableLookup, TransactionMessage}; + +use crate::solana_program::address_lookup_table_account::AddressLookupTableAccount; +use crate::solana_program::instruction::Instruction; +use crate::solana_program::message::{AccountKeys, CompileError}; +use crate::solana_program::pubkey::Pubkey; +use crate::vault_transaction_message::compiled_keys::CompiledKeys; + +mod compiled_keys; + +pub trait VaultTransactionMessageExt { + /// This implementation is mostly a copy-paste from `solana_program::message::v0::Message::try_compile()`, + /// but it constructs a `TransactionMessage` meant to be passed to `vault_transaction_create`. + fn try_compile( + vault_key: &Pubkey, + instructions: &[Instruction], + address_lookup_table_accounts: &[AddressLookupTableAccount], + ) -> Result { + let mut compiled_keys = CompiledKeys::compile(instructions, Some(*vault_key)); + + let mut address_table_lookups = Vec::with_capacity(address_lookup_table_accounts.len()); + let mut loaded_addresses_list = Vec::with_capacity(address_lookup_table_accounts.len()); + for lookup_table_account in address_lookup_table_accounts { + if let Some((lookup, loaded_addresses)) = + compiled_keys.try_extract_table_lookup(lookup_table_account)? + { + address_table_lookups.push(lookup); + loaded_addresses_list.push(loaded_addresses); + } + } + + let (header, static_keys) = compiled_keys.try_into_message_components()?; + let dynamic_keys = loaded_addresses_list.into_iter().collect(); + let account_keys = AccountKeys::new(&static_keys, Some(&dynamic_keys)); + let instructions = account_keys.try_compile_instructions(instructions)?; + + let num_static_keys: u8 = static_keys + .len() + .try_into() + .map_err(|_| CompileError::AccountIndexOverflow)?; + + Ok(TransactionMessage { + num_signers: header.num_required_signatures, + num_writable_signers: header.num_required_signatures + - header.num_readonly_signed_accounts, + num_writable_non_signers: num_static_keys + - header.num_required_signatures + - header.num_readonly_unsigned_accounts, + account_keys: static_keys.into(), + instructions: instructions + .into_iter() + .map(|ix| CompiledInstruction { + program_id_index: ix.program_id_index, + account_indexes: ix.accounts.into(), + data: ix.data.into(), + }) + .collect::>() + .into(), + address_table_lookups: address_table_lookups + .into_iter() + .map(|lookup| MessageAddressTableLookup { + account_key: lookup.account_key, + writable_indexes: lookup.writable_indexes.into(), + readonly_indexes: lookup.readonly_indexes.into(), + }) + .collect::>() + .into(), + }) + } +} + +impl VaultTransactionMessageExt for TransactionMessage {}