TLV Account Resolution: Add support for keys in data (#6847)

* tlv-account-resolution: add `KeyData` configuration

* tlv-account-resolution: add `KeyData` support to `ExtraAccountMeta`

* tlv-account-resolution: add tests for `KeyData`

* rename to `PubkeyData`

* update docs

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

@@ -4,11 +4,13 @@
 //! collection of seeds
 
 use {
-    crate::{error::AccountResolutionError, seeds::Seed},
+    crate::{error::AccountResolutionError, pubkey_data::PubkeyData, seeds::Seed},
     bytemuck::{Pod, Zeroable},
     solana_program::{
-        account_info::AccountInfo, instruction::AccountMeta, program_error::ProgramError,
-        pubkey::Pubkey,
+        account_info::AccountInfo,
+        instruction::AccountMeta,
+        program_error::ProgramError,
+        pubkey::{Pubkey, PUBKEY_BYTES},
     },
     spl_pod::primitives::PodBool,
 };
@@ -66,18 +68,66 @@ where
     Ok(Pubkey::find_program_address(&pda_seeds, program_id).0)
 }
 
+/// Resolve a pubkey from a pubkey data configuration.
+fn resolve_key_data<'a, F>(
+    key_data: &PubkeyData,
+    instruction_data: &[u8],
+    get_account_key_data_fn: F,
+) -> Result<Pubkey, ProgramError>
+where
+    F: Fn(usize) -> Option<(&'a Pubkey, Option<&'a [u8]>)>,
+{
+    match key_data {
+        PubkeyData::Uninitialized => Err(ProgramError::InvalidAccountData),
+        PubkeyData::InstructionData { index } => {
+            let key_start = *index as usize;
+            let key_end = key_start + PUBKEY_BYTES;
+            if key_end > instruction_data.len() {
+                return Err(AccountResolutionError::InstructionDataTooSmall.into());
+            }
+            Ok(Pubkey::new_from_array(
+                instruction_data[key_start..key_end].try_into().unwrap(),
+            ))
+        }
+        PubkeyData::AccountData {
+            account_index,
+            data_index,
+        } => {
+            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 + PUBKEY_BYTES;
+            if account_data.len() < arg_end {
+                return Err(AccountResolutionError::AccountDataTooSmall.into());
+            }
+            Ok(Pubkey::new_from_array(
+                account_data[arg_start..arg_end].try_into().unwrap(),
+            ))
+        }
+    }
+}
+
 /// `Pod` type for defining a required account in a validation account.
 ///
-/// This can either be a standard `AccountMeta` or a PDA.
+/// This can be any of the following:
+///
+/// * A standard `AccountMeta`
+/// * A PDA (with seed configurations)
+/// * A pubkey stored in some data (account or instruction data)
+///
 /// Can be used in TLV-encoded data.
 #[repr(C)]
 #[derive(Clone, Copy, Debug, Default, PartialEq, Pod, Zeroable)]
 pub struct ExtraAccountMeta {
     /// Discriminator to tell whether this represents a standard
-    /// `AccountMeta` or a PDA
+    /// `AccountMeta`, PDA, or pubkey data.
     pub discriminator: u8,
-    /// This `address_config` field can either be the pubkey of the account
-    /// or the seeds used to derive the pubkey from provided inputs
+    /// This `address_config` field can either be the pubkey of the account,
+    /// the seeds used to derive the pubkey from provided inputs (PDA), or the
+    /// data used to derive the pubkey (account or instruction data).
     pub address_config: [u8; 32],
     /// Whether the account should sign
     pub is_signer: PodBool,
@@ -118,6 +168,20 @@ impl ExtraAccountMeta {
         })
     }
 
+    /// Create a `ExtraAccountMeta` from a pubkey data configuration.
+    pub fn new_with_pubkey_data(
+        key_data: &PubkeyData,
+        is_signer: bool,
+        is_writable: bool,
+    ) -> Result<Self, ProgramError> {
+        Ok(Self {
+            discriminator: 2,
+            address_config: PubkeyData::pack_into_address_config(key_data)?,
+            is_signer: is_signer.into(),
+            is_writable: is_writable.into(),
+        })
+    }
+
     /// Create a `ExtraAccountMeta` from a list of seed configurations,
     /// representing a PDA for an external program
     ///
