Skip to content

Commit

Permalink
add support for account data in seeds
Browse files Browse the repository at this point in the history
  • Loading branch information
buffalojoec committed Aug 31, 2023
1 parent 7b61198 commit fc2df04
Show file tree
Hide file tree
Showing 6 changed files with 245 additions and 58 deletions.
70 changes: 60 additions & 10 deletions libraries/tlv-account-resolution/src/account.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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<Pubkey, ProgramError> {
get_account_data_fn: G,
) -> Result<Pubkey, ProgramError>
where
A: Addressable,
G: Fn(usize) -> Option<&'a [u8]>,
{
let mut pda_seeds: Vec<&[u8]> = vec![];
for config in seeds {
match config {
Expand All @@ -39,7 +44,22 @@ fn resolve_pda(
let account_meta = accounts
.get(account_index)
.ok_or::<ProgramError>(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::<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]);
}
}
}
Expand Down Expand Up @@ -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<AccountMeta, ProgramError> {
get_account_data_fn: G,
) -> Result<AccountMeta, ProgramError>
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(),
})
Expand Down Expand Up @@ -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
}
}
9 changes: 9 additions & 0 deletions libraries/tlv-account-resolution/src/error.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
90 changes: 78 additions & 12 deletions libraries/tlv-account-resolution/src/seeds.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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
Expand All @@ -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,
}
}

Expand Down Expand Up @@ -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(())
}
Expand Down Expand Up @@ -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),
}
}
Expand Down Expand Up @@ -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::*;
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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]
Expand All @@ -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

Expand All @@ -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();
Expand Down
Loading

0 comments on commit fc2df04

Please sign in to comment.