瀏覽代碼

p-token: Fix error codes (#52)

* Fix error codes

* Fix formatting

* Add validation on load

* Fix formatting

* Fix typo

* Relax instruction data length check

* Another relax length check

* More relax instruction data length check

* Move account validation

* Fix rent exempt check

* Update pinocchio reference

* Improve coption validation

* Streamline initialize check

* Add unpack helpers
Fernando Otero 5 月之前
父節點
當前提交
4c3fabd685
共有 37 個文件被更改,包括 229 次插入198 次删除
  1. 2 2
      Cargo.lock
  2. 2 2
      interface/Cargo.toml
  3. 3 3
      interface/src/instruction.rs
  4. 16 6
      interface/src/state/account.rs
  5. 15 0
      interface/src/state/account_state.rs
  6. 9 5
      interface/src/state/mint.rs
  7. 3 3
      interface/src/state/mod.rs
  8. 10 6
      interface/src/state/multisig.rs
  9. 1 1
      p-token/Cargo.toml
  10. 4 4
      p-token/src/entrypoint.rs
  11. 2 6
      p-token/src/processor/amount_to_ui_amount.rs
  12. 3 7
      p-token/src/processor/approve.rs
  13. 4 13
      p-token/src/processor/approve_checked.rs
  14. 3 2
      p-token/src/processor/batch.rs
  15. 3 7
      p-token/src/processor/burn.rs
  16. 4 13
      p-token/src/processor/burn_checked.rs
  17. 10 4
      p-token/src/processor/initialize_account2.rs
  18. 10 4
      p-token/src/processor/initialize_account3.rs
  19. 1 1
      p-token/src/processor/initialize_immutable_owner.rs
  20. 3 2
      p-token/src/processor/initialize_multisig.rs
  21. 4 2
      p-token/src/processor/initialize_multisig2.rs
  22. 3 7
      p-token/src/processor/mint_to.rs
  23. 4 13
      p-token/src/processor/mint_to_checked.rs
  24. 31 0
      p-token/src/processor/mod.rs
  25. 9 3
      p-token/src/processor/revoke.rs
  26. 17 15
      p-token/src/processor/set_authority.rs
  27. 1 1
      p-token/src/processor/shared/approve.rs
  28. 6 7
      p-token/src/processor/shared/burn.rs
  29. 2 2
      p-token/src/processor/shared/initialize_account.rs
  30. 29 28
      p-token/src/processor/shared/initialize_mint.rs
  31. 1 1
      p-token/src/processor/shared/initialize_multisig.rs
  32. 1 1
      p-token/src/processor/shared/mint_to.rs
  33. 3 3
      p-token/src/processor/shared/toggle_account_state.rs
  34. 2 2
      p-token/src/processor/shared/transfer.rs
  35. 3 7
      p-token/src/processor/transfer.rs
  36. 4 13
      p-token/src/processor/transfer_checked.rs
  37. 1 2
      p-token/src/processor/ui_amount_to_amount.rs

+ 2 - 2
Cargo.lock

@@ -2805,9 +2805,9 @@ checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184"
 
 [[package]]
 name = "pinocchio"
-version = "0.8.0"
+version = "0.8.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "2281a8e9e947c480ce0b64e2d6c6349d72890ba5db6503d5568edf96f304cba9"
+checksum = "7c33b58567c11b07749cefbb8320ac023f3387c57807aeb8e3b1262501b6e9f0"
 
 [[package]]
 name = "pinocchio-log"

+ 2 - 2
interface/Cargo.toml

@@ -12,8 +12,8 @@ readme = "./README.md"
 crate-type = ["rlib"]
 
 [dependencies]
-pinocchio = "0.8"
-pinocchio-pubkey = "0.2.4"
+pinocchio = "0.8.4"
+pinocchio-pubkey = "0.2"
 
 [dev-dependencies]
 strum = "0.27"

+ 3 - 3
interface/src/instruction.rs

@@ -1,6 +1,6 @@
 //! Instruction types.
 
-use pinocchio::program_error::ProgramError;
+use {crate::error::TokenError, pinocchio::program_error::ProgramError};
 
 /// Instructions supported by the token program.
 #[repr(u8)]
@@ -531,7 +531,7 @@ impl TryFrom<u8> for TokenInstruction {
         match value {
             // SAFETY: `value` is guaranteed to be in the range of the enum variants.
             0..=24 | 38 | 255 => Ok(unsafe { core::mem::transmute::<u8, TokenInstruction>(value) }),
-            _ => Err(ProgramError::InvalidInstructionData),
+            _ => Err(TokenError::InvalidInstruction.into()),
         }
     }
 }
@@ -559,7 +559,7 @@ impl TryFrom<u8> for AuthorityType {
         match value {
             // SAFETY: `value` is guaranteed to be in the range of the enum variants.
             0..=3 => Ok(unsafe { core::mem::transmute::<u8, AuthorityType>(value) }),
-            _ => Err(ProgramError::InvalidInstructionData),
+            _ => Err(TokenError::InvalidInstruction.into()),
         }
     }
 }

+ 16 - 6
interface/src/state/account.rs

@@ -1,6 +1,6 @@
 use {
     super::{account_state::AccountState, COption, Initializable, Transmutable},
-    pinocchio::pubkey::Pubkey,
+    pinocchio::{program_error::ProgramError, pubkey::Pubkey},
 };
 
 /// Incinerator address.
