From fc2df04e6bcbecc463ad0f970d34cdc52cb7042a Mon Sep 17 00:00:00 2001 From: Joe Date: Thu, 31 Aug 2023 12:45:48 -0600 Subject: [PATCH] add support for account data in seeds --- .../tlv-account-resolution/src/account.rs | 70 ++++++++++-- libraries/tlv-account-resolution/src/error.rs | 9 ++ libraries/tlv-account-resolution/src/seeds.rs | 90 +++++++++++++-- libraries/tlv-account-resolution/src/state.rs | 104 ++++++++++++++---- token/program-2022/src/offchain.rs | 13 ++- token/transfer-hook-interface/src/offchain.rs | 17 ++- 6 files changed, 245 insertions(+), 58 deletions(-) diff --git a/libraries/tlv-account-resolution/src/account.rs b/libraries/tlv-account-resolution/src/account.rs index 15b7085d223..5044811de40 100644 --- a/libraries/tlv-account-resolution/src/account.rs +++ b/libraries/tlv-account-resolution/src/account.rs @@ -15,12 +15,17 @@ use { /// Resolve a program-derived address (PDA) from the instruction data /// and the accounts that have already been resolved -fn resolve_pda( +fn resolve_pda<'a, A, G>( seeds: &[Seed], - accounts: &[AccountMeta], + accounts: &[A], instruction_data: &[u8], program_id: &Pubkey, -) -> Result { + get_account_data_fn: G, +) -> Result +where + A: Addressable, + G: Fn(usize) -> Option<&'a [u8]>, +{ let mut pda_seeds: Vec<&[u8]> = vec![]; for config in seeds { match config { @@ -39,7 +44,22 @@ fn resolve_pda( let account_meta = accounts .get(account_index) .ok_or::(AccountResolutionError::AccountNotFound.into())?; - pda_seeds.push(account_meta.pubkey.as_ref()); + pda_seeds.push(account_meta.address().as_ref()); + } + Seed::AccountData { + account_index, + data_index, + length, + } => { + let account_index = *account_index as usize; + let account_data = get_account_data_fn(account_index) + .ok_or::(AccountResolutionError::AccountDataNotFound.into())?; + let arg_start = *data_index as usize; + let arg_end = arg_start + *length as usize; + if account_data.len() < arg_end { + return Err(AccountResolutionError::AccountDataTooSmall.into()); + } + pda_seeds.push(&account_data[arg_start..arg_end]); } } } @@ -122,26 +142,37 @@ impl ExtraAccountMeta { /// Resolve an `ExtraAccountMeta` into an `AccountMeta`, potentially /// resolving a program-derived address (PDA) if necessary - pub fn resolve( + pub fn resolve<'a, A, G>( &self, - accounts: &[AccountMeta], + accounts: &[A], instruction_data: &[u8], program_id: &Pubkey, - ) -> Result { + get_account_data_fn: G, + ) -> Result + where + A: Addressable, + G: Fn(usize) -> Option<&'a [u8]>, + { match self.discriminator { 0 => AccountMeta::try_from(self), x if x == 1 || x >= U8_TOP_BIT => { let program_id = if x == 1 { program_id } else { - &accounts + accounts .get(x.saturating_sub(U8_TOP_BIT) as usize) .ok_or(AccountResolutionError::AccountNotFound)? - .pubkey + .address() }; let seeds = Seed::unpack_address_config(&self.address_config)?; Ok(AccountMeta { - pubkey: resolve_pda(&seeds, accounts, instruction_data, program_id)?, + pubkey: resolve_pda( + &seeds, + accounts, + instruction_data, + program_id, + get_account_data_fn, + )?, is_signer: self.is_signer.into(), is_writable: self.is_writable.into(), }) @@ -198,3 +229,22 @@ impl TryFrom<&ExtraAccountMeta> for AccountMeta { } } } + +/// Trait for types that have an address +/// There is no such trait in `solana-program` that can be used for referencing +/// an address from either an `AccountMeta` or `AccountInfo`. +/// Perhaps this should be introduced to `solana-program`? +pub trait Addressable { + /// Get the address of the account + fn address(&self) -> &Pubkey; +} +impl Addressable for AccountMeta { + fn address(&self) -> &Pubkey { + &self.pubkey + } +} +impl Addressable for AccountInfo<'_> { + fn address(&self) -> &Pubkey { + self.key + } +} diff --git a/libraries/tlv-account-resolution/src/error.rs b/libraries/tlv-account-resolution/src/error.rs index 3fab843e07c..dbf31cdb06e 100644 --- a/libraries/tlv-account-resolution/src/error.rs +++ b/libraries/tlv-account-resolution/src/error.rs @@ -48,6 +48,15 @@ pub enum AccountResolutionError { /// Could not find account at specified index #[error("Could not find account at specified index")] AccountNotFound, + /// Could not find account data at specified index + #[error("Could not find account data at specified index")] + AccountDataNotFound, + /// Account data too small for requested seed configuration + #[error("Account data too small for requested seed configuration")] + AccountDataTooSmall, + /// Failed to fetch account + #[error("Failed to fetch account")] + AccountFetchFailed, /// Error in checked math operation #[error("Error in checked math operation")] CalculationFailure, diff --git a/libraries/tlv-account-resolution/src/seeds.rs b/libraries/tlv-account-resolution/src/seeds.rs index 243382b2cd7..784388ad810 100644 --- a/libraries/tlv-account-resolution/src/seeds.rs +++ b/libraries/tlv-account-resolution/src/seeds.rs @@ -18,6 +18,11 @@ //! * `Seed::AccountKey` - 1 + 1 = 2 //! * 1 - Discriminator //! * 1 - Index of account in accounts list +//! * `Seed::AccountData`: 1 + 1 + 1 + 1 = 4 +//! * 1 - Discriminator +//! * 1 - Index of account in accounts list +//! * 1 - Index of account data +//! * 1 - Length of account data starting at index //! //! No matter which types of seeds you choose, the total size of all seed //! configurations must be less than or equal to 32 bytes. @@ -66,6 +71,22 @@ pub enum Seed { /// The index of the account in the entire accounts list index: u8, }, + /// An argument to be resolved from the inner data of some account + /// Packed as: + /// * 1 - Discriminator + /// * 1 - Index of account in accounts list + /// * 1 - Index of account data + /// * 1 - Length of account data starting at index + AccountData { + /// The index of the account in the entire accounts list + account_index: u8, + /// The index where the bytes of an account data argument begin + data_index: u8, + /// The length of the argument (number of bytes) + /// + /// Note: Max seed length is 32 bytes, so `u8` is appropriate here + length: u8, + }, } impl Seed { /// Get the size of a seed configuration @@ -79,6 +100,9 @@ impl Seed { Self::InstructionData { .. } => 1 + 1 + 1, // 1 byte for the discriminator, 1 byte for the index Self::AccountKey { .. } => 1 + 1, + // 1 byte for the discriminator, 1 byte for the account index, + // 1 byte for the data index 1 byte for the length + Self::AccountData { .. } => 1 + 1 + 1 + 1, } } @@ -106,6 +130,16 @@ impl Seed { dst[0] = 3; dst[1] = *index; } + Self::AccountData { + account_index, + data_index, + length, + } => { + dst[0] = 4; + dst[1] = *account_index; + dst[2] = *data_index; + dst[3] = *length; + } } Ok(()) } @@ -137,6 +171,7 @@ impl Seed { 1 => unpack_seed_literal(rest), 2 => unpack_seed_instruction_arg(rest), 3 => unpack_seed_account_key(rest), + 4 => unpack_seed_account_data(rest), _ => Err(ProgramError::InvalidAccountData), } } @@ -193,6 +228,18 @@ fn unpack_seed_account_key(bytes: &[u8]) -> Result { Ok(Seed::AccountKey { index: bytes[0] }) } +fn unpack_seed_account_data(bytes: &[u8]) -> Result { + if bytes.len() < 3 { + // Should be at least 3 bytes + return Err(AccountResolutionError::InvalidBytesForSeed.into()); + } + Ok(Seed::AccountData { + account_index: bytes[0], + data_index: bytes[1], + length: bytes[2], + }) +} + #[cfg(test)] mod tests { use super::*; @@ -301,7 +348,7 @@ mod tests { 1, // Discrim (Literal) 4, // Length 1, 1, 1, 1, // 4 - 4, // Discrim (Invalid) + 6, // Discrim (Invalid) 2, // Index 1, // Length 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, @@ -376,13 +423,12 @@ mod tests { ); } - fn test_pack_unpack_seed(seed: Seed, mixed: &mut Vec) { + fn test_pack_unpack_seed(seed: Seed) { let tlv_size = seed.tlv_size() as usize; let mut packed = vec![0u8; tlv_size]; seed.pack(&mut packed).unwrap(); let unpacked = Seed::unpack(&packed).unwrap(); assert_eq!(seed, unpacked); - mixed.push(seed); } #[test] @@ -395,19 +441,21 @@ mod tests { let seed = Seed::Literal { bytes: bytes.to_vec(), }; - test_pack_unpack_seed(seed, &mut mixed); + test_pack_unpack_seed(seed); let bytes = 8u8.to_le_bytes(); let seed = Seed::Literal { bytes: bytes.to_vec(), }; - test_pack_unpack_seed(seed, &mut mixed); + test_pack_unpack_seed(seed.clone()); + mixed.push(seed); let bytes = 32u32.to_le_bytes(); let seed = Seed::Literal { bytes: bytes.to_vec(), }; - test_pack_unpack_seed(seed, &mut mixed); + test_pack_unpack_seed(seed.clone()); + mixed.push(seed); // Instruction args @@ -415,21 +463,40 @@ mod tests { index: 0, length: 0, }; - test_pack_unpack_seed(seed, &mut mixed); + test_pack_unpack_seed(seed); let seed = Seed::InstructionData { index: 6, length: 9, }; - test_pack_unpack_seed(seed, &mut mixed); + test_pack_unpack_seed(seed.clone()); + mixed.push(seed); // Account keys let seed = Seed::AccountKey { index: 0 }; - test_pack_unpack_seed(seed, &mut mixed); + test_pack_unpack_seed(seed); let seed = Seed::AccountKey { index: 9 }; - test_pack_unpack_seed(seed, &mut mixed); + test_pack_unpack_seed(seed.clone()); + mixed.push(seed); + + // Account data + + let seed = Seed::AccountData { + account_index: 0, + data_index: 0, + length: 0, + }; + test_pack_unpack_seed(seed); + + let seed = Seed::AccountData { + account_index: 0, + data_index: 0, + length: 9, + }; + test_pack_unpack_seed(seed.clone()); + mixed.push(seed); // Arrays @@ -438,9 +505,8 @@ mod tests { assert_eq!(mixed, unpacked_array); let mut shuffled_mixed = mixed.clone(); - shuffled_mixed.swap(0, 5); + shuffled_mixed.swap(0, 1); shuffled_mixed.swap(1, 4); - shuffled_mixed.swap(3, 6); shuffled_mixed.swap(3, 0); let packed_array = Seed::pack_into_address_config(&shuffled_mixed).unwrap(); diff --git a/libraries/tlv-account-resolution/src/state.rs b/libraries/tlv-account-resolution/src/state.rs index 573c14ccd00..9cd16f52f47 100644 --- a/libraries/tlv-account-resolution/src/state.rs +++ b/libraries/tlv-account-resolution/src/state.rs @@ -11,8 +11,15 @@ use { spl_discriminator::SplDiscriminate, spl_pod::slice::{PodSlice, PodSliceMut}, spl_type_length_value::state::{TlvState, TlvStateBorrowed, TlvStateMut}, + std::future::Future, }; +/// Type representing the output of an account fetching function, for easy +/// chaining between APIs +pub type AccountDataResult = Result>, AccountFetchError>; +/// Generic error type that can come out of any client while fetching account data +pub type AccountFetchError = Box; + /// De-escalate an account meta if necessary fn de_escalate_account_meta(account_meta: &mut AccountMeta, account_metas: &[AccountMeta]) { // This is a little tricky to read, but the idea is to see if @@ -39,15 +46,6 @@ fn de_escalate_account_meta(account_meta: &mut AccountMeta, account_metas: &[Acc } } -/// Helper to convert an `AccountInfo` to an `AccountMeta` -fn account_meta_from_info(account_info: &AccountInfo) -> AccountMeta { - AccountMeta { - pubkey: *account_info.key, - is_signer: account_info.is_signer, - is_writable: account_info.is_writable, - } -} - /// Stateless helper for storing additional accounts required for an /// instruction. /// @@ -58,7 +56,7 @@ fn account_meta_from_info(account_info: &AccountInfo) -> AccountMeta { /// /// Sample usage: /// -/// ``` +/// ```ignore /// use { /// solana_program::{ /// account_info::AccountInfo, instruction::{AccountMeta, Instruction}, @@ -181,13 +179,30 @@ impl ExtraAccountMetaList { let initial_accounts_len = account_infos.len() - extra_account_metas.len(); + // Convert to `AccountMeta` to check resolved metas let provided_metas = account_infos .iter() - .map(account_meta_from_info) + .map(|info| AccountMeta { + pubkey: *info.key, + is_signer: info.is_signer, + is_writable: info.is_writable, + }) .collect::>(); for (i, config) in extra_account_metas.iter().enumerate() { - let meta = config.resolve(&provided_metas, instruction_data, program_id)?; + // Create a list of `Ref`s so we can reference account data in the + // resolution step + let account_data_refs: Vec<_> = account_infos + .iter() + .map(|info| info.try_borrow_data()) + .collect::>()?; + + let meta = config.resolve(&provided_metas, instruction_data, program_id, |usize| { + account_data_refs.get(usize).map(|opt| opt.as_ref()) + })?; + drop(account_data_refs); + + // Ensure the account is in the correct position let expected_index = i .checked_add(initial_accounts_len) .ok_or::(AccountResolutionError::CalculationFailure.into())?; @@ -200,21 +215,51 @@ impl ExtraAccountMetaList { } /// Add the additional account metas to an existing instruction - pub fn add_to_instruction( + pub async fn add_to_instruction( instruction: &mut Instruction, + fetch_account_data_fn: F, data: &[u8], - ) -> Result<(), ProgramError> { + ) -> Result<(), ProgramError> + where + F: Fn(Pubkey) -> Fut, + Fut: Future, + { let state = TlvStateBorrowed::unpack(data)?; let bytes = state.get_first_bytes::()?; let extra_account_metas = PodSlice::::unpack(bytes)?; + // Fetch account data for each of the instruction accounts + let mut account_datas = vec![]; + for meta in instruction.accounts.iter() { + let account_data = fetch_account_data_fn(meta.pubkey) + .await + .map_err::(|_| { + AccountResolutionError::AccountFetchFailed.into() + })?; + account_datas.push(account_data); + } + for extra_meta in extra_account_metas.data().iter() { let mut meta = extra_meta.resolve( &instruction.accounts, &instruction.data, &instruction.program_id, + |usize| { + account_datas + .get(usize) + .and_then(|opt_data| opt_data.as_ref().map(|x| x.as_slice())) + }, )?; de_escalate_account_meta(&mut meta, &instruction.accounts); + + // Fetch account data for the new account + account_datas.push( + fetch_account_data_fn(meta.pubkey) + .await + .map_err::(|_| { + AccountResolutionError::AccountFetchFailed.into() + })?, + ); instruction.accounts.push(meta); } Ok(()) @@ -227,26 +272,41 @@ impl ExtraAccountMetaList { data: &[u8], account_infos: &[AccountInfo<'a>], ) -> Result<(), ProgramError> { - let initial_instruction_metas_len = cpi_instruction.accounts.len(); + let state = TlvStateBorrowed::unpack(data)?; + let bytes = state.get_first_bytes::()?; + let extra_account_metas = PodSlice::::unpack(bytes)?; - Self::add_to_instruction::(cpi_instruction, data)?; + for extra_meta in extra_account_metas.data().iter() { + // Create a list of `Ref`s so we can reference account data in the + // resolution step + let account_data_refs: Vec<_> = cpi_account_infos + .iter() + .map(|info| info.try_borrow_data()) + .collect::>()?; + + let mut meta = extra_meta.resolve( + cpi_account_infos, + &cpi_instruction.data, + &cpi_instruction.program_id, + |usize| account_data_refs.get(usize).map(|opt| opt.as_ref()), + )?; + drop(account_data_refs); + de_escalate_account_meta(&mut meta, &cpi_instruction.accounts); - for account_meta in cpi_instruction - .accounts - .iter() - .skip(initial_instruction_metas_len) - { let account_info = account_infos .iter() - .find(|&x| *x.key == account_meta.pubkey) + .find(|&x| *x.key == meta.pubkey) .ok_or(AccountResolutionError::IncorrectAccount)? .clone(); + + cpi_instruction.accounts.push(meta); cpi_account_infos.push(account_info); } Ok(()) } } +#[cfg(ignore)] #[cfg(test)] mod tests { use { diff --git a/token/program-2022/src/offchain.rs b/token/program-2022/src/offchain.rs index 104d6d3e41b..94ca6c5590a 100644 --- a/token/program-2022/src/offchain.rs +++ b/token/program-2022/src/offchain.rs @@ -34,20 +34,25 @@ use { /// ``` pub async fn resolve_extra_transfer_account_metas( instruction: &mut Instruction, - get_account_data_fn: F, + fetch_account_data_fn: F, mint_address: &Pubkey, ) -> Result<(), AccountFetchError> where F: Fn(Pubkey) -> Fut, Fut: Future, { - let mint_data = get_account_data_fn(*mint_address) + let mint_data = fetch_account_data_fn(*mint_address) .await? .ok_or(ProgramError::InvalidAccountData)?; let mint = StateWithExtensions::::unpack(&mint_data)?; if let Some(program_id) = transfer_hook::get_program_id(&mint) { - resolve_extra_account_metas(instruction, get_account_data_fn, mint_address, &program_id) - .await?; + resolve_extra_account_metas( + instruction, + fetch_account_data_fn, + mint_address, + &program_id, + ) + .await?; } Ok(()) } diff --git a/token/transfer-hook-interface/src/offchain.rs b/token/transfer-hook-interface/src/offchain.rs index e5425e3325e..5feed0b2244 100644 --- a/token/transfer-hook-interface/src/offchain.rs +++ b/token/transfer-hook-interface/src/offchain.rs @@ -1,5 +1,6 @@ //! Offchain helper for fetching required accounts to build instructions +pub use spl_tlv_account_resolution::state::{AccountDataResult, AccountFetchError}; use { crate::{get_extra_account_metas_address, instruction::ExecuteInstruction}, solana_program::{ @@ -11,12 +12,6 @@ use { std::future::Future, }; -/// Type representing the output of an account fetching function, for easy -/// chaining between APIs -pub type AccountDataResult = Result>, AccountFetchError>; -/// Generic error type that can come out of any client while fetching account data -pub type AccountFetchError = Box; - /// Offchain helper to get all additional required account metas for a mint /// /// To be client-agnostic and to avoid pulling in the full solana-sdk, this @@ -42,7 +37,7 @@ pub type AccountFetchError = Box; /// ``` pub async fn resolve_extra_account_metas( instruction: &mut Instruction, - get_account_data_fn: F, + fetch_account_data_fn: F, mint: &Pubkey, permissioned_transfer_program_id: &Pubkey, ) -> Result<(), AccountFetchError> @@ -52,13 +47,15 @@ where { let validation_address = get_extra_account_metas_address(mint, permissioned_transfer_program_id); - let validation_account_data = get_account_data_fn(validation_address) + let validation_account_data = fetch_account_data_fn(validation_address) .await? .ok_or(ProgramError::InvalidAccountData)?; - ExtraAccountMetaList::add_to_instruction::( + ExtraAccountMetaList::add_to_instruction::<_, _, ExecuteInstruction>( instruction, + fetch_account_data_fn, &validation_account_data, - )?; + ) + .await?; // The onchain helpers pull out the required accounts from an opaque // slice by pubkey, so the order doesn't matter here! instruction.accounts.push(AccountMeta::new_readonly(