@@ -173,6 +237,14 @@ impl ExtraAccountMeta {
                     is_writable: self.is_writable.into(),
                 })
             }
+            2 => {
+                let key_data = PubkeyData::unpack(&self.address_config)?;
+                Ok(AccountMeta {
+                    pubkey: resolve_key_data(&key_data, instruction_data, get_account_key_data_fn)?,
+                    is_signer: self.is_signer.into(),
+                    is_writable: self.is_writable.into(),
+                })
+            }
             _ => Err(ProgramError::InvalidAccountData),
         }
     }

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

@@ -60,4 +60,13 @@ pub enum AccountResolutionError {
     /// Failed to fetch account
     #[error("Failed to fetch account")]
     AccountFetchFailed,
+    /// Not enough bytes available to pack pubkey data configuration.
+    #[error("Not enough bytes available to pack pubkey data configuration")]
+    NotEnoughBytesForPubkeyData,
+    /// The provided bytes are not valid for a pubkey data configuration
+    #[error("The provided bytes are not valid for a pubkey data configuration")]
+    InvalidBytesForPubkeyData,
+    /// Tried to pack an invalid pubkey data configuration
+    #[error("Tried to pack an invalid pubkey data configuration")]
+    InvalidPubkeyDataConfig,
 }

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

@@ -9,6 +9,7 @@
 
 pub mod account;
 pub mod error;
+pub mod pubkey_data;
 pub mod seeds;
 pub mod state;
 

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

