Support Account Data for Seeds in Account Resolution (#5182)

* add support for account data in seeds

* modified tests and docs

* fix a lil test bug

* reformatting feedback

* reworking generics

* scoping over dropping

* replace addressable trait with closure

* cleanup tests

* enable doctest

* bump dev dep versions

* clean up types

libraries/tlv-account-resolution/Cargo.toml

@@ -19,6 +19,9 @@ spl-type-length-value = { version = "0.3", path = "../type-length-value" }
 spl-pod = { version = "0.1", path = "../pod" }
 
 [dev-dependencies]
+futures = "0.3.28"
+futures-util = "0.3"
+solana-client = "1.16.13"
 solana-program-test = "1.16.13"
 solana-sdk = "1.16.13"
 spl-discriminator = { version = "0.1", path = "../discriminator" }

libraries/tlv-account-resolution/README.md

@@ -59,9 +59,21 @@ let mut buffer = vec![0; account_size];
 ExtraAccountMetaList::init::<MyInstruction>(&mut buffer, &extra_metas).unwrap();
 
 // Off-chain, you can add the additional accounts directly from the account data
+// You need to provide the resolver a way to fetch account data off-chain
+let client = RpcClient::new_mock("succeeds".to_string());
 let program_id = Pubkey::new_unique();
 let mut instruction = Instruction::new_with_bytes(program_id, &[0, 1, 2], vec![]);
-ExtraAccountMetaList::add_to_instruction::<MyInstruction>(&mut instruction, &buffer).unwrap();
+ExtraAccountMetaList::add_to_instruction::<_, _, MyInstruction>(
+    &mut instruction,
+    |address: &Pubkey| {
+        client
+            .get_account(address)
+            .map_ok(|acct| Some(acct.data))
+        },
+    &buffer,
+)
+.await
+.unwrap();
 
 // On-chain, you can add the additional accounts *and* account infos
 let mut cpi_instruction = Instruction::new_with_bytes(program_id, &[0, 1, 2], vec![]);

libraries/tlv-account-resolution/src/account.rs

@@ -15,12 +15,15 @@ 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, F>(
     seeds: &[Seed],
-    accounts: &[AccountMeta],
     instruction_data: &[u8],
     program_id: &Pubkey,
-) -> Result<Pubkey, ProgramError> {
+    get_account_key_data_fn: F,
+) -> Result<Pubkey, ProgramError>
+where
+    F: Fn(usize) -> Option<(&'a Pubkey, Option<&'a [u8]>)>,
+{
     let mut pda_seeds: Vec<&[u8]> = vec![];
     for config in seeds {
         match config {
@@ -36,10 +39,27 @@ fn resolve_pda(
             }
             Seed::AccountKey { index } => {
                 let account_index = *index as usize;
-                let account_meta = accounts
-                    .get(account_index)
-                    .ok_or::<ProgramError>(AccountResolutionError::AccountNotFound.into())?;
-                pda_seeds.push(account_meta.pubkey.as_ref());
+                let address = get_account_key_data_fn(account_index)
+                    .ok_or::<ProgramError>(AccountResolutionError::AccountNotFound.into())?
+                    .0;
+                pda_seeds.push(address.as_ref());
+            }
+            Seed::AccountData {
+                account_index,
+                data_index,
+                length,
+            } => {
+                let account_index = *account_index as usize;
+                let account_data = get_account_key_data_fn(account_index)
+                    .ok_or::<ProgramError>(AccountResolutionError::AccountNotFound.into())?
+                    .1
+                    .ok_or::<ProgramError>(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,33 @@ impl ExtraAccountMeta {
 
     /// Resolve an `ExtraAccountMeta` into an `AccountMeta`, potentially
     /// resolving a program-derived address (PDA) if necessary
-    pub fn resolve(
+    pub fn resolve<'a, F>(
         &self,
-        accounts: &[AccountMeta],
         instruction_data: &[u8],
         program_id: &Pubkey,
-    ) -> Result<AccountMeta, ProgramError> {
+        get_account_key_data_fn: F,
+    ) -> Result<AccountMeta, ProgramError>
+    where
+        F: Fn(usize) -> Option<(&'a Pubkey, 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
-                        .get(x.saturating_sub(U8_TOP_BIT) as usize)
-                        .ok_or(AccountResolutionError::AccountNotFound)?
-                        .pubkey
+                    get_account_key_data_fn(x.saturating_sub(U8_TOP_BIT) as usize)
+                        .ok_or::<ProgramError>(AccountResolutionError::AccountNotFound.into())?
+                        .0
                 };
                 let seeds = Seed::unpack_address_config(&self.address_config)?;
                 Ok(AccountMeta {
-                    pubkey: resolve_pda(&seeds, accounts, instruction_data, program_id)?,
+                    pubkey: resolve_pda(
+                        &seeds,
+                        instruction_data,
+                        program_id,
+                        get_account_key_data_fn,
+                    )?,
                     is_signer: self.is_signer.into(),
                     is_writable: self.is_writable.into(),
                 })

libraries/tlv-account-resolution/src/error.rs

@@ -51,4 +51,13 @@ pub enum AccountResolutionError {
     /// Error in checked math operation
     #[error("Error in checked math operation")]
     CalculationFailure,
+    /// 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,
 }

libraries/tlv-account-resolution/src/seeds.rs

@@ -13,11 +13,16 @@
 //!         * N - Literal bytes themselves
 //!     * `Seed::InstructionData`: 1 + 1 + 1 = 3
 //!         * 1 - Discriminator
-//!         * 1 - Index of instruction data
+//!         * 1 - Start index of instruction data
 //!         * 1 - Length of instruction data starting at index
 //!     * `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 - Start 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.
@@ -46,7 +51,7 @@ pub enum Seed {
     /// data
     /// Packed as:
     ///     * 1 - Discriminator
-    ///     * 1 - Index of instruction data
+    ///     * 1 - Start index of instruction data
     ///     * 1 - Length of instruction data starting at index
     InstructionData {
         /// The index where the bytes of an instruction argument begin
@@ -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 - Start 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<Seed, ProgramError> {
     Ok(Seed::AccountKey { index: bytes[0] })
 }
 
+fn unpack_seed_account_data(bytes: &[u8]) -> Result<Seed, ProgramError> {
+    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<Seed>) {
+    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,41 +441,62 @@ 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
 
         let seed = Seed::InstructionData {
             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();