Support unpacking token accounts fields partially (#2970)

* Support unpacking token accounts fields partially

Co-authored-by: Tyera Eulberg <teulberg@gmail.com>
This commit is contained in:
Lijun Wang 2022-03-02 11:34:55 -08:00 committed by GitHub
parent 0c0168f8a9
commit bbbf250a67
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 152 additions and 2 deletions

1
Cargo.lock generated
View File

@ -3903,6 +3903,7 @@ name = "spl-token"
version = "3.3.0"
dependencies = [
"arrayref",
"bytemuck",
"lazy_static",
"num-derive",
"num-traits",

View File

@ -1,7 +1,7 @@
//! State transition types
use {
crate::instruction::MAX_SIGNERS,
crate::{extension::AccountType, instruction::MAX_SIGNERS},
arrayref::{array_mut_ref, array_ref, array_refs, mut_array_refs},
num_enum::{IntoPrimitive, TryFromPrimitive},
solana_program::{
@ -10,6 +10,7 @@ use {
program_pack::{IsInitialized, Pack, Sealed},
pubkey::Pubkey,
},
spl_token::state::GenericTokenAccount,
};
/// Mint data.
@ -289,6 +290,15 @@ fn unpack_coption_u64(src: &[u8; 12]) -> Result<COption<u64>, ProgramError> {
}
}
// `spl_token_program_2022::extension::AccountType::Account` ordinal value
const ACCOUNTTYPE_ACCOUNT: u8 = AccountType::Account as u8;
impl GenericTokenAccount for Account {
fn valid_account_data(account_data: &[u8]) -> bool {
spl_token::state::Account::valid_account_data(account_data)
|| ACCOUNTTYPE_ACCOUNT == *account_data.get(Account::LEN).unwrap_or(&0)
}
}
#[cfg(test)]
pub(crate) mod test {
use super::*;
@ -402,4 +412,56 @@ pub(crate) mod test {
let unpacked = Multisig::unpack(&packed).unwrap();
assert_eq!(unpacked, check);
}
#[test]
fn test_unpack_token_owner() {
// Account data length < Account::LEN, unpack will not return a key
let src: [u8; 12] = [0; 12];
let result = Account::unpack_account_owner(&src);
assert_eq!(result, Option::None);
// The right account data size, unpack will return some key
let src: [u8; Account::LEN] = [0; Account::LEN];
let result = Account::unpack_account_owner(&src);
assert!(result.is_some());
// Account data length > account data size, but not a valid extension,
// unpack will not return a key
let src: [u8; Account::LEN + 5] = [0; Account::LEN + 5];
let result = Account::unpack_account_owner(&src);
assert_eq!(result, Option::None);
// Account data length > account data size with a valid extension,
// expect some key returned
let mut src: [u8; Account::LEN + 5] = [0; Account::LEN + 5];
src[Account::LEN] = 2;
let result = Account::unpack_account_owner(&src);
assert!(result.is_some());
}
#[test]
fn test_unpack_token_mint() {
// Account data length < Account::LEN, unpack will not return a key
let src: [u8; 12] = [0; 12];
let result = Account::unpack_account_mint(&src);
assert_eq!(result, Option::None);
// The right account data size, unpack will return some key
let src: [u8; Account::LEN] = [0; Account::LEN];
let result = Account::unpack_account_mint(&src);
assert!(result.is_some());
// Account data length > account data size, but not a valid extension,
// unpack will not return a key
let src: [u8; Account::LEN + 5] = [0; Account::LEN + 5];
let result = Account::unpack_account_mint(&src);
assert_eq!(result, Option::None);
// Account data length > account data size with a valid extension,
// expect some key returned
let mut src: [u8; Account::LEN + 5] = [0; Account::LEN + 5];
src[Account::LEN] = 2;
let result = Account::unpack_account_mint(&src);
assert!(result.is_some());
}
}

View File

@ -14,6 +14,7 @@ test-bpf = []
[dependencies]
arrayref = "0.3.6"
bytemuck = "1.7.2"
num-derive = "0.3"
num-traits = "0.2"
num_enum = "0.5.4"

View File

@ -7,7 +7,7 @@ use solana_program::{
program_error::ProgramError,
program_option::COption,
program_pack::{IsInitialized, Pack, Sealed},
pubkey::Pubkey,
pubkey::{Pubkey, PUBKEY_BYTES},
};
/// Mint data.
@ -287,6 +287,56 @@ fn unpack_coption_u64(src: &[u8; 12]) -> Result<COption<u64>, ProgramError> {
}
}
const SPL_TOKEN_ACCOUNT_MINT_OFFSET: usize = 0;
const SPL_TOKEN_ACCOUNT_OWNER_OFFSET: usize = 32;
/// A trait for token Account structs to enable efficiently unpacking various fields
/// without unpacking the complete state.
pub trait GenericTokenAccount {
/// Check if the account data is a valid token account
fn valid_account_data(account_data: &[u8]) -> bool;
/// Call after account length has already been verified to unpack the account owner
fn unpack_account_owner_unchecked(account_data: &[u8]) -> &Pubkey {
Self::unpack_pubkey_unchecked(account_data, SPL_TOKEN_ACCOUNT_OWNER_OFFSET)
}
/// Call after account length has already been verified to unpack the account mint
fn unpack_account_mint_unchecked(account_data: &[u8]) -> &Pubkey {
Self::unpack_pubkey_unchecked(account_data, SPL_TOKEN_ACCOUNT_MINT_OFFSET)
}
/// Call after account length has already been verified to unpack a Pubkey at
/// the specified offset. Panics if `account_data.len()` is less than `PUBKEY_BYTES`
fn unpack_pubkey_unchecked(account_data: &[u8], offset: usize) -> &Pubkey {
bytemuck::from_bytes(&account_data[offset..offset + PUBKEY_BYTES])
}
/// Unpacks an account's owner from opaque account data.
fn unpack_account_owner(account_data: &[u8]) -> Option<&Pubkey> {
if Self::valid_account_data(account_data) {
Some(Self::unpack_account_owner_unchecked(account_data))
} else {
None
}
}
/// Unpacks an account's mint from opaque account data.
fn unpack_account_mint(account_data: &[u8]) -> Option<&Pubkey> {
if Self::valid_account_data(account_data) {
Some(Self::unpack_account_mint_unchecked(account_data))
} else {
None
}
}
}
impl GenericTokenAccount for Account {
fn valid_account_data(account_data: &[u8]) -> bool {
account_data.len() == Account::LEN
}
}
#[cfg(test)]
mod tests {
use super::*;
@ -360,4 +410,40 @@ mod tests {
let result = unpack_coption_u64(&src).unwrap_err();
assert_eq!(result, ProgramError::InvalidAccountData);
}
#[test]
fn test_unpack_token_owner() {
// Account data length < Account::LEN, unpack will not return a key
let src: [u8; 12] = [0; 12];
let result = Account::unpack_account_owner(&src);
assert_eq!(result, Option::None);
// The right account data size, unpack will return some key
let src: [u8; Account::LEN] = [0; Account::LEN];
let result = Account::unpack_account_owner(&src);
assert!(result.is_some());
// Account data length > account data size, unpack will not return a key
let src: [u8; Account::LEN + 5] = [0; Account::LEN + 5];
let result = Account::unpack_account_owner(&src);
assert_eq!(result, Option::None);
}
#[test]
fn test_unpack_token_mint() {
// Account data length < Account::LEN, unpack will not return a key
let src: [u8; 12] = [0; 12];
let result = Account::unpack_account_mint(&src);
assert_eq!(result, Option::None);
// The right account data size, unpack will return some key
let src: [u8; Account::LEN] = [0; Account::LEN];
let result = Account::unpack_account_mint(&src);
assert!(result.is_some());
// Account data length > account data size, unpack will not return a key
let src: [u8; Account::LEN + 5] = [0; Account::LEN + 5];
let result = Account::unpack_account_mint(&src);
assert_eq!(result, Option::None);
}
}