@@ -27,7 +27,7 @@ pub struct Account {
     delegate: COption<Pubkey>,
 
     /// The account's state.
-    pub state: AccountState,
+    state: u8,
 
     /// Indicates whether this account represents a native token or not.
     is_native: [u8; 4],
@@ -46,6 +46,16 @@ pub struct Account {
 }
 
 impl Account {
+    #[inline(always)]
+    pub fn set_account_state(&mut self, state: AccountState) {
+        self.state = state as u8;
+    }
+
+    #[inline(always)]
+    pub fn account_state(&self) -> Result<AccountState, ProgramError> {
+        AccountState::try_from(self.state)
+    }
+
     #[inline(always)]
     pub fn set_amount(&mut self, amount: u64) {
         self.amount = amount.to_le_bytes();
@@ -131,8 +141,8 @@ impl Account {
     }
 
     #[inline(always)]
-    pub fn is_frozen(&self) -> bool {
-        self.state == AccountState::Frozen
+    pub fn is_frozen(&self) -> Result<bool, ProgramError> {
+        AccountState::try_from(self.state).map(|state| state == AccountState::Frozen)
     }
 
     #[inline(always)]
@@ -147,7 +157,7 @@ impl Transmutable for Account {
 
 impl Initializable for Account {
     #[inline(always)]
-    fn is_initialized(&self) -> bool {
-        self.state != AccountState::Uninitialized
+    fn is_initialized(&self) -> Result<bool, ProgramError> {
+        AccountState::try_from(self.state).map(|state| state != AccountState::Uninitialized)
     }
 }

+ 15 - 0
interface/src/state/account_state.rs

@@ -1,3 +1,5 @@
+use pinocchio::program_error::ProgramError;
+
 #[repr(u8)]
 #[derive(Clone, Copy, Debug, PartialEq)]
 pub enum AccountState {
@@ -13,3 +15,16 @@ pub enum AccountState {
     /// this account.
     Frozen,
 }
+
+impl TryFrom<u8> for AccountState {
+    type Error = ProgramError;
+
+    #[inline(always)]
+    fn try_from(value: u8) -> Result<Self, Self::Error> {
+        match value {
+            // SAFETY: `value` is guaranteed to be in the range of the enum variants.
+            0..=2 => Ok(unsafe { core::mem::transmute::<u8, AccountState>(value) }),
+            _ => Err(ProgramError::InvalidAccountData),
+        }
+    }
+}

+ 9 - 5
interface/src/state/mint.rs

@@ -1,6 +1,6 @@
 use {
     super::{COption, Initializable, Transmutable},
-    pinocchio::pubkey::Pubkey,
+    pinocchio::{program_error::ProgramError, pubkey::Pubkey},
 };
 
 /// Internal representation of a mint data.
@@ -10,7 +10,7 @@ pub struct Mint {
     /// be provided during mint creation. If no mint authority is present
     /// then the mint has a fixed supply and no further tokens may be
     /// minted.
-    pub mint_authority: COption<Pubkey>,
+    mint_authority: COption<Pubkey>,
 
     /// Total supply of tokens.
     supply: [u8; 8],
@@ -24,7 +24,7 @@ pub struct Mint {
     // Indicates whether the freeze authority is present or not.
     //freeze_authority_option: [u8; 4],
     /// Optional authority to freeze token accounts.
-    pub freeze_authority: COption<Pubkey>,
+    freeze_authority: COption<Pubkey>,
 }
 
 impl Mint {
@@ -91,7 +91,11 @@ impl Transmutable for Mint {
 
 impl Initializable for Mint {
     #[inline(always)]
-    fn is_initialized(&self) -> bool {
-        self.is_initialized == 1
+    fn is_initialized(&self) -> Result<bool, ProgramError> {
+        match self.is_initialized {
+            0 => Ok(false),
+            1 => Ok(true),
+            _ => Err(ProgramError::InvalidAccountData),
+        }
     }
 }

+ 3 - 3
interface/src/state/mod.rs

@@ -23,7 +23,7 @@ pub trait Transmutable {
 /// Trait to represent a type that can be initialized.
 pub trait Initializable {
     /// Return `true` if the object is initialized.
-    fn is_initialized(&self) -> bool;
+    fn is_initialized(&self) -> Result<bool, ProgramError>;
 }
 
 /// Return a reference for an initialized `T` from the given bytes.
@@ -35,7 +35,7 @@ pub trait Initializable {
 pub unsafe fn load<T: Initializable + Transmutable>(bytes: &[u8]) -> Result<&T, ProgramError> {
     load_unchecked(bytes).and_then(|t: &T| {
         // checks if the data is initialized
-        if t.is_initialized() {
+        if t.is_initialized()? {
             Ok(t)
         } else {
             Err(ProgramError::UninitializedAccount)
@@ -69,7 +69,7 @@ pub unsafe fn load_mut<T: Initializable + Transmutable>(
 ) -> Result<&mut T, ProgramError> {
     load_mut_unchecked(bytes).and_then(|t: &mut T| {
         // checks if the data is initialized
-        if t.is_initialized() {
+        if t.is_initialized()? {
             Ok(t)
         } else {
             Err(ProgramError::UninitializedAccount)

+ 10 - 6
interface/src/state/multisig.rs

@@ -1,6 +1,6 @@
 use {
     super::{Initializable, Transmutable},
-    pinocchio::pubkey::Pubkey,
+    pinocchio::{program_error::ProgramError, pubkey::Pubkey},
 };
 
 /// Minimum number of multisignature signers (min N)
@@ -18,10 +18,10 @@ pub struct Multisig {
     /// Number of valid signers.
     pub n: u8,
 
-    /// Is `true` if this structure has been initialized
+    /// Is `true` if this structure has been initialized.
     is_initialized: u8,
 
-    /// Signer public keys
+    /// Signer public keys.
     pub signers: [Pubkey; MAX_SIGNERS as usize],
 }
 
@@ -39,13 +39,17 @@ impl Multisig {
 }
 
 impl Transmutable for Multisig {
-    /// The length of the `Mint` account data.
+    /// The length of the `Multisig` account data.
     const LEN: usize = core::mem::size_of::<Multisig>();
 }
 
 impl Initializable for Multisig {
     #[inline(always)]
-    fn is_initialized(&self) -> bool {
-        self.is_initialized == 1
+    fn is_initialized(&self) -> Result<bool, ProgramError> {
+        match self.is_initialized {
+            0 => Ok(false),
+            1 => Ok(true),
+            _ => Err(ProgramError::InvalidAccountData),
+        }
     }
 }

+ 1 - 1
p-token/Cargo.toml

@@ -15,7 +15,7 @@ crate-type = ["cdylib"]
 logging = []
 
 [dependencies]
-pinocchio = "0.8"
+pinocchio = "0.8.4"
 pinocchio-log = { version = "0.4", default-features = false }
 spl-token-interface = { version = "^0", path = "../interface" }
 

+ 4 - 4
p-token/src/entrypoint.rs

@@ -36,7 +36,7 @@ pub fn process_instruction(
     instruction_data: &[u8],
 ) -> ProgramResult {
     let [discriminator, remaining @ ..] = instruction_data else {
-        return Err(ProgramError::InvalidInstructionData);
+        return Err(TokenError::InvalidInstruction.into());
     };
 
     let result = if *discriminator == 255 {
@@ -78,7 +78,7 @@ pub(crate) fn inner_process_instruction(
     instruction_data: &[u8],
 ) -> ProgramResult {
     let [discriminator, instruction_data @ ..] = instruction_data else {
-        return Err(ProgramError::InvalidInstructionData);
+        return Err(TokenError::InvalidInstruction.into());
     };
 
     match *discriminator {
@@ -194,7 +194,7 @@ fn inner_process_remaining_instruction(
             #[cfg(feature = "logging")]
             pinocchio::msg!("Instruction: Revoke");
 
-            process_revoke(accounts, instruction_data)
+            process_revoke(accounts)
         }
         // 6 - SetAuthority
         6 => {
@@ -280,6 +280,6 @@ fn inner_process_remaining_instruction(
 
             process_withdraw_excess_lamports(accounts)
         }
-        _ => Err(ProgramError::InvalidInstructionData),
+        _ => Err(TokenError::InvalidInstruction.into()),
     }
 }

+ 2 - 6
p-token/src/processor/amount_to_ui_amount.rs

@@ -1,5 +1,5 @@
 use {
-    super::{check_account_owner, MAX_FORMATTED_DIGITS},
+    super::{check_account_owner, unpack_amount, MAX_FORMATTED_DIGITS},
     core::str::from_utf8_unchecked,
     pinocchio::{
         account_info::AccountInfo, program::set_return_data, program_error::ProgramError,
@@ -16,11 +16,7 @@ pub fn process_amount_to_ui_amount(
     accounts: &[AccountInfo],
     instruction_data: &[u8],
 ) -> ProgramResult {
-    let amount = u64::from_le_bytes(
-        instruction_data
-            .try_into()
-            .map_err(|_error| ProgramError::InvalidInstructionData)?,
-    );
+    let amount = unpack_amount(instruction_data)?;
 
     let mint_info = accounts.first().ok_or(ProgramError::NotEnoughAccountKeys)?;
     check_account_owner(mint_info)?;

+ 3 - 7
p-token/src/processor/approve.rs

@@ -1,15 +1,11 @@
 use {
-    super::shared,
-    pinocchio::{account_info::AccountInfo, program_error::ProgramError, ProgramResult},
+    super::{shared, unpack_amount},
+    pinocchio::{account_info::AccountInfo, ProgramResult},
 };
 
 #[inline(always)]
 pub fn process_approve(accounts: &[AccountInfo], instruction_data: &[u8]) -> ProgramResult {
-    let amount = u64::from_le_bytes(
-        instruction_data
-            .try_into()
-            .map_err(|_error| ProgramError::InvalidInstructionData)?,
-    );
+    let amount = unpack_amount(instruction_data)?;
 
     shared::approve::process_approve(accounts, amount, None)
 }

+ 4 - 13
p-token/src/processor/approve_checked.rs

@@ -1,20 +1,11 @@
 use {
-    super::shared,
-    pinocchio::{account_info::AccountInfo, program_error::ProgramError, ProgramResult},
+    super::{shared, unpack_amount_and_decimals},
+    pinocchio::{account_info::AccountInfo, ProgramResult},
 };
 
 #[inline(always)]
 pub fn process_approve_checked(accounts: &[AccountInfo], instruction_data: &[u8]) -> ProgramResult {
-    // expected u64 (8) + u8 (1)
-    let (amount, decimals) = if instruction_data.len() == 9 {
-        let (amount, decimals) = instruction_data.split_at(core::mem::size_of::<u64>());
-        (
-            u64::from_le_bytes(amount.try_into().unwrap()),
-            decimals.first().copied(),
-        )
-    } else {
-        return Err(ProgramError::InvalidInstructionData);
-    };
+    let (amount, decimals) = unpack_amount_and_decimals(instruction_data)?;
 
-    shared::approve::process_approve(accounts, amount, decimals)
+    shared::approve::process_approve(accounts, amount, Some(decimals))
 }

+ 3 - 2
p-token/src/processor/batch.rs

@@ -1,6 +1,7 @@
 use {
     crate::entrypoint::inner_process_instruction,
     pinocchio::{account_info::AccountInfo, program_error::ProgramError, ProgramResult},
+    spl_token_interface::error::TokenError,
 };
 
 /// The size of the batch instruction header.
@@ -17,7 +18,7 @@ pub fn process_batch(mut accounts: &[AccountInfo], mut instruction_data: &[u8])
 
         if instruction_data.len() < IX_HEADER_SIZE {
             // The instruction data must have at least two bytes.
-            return Err(ProgramError::InvalidInstructionData);
+            return Err(TokenError::InvalidInstruction.into());
         }
 
         // SAFETY: The instruction data is guaranteed to have at least two bytes
@@ -27,7 +28,7 @@ pub fn process_batch(mut accounts: &[AccountInfo], mut instruction_data: &[u8])
         let data_offset = IX_HEADER_SIZE + unsafe { *instruction_data.get_unchecked(1) as usize };
 
         if instruction_data.len() < data_offset || data_offset == IX_HEADER_SIZE {
-            return Err(ProgramError::InvalidInstructionData);
+            return Err(TokenError::InvalidInstruction.into());
         }
 
         if accounts.len() < expected_accounts {

+ 3 - 7
p-token/src/processor/burn.rs

@@ -1,15 +1,11 @@
 use {
-    super::shared,
-    pinocchio::{account_info::AccountInfo, program_error::ProgramError, ProgramResult},
+    super::{shared, unpack_amount},
+    pinocchio::{account_info::AccountInfo, ProgramResult},
 };
 
 #[inline(always)]
 pub fn process_burn(accounts: &[AccountInfo], instruction_data: &[u8]) -> ProgramResult {
-    let amount = u64::from_le_bytes(
-        instruction_data
-            .try_into()
-            .map_err(|_error| ProgramError::InvalidInstructionData)?,
-    );
+    let amount = unpack_amount(instruction_data)?;
 
     shared::burn::process_burn(accounts, amount, None)
 }

+ 4 - 13
p-token/src/processor/burn_checked.rs

@@ -1,20 +1,11 @@
 use {
-    super::shared,
-    pinocchio::{account_info::AccountInfo, program_error::ProgramError, ProgramResult},
+    super::{shared, unpack_amount_and_decimals},
+    pinocchio::{account_info::AccountInfo, ProgramResult},
 };
 
 #[inline(always)]
 pub fn process_burn_checked(accounts: &[AccountInfo], instruction_data: &[u8]) -> ProgramResult {
-    // expected u64 (8) + u8 (1)
-    let (amount, decimals) = if instruction_data.len() == 9 {
-        let (amount, decimals) = instruction_data.split_at(core::mem::size_of::<u64>());
-        (
-            u64::from_le_bytes(amount.try_into().unwrap()),
-            decimals.first().copied(),
-        )
-    } else {
-        return Err(ProgramError::InvalidInstructionData);
-    };
+    let (amount, decimals) = unpack_amount_and_decimals(instruction_data)?;
 
-    shared::burn::process_burn(accounts, amount, decimals)
+    shared::burn::process_burn(accounts, amount, Some(decimals))
 }

+ 10 - 4
p-token/src/processor/initialize_account2.rs

@@ -1,8 +1,11 @@
 use {
     super::shared,
     pinocchio::{
-        account_info::AccountInfo, program_error::ProgramError, pubkey::Pubkey, ProgramResult,
+        account_info::AccountInfo,
+        pubkey::{Pubkey, PUBKEY_BYTES},
+        ProgramResult,
     },
+    spl_token_interface::error::TokenError,
 };
 
 #[inline(always)]
@@ -10,9 +13,12 @@ pub fn process_initialize_account2(
     accounts: &[AccountInfo],
     instruction_data: &[u8],
 ) -> ProgramResult {
-    let owner: &Pubkey = instruction_data
-        .try_into()
-        .map_err(|_error| ProgramError::InvalidInstructionData)?;
+    let owner = if instruction_data.len() >= PUBKEY_BYTES {
+        // SAFETY: The minimum size of the instruction data is `PUBKEY_BYTES` bytes.
+        unsafe { &*(instruction_data.as_ptr() as *const Pubkey) }
+    } else {
+        return Err(TokenError::InvalidInstruction.into());
+    };
 
     shared::initialize_account::process_initialize_account(accounts, Some(owner), true)
 }

+ 10 - 4
p-token/src/processor/initialize_account3.rs

@@ -1,8 +1,11 @@
 use {
     super::shared,
     pinocchio::{
-        account_info::AccountInfo, program_error::ProgramError, pubkey::Pubkey, ProgramResult,
+        account_info::AccountInfo,
+        pubkey::{Pubkey, PUBKEY_BYTES},
+        ProgramResult,
     },
+    spl_token_interface::error::TokenError,
 };
 
 #[inline(always)]
@@ -10,9 +13,12 @@ pub fn process_initialize_account3(
     accounts: &[AccountInfo],
     instruction_data: &[u8],
 ) -> ProgramResult {
-    let owner: &Pubkey = instruction_data
-        .try_into()
-        .map_err(|_error| ProgramError::InvalidInstructionData)?;
+    let owner = if instruction_data.len() >= PUBKEY_BYTES {
+        // SAFETY: The minimum size of the instruction data is `PUBKEY_BYTES` bytes.
+        unsafe { &*(instruction_data.as_ptr() as *const Pubkey) }
+    } else {
+        return Err(TokenError::InvalidInstruction.into());
+    };
 
     shared::initialize_account::process_initialize_account(accounts, Some(owner), false)
 }

+ 1 - 1
p-token/src/processor/initialize_immutable_owner.rs

@@ -13,7 +13,7 @@ pub fn process_initialize_immutable_owner(accounts: &[AccountInfo]) -> ProgramRe
     // SAFETY: single immutable borrow to `token_account_info` account data.
     let account = unsafe { load_unchecked::<Account>(token_account_info.borrow_data_unchecked())? };
 
-    if account.is_initialized() {
+    if account.is_initialized()? {
         return Err(TokenError::AlreadyInUse.into());
     }
     // Please upgrade to SPL Token 2022 for immutable owner support.

+ 3 - 2
p-token/src/processor/initialize_multisig.rs

@@ -1,6 +1,7 @@
 use {
     super::shared,
-    pinocchio::{account_info::AccountInfo, program_error::ProgramError, ProgramResult},
+    pinocchio::{account_info::AccountInfo, ProgramResult},
+    spl_token_interface::error::TokenError,
 };
 
 #[inline(always)]
@@ -10,7 +11,7 @@ pub fn process_initialize_multisig(
 ) -> ProgramResult {
     let m = instruction_data
         .first()
-        .ok_or(ProgramError::InvalidInstructionData)?;
+        .ok_or(TokenError::InvalidInstruction)?;
 
     shared::initialize_multisig::process_initialize_multisig(accounts, *m, true)
 }

+ 4 - 2
p-token/src/processor/initialize_multisig2.rs

@@ -1,6 +1,7 @@
 use {
     super::shared,
-    pinocchio::{account_info::AccountInfo, program_error::ProgramError, ProgramResult},
+    pinocchio::{account_info::AccountInfo, ProgramResult},
+    spl_token_interface::error::TokenError,
 };
 
 pub fn process_initialize_multisig2(
@@ -9,6 +10,7 @@ pub fn process_initialize_multisig2(
 ) -> ProgramResult {
     let m = instruction_data
         .first()
-        .ok_or(ProgramError::InvalidInstructionData)?;
+        .ok_or(TokenError::InvalidInstruction)?;
+
     shared::initialize_multisig::process_initialize_multisig(accounts, *m, false)
 }

+ 3 - 7
p-token/src/processor/mint_to.rs

@@ -1,15 +1,11 @@
 use {
-    super::shared,
-    pinocchio::{account_info::AccountInfo, program_error::ProgramError, ProgramResult},
+    super::{shared, unpack_amount},
+    pinocchio::{account_info::AccountInfo, ProgramResult},
 };
 
 #[inline(always)]
 pub fn process_mint_to(accounts: &[AccountInfo], instruction_data: &[u8]) -> ProgramResult {
-    let amount = u64::from_le_bytes(
-        instruction_data
-            .try_into()
-            .map_err(|_error| ProgramError::InvalidInstructionData)?,
-    );
+    let amount = unpack_amount(instruction_data)?;
 
     shared::mint_to::process_mint_to(accounts, amount, None)
 }

+ 4 - 13
p-token/src/processor/mint_to_checked.rs

@@ -1,20 +1,11 @@
 use {
-    super::shared,
-    pinocchio::{account_info::AccountInfo, program_error::ProgramError, ProgramResult},
+    super::{shared, unpack_amount_and_decimals},
+    pinocchio::{account_info::AccountInfo, ProgramResult},
 };
 
 #[inline(always)]
 pub fn process_mint_to_checked(accounts: &[AccountInfo], instruction_data: &[u8]) -> ProgramResult {
-    // expected u64 (8) + u8 (1)
-    let (amount, decimals) = if instruction_data.len() == 9 {
-        let (amount, decimals) = instruction_data.split_at(core::mem::size_of::<u64>());
-        (
-            u64::from_le_bytes(amount.try_into().unwrap()),
-            decimals.first().copied(),
-        )
-    } else {
-        return Err(ProgramError::InvalidInstructionData);
-    };
+    let (amount, decimals) = unpack_amount_and_decimals(instruction_data)?;
 
-    shared::mint_to::process_mint_to(accounts, amount, decimals)
+    shared::mint_to::process_mint_to(accounts, amount, Some(decimals))
 }

+ 31 - 0
p-token/src/processor/mod.rs

@@ -64,6 +64,9 @@ pub use {
     withdraw_excess_lamports::process_withdraw_excess_lamports,
 };
 
+/// Number of bytes in a `u64`.
+const U64_BYTES: usize = core::mem::size_of::<u64>();
+
 /// Maximum number of digits in a formatted `u64`.
 ///
 /// The maximum number of digits is equal to the maximum number
@@ -180,3 +183,31 @@ fn try_ui_amount_into_amount(ui_amount: &str, decimals: u8) -> Result<u64, Progr
             .map_err(|_| ProgramError::InvalidArgument)
     }
 }
+
+/// Unpacks a `u64` amount from the instruction data.
+#[inline(always)]
+const fn unpack_amount(instruction_data: &[u8]) -> Result<u64, TokenError> {
+    // expected u64 (8)
+    if instruction_data.len() >= U64_BYTES {
+        // SAFETY: The minimum size of the instruction data is `U64_BYTES` bytes.
+        Ok(unsafe { u64::from_le_bytes(*(instruction_data.as_ptr() as *const [u8; U64_BYTES])) })
+    } else {
+        Err(TokenError::InvalidInstruction)
+    }
+}
+
+/// Unpacks a `u64` amount and an optional `u8` from the instruction data.
+#[inline(always)]
+const fn unpack_amount_and_decimals(instruction_data: &[u8]) -> Result<(u64, u8), TokenError> {
+    // expected u64 (8) + u8 (1)
+    if instruction_data.len() >= 9 {
+        let (amount, decimals) = instruction_data.split_at(U64_BYTES);
+        Ok((
+            // SAFETY: The size of `amount` is `U64_BYTES` bytes.
+            unsafe { u64::from_le_bytes(*(amount.as_ptr() as *const [u8; U64_BYTES])) },
+            decimals[0],
+        ))
+    } else {
+        Err(TokenError::InvalidInstruction)
+    }
+}

+ 9 - 3
p-token/src/processor/revoke.rs

@@ -8,8 +8,8 @@ use {
 };
 
 #[inline(always)]
-pub fn process_revoke(accounts: &[AccountInfo], _instruction_data: &[u8]) -> ProgramResult {
-    let [source_account_info, owner_info, remaining @ ..] = accounts else {
+pub fn process_revoke(accounts: &[AccountInfo]) -> ProgramResult {
+    let [source_account_info, remaining @ ..] = accounts else {
         return Err(ProgramError::NotEnoughAccountKeys);
     };
 
@@ -18,7 +18,13 @@ pub fn process_revoke(accounts: &[AccountInfo], _instruction_data: &[u8]) -> Pro
     let source_account =
         unsafe { load_mut::<Account>(source_account_info.borrow_mut_data_unchecked())? };
 
-    if source_account.is_frozen() {
+    // Unpacking the remaining accounts to get the owner account at this point
+    // to maintain the same order as SPL Token.
+    let [owner_info, remaining @ ..] = remaining else {
+        return Err(ProgramError::NotEnoughAccountKeys);
+    };
+
+    if source_account.is_frozen()? {
         return Err(TokenError::AccountFrozen.into());
     }
 

+ 17 - 15
p-token/src/processor/set_authority.rs

@@ -14,21 +14,23 @@ use {
 pub fn process_set_authority(accounts: &[AccountInfo], instruction_data: &[u8]) -> ProgramResult {
     // Validates the instruction data.
 
-    // SAFETY: The expected size of the instruction data is either 2 or 34 bytes:
-    //   - authority_type (1 byte)
-    //   - option + new_authority (1 byte + 32 bytes)
-    let (authority_type, new_authority) = unsafe {
-        match instruction_data.len() {
-            2 if *instruction_data.get_unchecked(1) == 0 => (
-                AuthorityType::try_from(*instruction_data.get_unchecked(0))?,
-                None,
-            ),
-            34 if *instruction_data.get_unchecked(1) == 1 => (
-                AuthorityType::try_from(*instruction_data.get_unchecked(0))?,
-                Some(&*(instruction_data.as_ptr().add(2) as *const Pubkey)),
-            ),
-            _ => return Err(ProgramError::InvalidInstructionData),
+    let (authority_type, new_authority) = if instruction_data.len() >= 2 {
+        // SAFETY: The expected size of the instruction data is either 2 or 34 bytes:
+        //   - authority_type (1 byte)
+        //   - option + new_authority (1 byte + 32 bytes)
+        unsafe {
+            let authority_type = AuthorityType::try_from(*instruction_data.get_unchecked(0))?;
+            let new_authority = if *instruction_data.get_unchecked(1) == 0 {
+                None
+            } else if *instruction_data.get_unchecked(1) == 1 && instruction_data.len() >= 34 {
+                Some(&*(instruction_data.as_ptr().add(2) as *const Pubkey))
+            } else {
+                return Err(TokenError::InvalidInstruction.into());
+            };
+            (authority_type, new_authority)
         }
+    } else {
+        return Err(TokenError::InvalidInstruction.into());
     };
 
     // Validates the accounts.
@@ -42,7 +44,7 @@ pub fn process_set_authority(accounts: &[AccountInfo], instruction_data: &[u8])
         // `load_mut` validates that the account is initialized.
         let account = unsafe { load_mut::<Account>(account_info.borrow_mut_data_unchecked())? };
 
-        if account.is_frozen() {
+        if account.is_frozen()? {
             return Err(TokenError::AccountFrozen.into());
         }
 

+ 1 - 1
p-token/src/processor/shared/approve.rs

@@ -51,7 +51,7 @@ pub fn process_approve(
     let source_account =
         unsafe { load_mut::<Account>(source_account_info.borrow_mut_data_unchecked())? };
 
-    if source_account.is_frozen() {
+    if source_account.is_frozen()? {
         return Err(TokenError::AccountFrozen.into());
     }
 

+ 6 - 7
p-token/src/processor/shared/burn.rs

@@ -21,8 +21,13 @@ pub fn process_burn(
     // `load_mut` validates that the account is initialized.
     let source_account =
         unsafe { load_mut::<Account>(source_account_info.borrow_mut_data_unchecked())? };
+    // SAFETY: single mutable borrow to `mint_info` account data and
+    // `load_mut` validates that the mint is initialized; additionally, an
+    // account cannot be both a token account and a mint, so if duplicates are
+    // passed in, one of them will fail the `load_mut` check.
+    let mint = unsafe { load_mut::<Mint>(mint_info.borrow_mut_data_unchecked())? };
 
-    if source_account.is_frozen() {
+    if source_account.is_frozen()? {
         return Err(TokenError::AccountFrozen.into());
     }
     if source_account.is_native() {
@@ -36,12 +41,6 @@ pub fn process_burn(
         .checked_sub(amount)
         .ok_or(TokenError::InsufficientFunds)?;
 
-    // SAFETY: single mutable borrow to `mint_info` account data and
-    // `load_mut` validates that the mint is initialized; additionally, an
-    // account cannot be both a token account and a mint, so if duplicates are
-    // passed in, one of them will fail the `load_mut` check.
-    let mint = unsafe { load_mut::<Mint>(mint_info.borrow_mut_data_unchecked())? };
-
     if mint_info.key() != &source_account.mint {
         return Err(TokenError::MintMismatch.into());
     }

+ 2 - 2
p-token/src/processor/shared/initialize_account.rs

@@ -62,7 +62,7 @@ pub fn process_initialize_account(
     let account =
         unsafe { load_mut_unchecked::<Account>(new_account_info.borrow_mut_data_unchecked())? };
 
-    if account.is_initialized() {
+    if account.is_initialized()? {
         return Err(TokenError::AlreadyInUse.into());
     }
 
@@ -80,7 +80,7 @@ pub fn process_initialize_account(
         };
     }
 
-    account.state = AccountState::Initialized;
+    account.set_account_state(AccountState::Initialized);
     account.mint = *mint_info.key();
     account.owner = *owner;
 

+ 29 - 28
p-token/src/processor/shared/initialize_mint.rs

@@ -1,5 +1,4 @@
 use {
-    core::mem::size_of,
     pinocchio::{
         account_info::AccountInfo,
         program_error::ProgramError,
@@ -21,24 +20,26 @@ pub fn process_initialize_mint(
 ) -> ProgramResult {
     // Validates the instruction data.
 
-    // SAFETY: The minimum size of the instruction data is either 34 or 66 bytes:
-    //   - decimals (1 byte)
-    //   - mint_authority (32 bytes)
-    //   - option + freeze_authority (1 byte + 32 bytes)
-    let (decimals, mint_authority, freeze_authority) = unsafe {
-        match instruction_data.len() {
-            34 if *instruction_data.get_unchecked(33) == 0 => (
-                *instruction_data.get_unchecked(0),
-                &*(instruction_data.as_ptr().add(1) as *const Pubkey),
-                None,
-            ),
-            66 if *instruction_data.get_unchecked(33) == 1 => (
-                *instruction_data.get_unchecked(0),
-                &*(instruction_data.as_ptr().add(1) as *const Pubkey),
-                Some(&*(instruction_data.as_ptr().add(34) as *const Pubkey)),
-            ),
-            _ => return Err(ProgramError::InvalidInstructionData),
+    let (decimals, mint_authority, freeze_authority) = if instruction_data.len() >= 34 {
+        // SAFETY: The minimum size of the instruction data is either 34 or 66 bytes:
+        //   - decimals (1 byte)
+        //   - mint_authority (32 bytes)
+        //   - option + freeze_authority (1 byte + 32 bytes)
+        unsafe {
+            let decimals = *instruction_data.get_unchecked(0);
+            let mint_authority = &*(instruction_data.as_ptr().add(1) as *const Pubkey);
+            let freeze_authority = if *instruction_data.get_unchecked(33) == 0 {
+                None
+            } else if *instruction_data.get_unchecked(33) == 1 && instruction_data.len() >= 66 {
+                Some(&*(instruction_data.as_ptr().add(34) as *const Pubkey))
+            } else {
+                return Err(TokenError::InvalidInstruction.into());
+            };
+
+            (decimals, mint_authority, freeze_authority)
         }
+    } else {
+        return Err(TokenError::InvalidInstruction.into());
     };
 
     // Validates the accounts.
@@ -55,24 +56,24 @@ pub fn process_initialize_mint(
         (mint_info, None)
     };
 
-    // SAFETY: single mutable borrow to `mint_info` account data.
-    let mint = unsafe { load_mut_unchecked::<Mint>(mint_info.borrow_mut_data_unchecked())? };
-
-    if mint.is_initialized() {
-        return Err(TokenError::AlreadyInUse.into());
-    }
-
-    // Check rent-exempt status of the mint account.
+    let mint_data_len = mint_info.data_len();
 
     let is_exempt = if let Some(rent_sysvar_info) = rent_sysvar_info {
         // SAFETY: single immutable borrow to `rent_sysvar_info`; account ID and length
         // are checked by `from_account_info_unchecked`.
         let rent = unsafe { Rent::from_account_info_unchecked(rent_sysvar_info)? };
-        rent.is_exempt(mint_info.lamports(), size_of::<Mint>())
+        rent.is_exempt(mint_info.lamports(), mint_data_len)
     } else {
-        Rent::get()?.is_exempt(mint_info.lamports(), size_of::<Mint>())
+        Rent::get()?.is_exempt(mint_info.lamports(), mint_data_len)
     };
 
+    // SAFETY: single mutable borrow to `mint_info` account data.
+    let mint = unsafe { load_mut_unchecked::<Mint>(mint_info.borrow_mut_data_unchecked())? };
+
+    if mint.is_initialized()? {
+        return Err(TokenError::AlreadyInUse.into());
+    }
+
     if !is_exempt {
         return Err(TokenError::NotRentExempt.into());
     }

+ 1 - 1
p-token/src/processor/shared/initialize_multisig.rs

@@ -46,7 +46,7 @@ pub fn process_initialize_multisig(
     let multisig =
         unsafe { load_mut_unchecked::<Multisig>(multisig_info.borrow_mut_data_unchecked())? };
 
-    if multisig.is_initialized() {
+    if multisig.is_initialized()? {
         return Err(TokenError::AlreadyInUse.into());
     }
 

+ 1 - 1
p-token/src/processor/shared/mint_to.rs

@@ -25,7 +25,7 @@ pub fn process_mint_to(
     let destination_account =
         unsafe { load_mut::<Account>(destination_account_info.borrow_mut_data_unchecked())? };
 
-    if destination_account.is_frozen() {
+    if destination_account.is_frozen()? {
         return Err(TokenError::AccountFrozen.into());
     }
 

+ 3 - 3
p-token/src/processor/shared/toggle_account_state.rs

@@ -18,7 +18,7 @@ pub fn process_toggle_account_state(accounts: &[AccountInfo], freeze: bool) -> P
     let source_account =
         unsafe { load_mut::<Account>(source_account_info.borrow_mut_data_unchecked())? };
 
-    if freeze == source_account.is_frozen() {
+    if freeze == source_account.is_frozen()? {
         return Err(TokenError::InvalidState.into());
     }
     if source_account.is_native() {
@@ -39,11 +39,11 @@ pub fn process_toggle_account_state(accounts: &[AccountInfo], freeze: bool) -> P
         None => Err(TokenError::MintCannotFreeze.into()),
     }?;
 
-    source_account.state = if freeze {
+    source_account.set_account_state(if freeze {
         AccountState::Frozen
     } else {
         AccountState::Initialized
-    };
+    });
 
     Ok(())
 }

+ 2 - 2
p-token/src/processor/shared/transfer.rs

@@ -77,7 +77,7 @@ pub fn process_transfer(
     //     destination accounts are not frozen, have the same mint, and the source
     //     account has enough tokens.
     let remaining_amount = if self_transfer {
-        if source_account.is_frozen() {
+        if source_account.is_frozen()? {
             return Err(TokenError::AccountFrozen.into());
         }
 
@@ -92,7 +92,7 @@ pub fn process_transfer(
         let destination_account =
             unsafe { load::<Account>(destination_account_info.borrow_data_unchecked())? };
 
-        if source_account.is_frozen() || destination_account.is_frozen() {
+        if source_account.is_frozen()? || destination_account.is_frozen()? {
             return Err(TokenError::AccountFrozen.into());
         }
 

+ 3 - 7
p-token/src/processor/transfer.rs

@@ -1,15 +1,11 @@
 use {
-    super::shared,
-    pinocchio::{account_info::AccountInfo, program_error::ProgramError, ProgramResult},
+    super::{shared, unpack_amount},
+    pinocchio::{account_info::AccountInfo, ProgramResult},
 };
 
 #[inline(always)]
 pub fn process_transfer(accounts: &[AccountInfo], instruction_data: &[u8]) -> ProgramResult {
-    let amount = u64::from_le_bytes(
-        instruction_data
-            .try_into()
-            .map_err(|_error| ProgramError::InvalidInstructionData)?,
-    );
+    let amount = unpack_amount(instruction_data)?;
 
     shared::transfer::process_transfer(accounts, amount, None)
 }

+ 4 - 13
p-token/src/processor/transfer_checked.rs

@@ -1,6 +1,6 @@
 use {
-    super::shared,
-    pinocchio::{account_info::AccountInfo, program_error::ProgramError, ProgramResult},
+    super::{shared, unpack_amount_and_decimals},
+    pinocchio::{account_info::AccountInfo, ProgramResult},
 };
 
 #[inline(always)]
@@ -8,16 +8,7 @@ pub fn process_transfer_checked(
     accounts: &[AccountInfo],
     instruction_data: &[u8],
 ) -> ProgramResult {
-    // expected u64 (8) + u8 (1)
-    let (amount, decimals) = if instruction_data.len() == 9 {
-        let (amount, decimals) = instruction_data.split_at(core::mem::size_of::<u64>());
-        (
-            u64::from_le_bytes(amount.try_into().unwrap()),
-            decimals.first().copied(),
-        )
-    } else {
-        return Err(ProgramError::InvalidInstructionData);
-    };
+    let (amount, decimals) = unpack_amount_and_decimals(instruction_data)?;
 
-    shared::transfer::process_transfer(accounts, amount, decimals)
+    shared::transfer::process_transfer(accounts, amount, Some(decimals))
 }

+ 1 - 2
p-token/src/processor/ui_amount_to_amount.rs

@@ -15,8 +15,7 @@ pub fn process_ui_amount_to_amount(
     accounts: &[AccountInfo],
     instruction_data: &[u8],
 ) -> ProgramResult {
-    let ui_amount =
-        from_utf8(instruction_data).map_err(|_error| ProgramError::InvalidInstructionData)?;
+    let ui_amount = from_utf8(instruction_data).map_err(|_error| TokenError::InvalidInstruction)?;
 
     let mint_info = accounts.first().ok_or(ProgramError::NotEnoughAccountKeys)?;
     check_account_owner(mint_info)?;