@@ -0,0 +1,185 @@
+//! Types for managing extra account meta keys that may be extracted from some
+//! data.
+//!
+//! This can be either account data from some account in the list of accounts
+//! or from the instruction data itself.
+
+#[cfg(feature = "serde-traits")]
+use serde::{Deserialize, Serialize};
+use {crate::error::AccountResolutionError, solana_program::program_error::ProgramError};
+
+/// Enum to describe a required key stored in some data.
+#[derive(Clone, Debug, PartialEq)]
+#[cfg_attr(feature = "serde-traits", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "serde-traits", serde(rename_all = "camelCase"))]
+pub enum PubkeyData {
+    /// Uninitialized configuration byte space.
+    Uninitialized,
+    /// A pubkey to be resolved from the instruction data.
+    ///
+    /// Packed as:
+    ///     * 1 - Discriminator
+    ///     * 1 - Start index of instruction data
+    ///
+    /// Note: Length is always 32 bytes.
+    InstructionData {
+        /// The index where the address bytes begin in the instruction data.
+        index: u8,
+    },
+    /// A pubkey 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
+    ///
+    /// Note: Length is always 32 bytes.
+    AccountData {
+        /// The index of the account in the entire accounts list.
+        account_index: u8,
+        /// The index where the address bytes begin in the account data.
+        data_index: u8,
+    },
+}
+impl PubkeyData {
+    /// Get the size of a pubkey data configuration.
+    pub fn tlv_size(&self) -> u8 {
+        match self {
+            Self::Uninitialized => 0,
+            // 1 byte for the discriminator, 1 byte for the index.
+            Self::InstructionData { .. } => 1 + 1,
+            // 1 byte for the discriminator, 1 byte for the account index,
+            // 1 byte for the data index.
+            Self::AccountData { .. } => 1 + 1 + 1,
+        }
+    }
+
+    /// Packs a pubkey data configuration into a slice.
+    pub fn pack(&self, dst: &mut [u8]) -> Result<(), ProgramError> {
+        // Because no `PubkeyData` variant is larger than 3 bytes, this check
+        // is sufficient for the data length.
+        if dst.len() != self.tlv_size() as usize {
+            return Err(AccountResolutionError::NotEnoughBytesForPubkeyData.into());
+        }
+        match &self {
+            Self::Uninitialized => {
+                return Err(AccountResolutionError::InvalidPubkeyDataConfig.into())
+            }
+            Self::InstructionData { index } => {
+                dst[0] = 1;
+                dst[1] = *index;
+            }
+            Self::AccountData {
+                account_index,
+                data_index,
+            } => {
+                dst[0] = 2;
+                dst[1] = *account_index;
+                dst[2] = *data_index;
+            }
+        }
+        Ok(())
+    }
+
+    /// Packs a pubkey data configuration into a 32-byte array, filling the
+    /// rest with 0s.
+    pub fn pack_into_address_config(key_data: &Self) -> Result<[u8; 32], ProgramError> {
+        let mut packed = [0u8; 32];
+        let tlv_size = key_data.tlv_size() as usize;
+        key_data.pack(&mut packed[..tlv_size])?;
+        Ok(packed)
+    }
+
+    /// Unpacks a pubkey data configuration from a slice.
+    pub fn unpack(bytes: &[u8]) -> Result<Self, ProgramError> {
+        let (discrim, rest) = bytes
+            .split_first()
+            .ok_or::<ProgramError>(ProgramError::InvalidAccountData)?;
+        match discrim {
+            0 => Ok(Self::Uninitialized),
+            1 => {
+                if rest.is_empty() {
+                    return Err(AccountResolutionError::InvalidBytesForPubkeyData.into());
+                }
+                Ok(Self::InstructionData { index: rest[0] })
+            }
+            2 => {
+                if rest.len() < 2 {
+                    return Err(AccountResolutionError::InvalidBytesForPubkeyData.into());
+                }
+                Ok(Self::AccountData {
+                    account_index: rest[0],
+                    data_index: rest[1],
+                })
+            }
+            _ => Err(ProgramError::InvalidAccountData),
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_pack() {
+        // Should fail if the length is too short.
+        let key = PubkeyData::InstructionData { index: 0 };
+        let mut packed = vec![0u8; key.tlv_size() as usize - 1];
+        assert_eq!(
+            key.pack(&mut packed).unwrap_err(),
+            AccountResolutionError::NotEnoughBytesForPubkeyData.into(),
+        );
+
+        // Should fail if the length is too long.
+        let key = PubkeyData::InstructionData { index: 0 };
+        let mut packed = vec![0u8; key.tlv_size() as usize + 1];
+        assert_eq!(
+            key.pack(&mut packed).unwrap_err(),
+            AccountResolutionError::NotEnoughBytesForPubkeyData.into(),
+        );
+
+        // Can't pack a `PubkeyData::Uninitialized`.
+        let key = PubkeyData::Uninitialized;
+        let mut packed = vec![0u8; key.tlv_size() as usize];
+        assert_eq!(
+            key.pack(&mut packed).unwrap_err(),
+            AccountResolutionError::InvalidPubkeyDataConfig.into(),
+        );
+    }
+
+    #[test]
+    fn test_unpack() {
+        // Can unpack zeroes.
+        let zeroes = [0u8; 32];
+        let key = PubkeyData::unpack(&zeroes).unwrap();
+        assert_eq!(key, PubkeyData::Uninitialized);
+
+        // Should fail for empty bytes.
+        let bytes = [];
+        assert_eq!(
+            PubkeyData::unpack(&bytes).unwrap_err(),
+            ProgramError::InvalidAccountData
+        );
+    }
+
+    fn test_pack_unpack_key(key: PubkeyData) {
+        let tlv_size = key.tlv_size() as usize;
+        let mut packed = vec![0u8; tlv_size];
+        key.pack(&mut packed).unwrap();
+        let unpacked = PubkeyData::unpack(&packed).unwrap();
+        assert_eq!(key, unpacked);
+    }
+
+    #[test]
+    fn test_pack_unpack() {
+        // Instruction data.
+        test_pack_unpack_key(PubkeyData::InstructionData { index: 0 });
+
+        // Account data.
+        test_pack_unpack_key(PubkeyData::AccountData {
+            account_index: 0,
+            data_index: 0,
+        });
+    }
+}