浏览代码

Support unpacking token accounts fields partially (#2970)

* Support unpacking token accounts fields partially

Co-authored-by: Tyera Eulberg <teulberg@gmail.com>
Lijun Wang 3 年之前
父节点
当前提交
c8190f3fb3
共有 2 个文件被更改,包括 88 次插入1 次删除
  1. 1 0
      program/Cargo.toml
  2. 87 1
      program/src/state.rs

+ 1 - 0
program/Cargo.toml

@@ -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"

+ 87 - 1
program/src/state.rs

@@ -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);
+    }
 }