Add unpack helpers

Author: febo

p-token/src/processor/amount_to_ui_amount.rs

@@ -1,5 +1,5 @@
 use {
-    super::{check_account_owner, MAX_FORMATTED_DIGITS, U64_BYTES},
+    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,12 +16,7 @@ pub fn process_amount_to_ui_amount(
     accounts: &[AccountInfo],
     instruction_data: &[u8],
 ) -> ProgramResult {
-    let amount = if instruction_data.len() >= U64_BYTES {
-        // SAFETY: The minimum size of the instruction data is `U64_BYTES` bytes.
-        unsafe { u64::from_le_bytes(*(instruction_data.as_ptr() as *const [u8; U64_BYTES])) }
-    } else {
-        return Err(TokenError::InvalidInstruction.into());
-    };
+    let amount = unpack_amount(instruction_data)?;
 
     let mint_info = accounts.first().ok_or(ProgramError::NotEnoughAccountKeys)?;
     check_account_owner(mint_info)?;

p-token/src/processor/approve.rs

@@ -1,17 +1,11 @@
 use {
-    super::{shared, U64_BYTES},
+    super::{shared, unpack_amount},
     pinocchio::{account_info::AccountInfo, ProgramResult},
-    spl_token_interface::error::TokenError,
 };
 
 #[inline(always)]
 pub fn process_approve(accounts: &[AccountInfo], instruction_data: &[u8]) -> ProgramResult {
-    let amount = if instruction_data.len() >= U64_BYTES {
-        // SAFETY: The minimum size of the instruction data is `U64_BYTES` bytes.
-        unsafe { u64::from_le_bytes(*(instruction_data.as_ptr() as *const [u8; U64_BYTES])) }
-    } else {
-        return Err(TokenError::InvalidInstruction.into());
-    };
+    let amount = unpack_amount(instruction_data)?;
 
     shared::approve::process_approve(accounts, amount, None)
 }

p-token/src/processor/approve_checked.rs

@@ -1,22 +1,11 @@
 use {
-    super::{shared, U64_BYTES},
+    super::{shared, unpack_amount_and_decimals},
     pinocchio::{account_info::AccountInfo, ProgramResult},
-    spl_token_interface::error::TokenError,
 };
 
 #[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(U64_BYTES);
-        (
-            // SAFETY: The size of `amount` is `U64_BYTES` bytes.
-            unsafe { u64::from_le_bytes(*(amount.as_ptr() as *const [u8; U64_BYTES])) },
-            decimals.first().copied(),
-        )
-    } else {
-        return Err(TokenError::InvalidInstruction.into());
-    };
+    let (amount, decimals) = unpack_amount_and_decimals(instruction_data)?;
 
-    shared::approve::process_approve(accounts, amount, decimals)
+    shared::approve::process_approve(accounts, amount, Some(decimals))
 }

p-token/src/processor/burn.rs

@@ -1,17 +1,11 @@
 use {
-    super::{shared, U64_BYTES},
+    super::{shared, unpack_amount},
     pinocchio::{account_info::AccountInfo, ProgramResult},
-    spl_token_interface::error::TokenError,
 };
 
 #[inline(always)]
 pub fn process_burn(accounts: &[AccountInfo], instruction_data: &[u8]) -> ProgramResult {
-    let amount = if instruction_data.len() >= U64_BYTES {
-        // SAFETY: The minimum size of the instruction data is `U64_BYTES` bytes.
-        unsafe { u64::from_le_bytes(*(instruction_data.as_ptr() as *const [u8; U64_BYTES])) }
-    } else {
-        return Err(TokenError::InvalidInstruction.into());
-    };
+    let amount = unpack_amount(instruction_data)?;
 
     shared::burn::process_burn(accounts, amount, None)
 }

p-token/src/processor/burn_checked.rs

@@ -1,22 +1,11 @@
 use {
-    super::{shared, U64_BYTES},
+    super::{shared, unpack_amount_and_decimals},
     pinocchio::{account_info::AccountInfo, ProgramResult},
-    spl_token_interface::error::TokenError,
 };
 
 #[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(U64_BYTES);
-        (
-            // SAFETY: The size of `amount` is `U64_BYTES` bytes.
-            unsafe { u64::from_le_bytes(*(amount.as_ptr() as *const [u8; U64_BYTES])) },
-            decimals.first().copied(),
-        )
-    } else {
-        return Err(TokenError::InvalidInstruction.into());
-    };
+    let (amount, decimals) = unpack_amount_and_decimals(instruction_data)?;
 
-    shared::burn::process_burn(accounts, amount, decimals)
+    shared::burn::process_burn(accounts, amount, Some(decimals))
 }

p-token/src/processor/mint_to.rs

@@ -1,17 +1,11 @@
 use {
-    super::{shared, U64_BYTES},
+    super::{shared, unpack_amount},
     pinocchio::{account_info::AccountInfo, ProgramResult},
-    spl_token_interface::error::TokenError,
 };
 
 #[inline(always)]
 pub fn process_mint_to(accounts: &[AccountInfo], instruction_data: &[u8]) -> ProgramResult {
-    let amount = if instruction_data.len() >= U64_BYTES {
-        // SAFETY: The minimum size of the instruction data is `U64_BYTES` bytes.
-        unsafe { u64::from_le_bytes(*(instruction_data.as_ptr() as *const [u8; U64_BYTES])) }
-    } else {
-        return Err(TokenError::InvalidInstruction.into());
-    };
+    let amount = unpack_amount(instruction_data)?;
 
     shared::mint_to::process_mint_to(accounts, amount, None)
 }

p-token/src/processor/mint_to_checked.rs

@@ -1,22 +1,11 @@
 use {
-    super::{shared, U64_BYTES},
+    super::{shared, unpack_amount_and_decimals},
     pinocchio::{account_info::AccountInfo, ProgramResult},
-    spl_token_interface::error::TokenError,
 };
 
 #[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(U64_BYTES);
-        (
-            // SAFETY: The size of `amount` is `U64_BYTES` bytes.
-            unsafe { u64::from_le_bytes(*(amount.as_ptr() as *const [u8; U64_BYTES])) },
-            decimals.first().copied(),
-        )
-    } else {
-        return Err(TokenError::InvalidInstruction.into());
-    };
+    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))
 }

p-token/src/processor/mod.rs

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

p-token/src/processor/transfer.rs

@@ -1,17 +1,11 @@
 use {
-    super::{shared, U64_BYTES},
+    super::{shared, unpack_amount},
     pinocchio::{account_info::AccountInfo, ProgramResult},
-    spl_token_interface::error::TokenError,
 };
 
 #[inline(always)]
 pub fn process_transfer(accounts: &[AccountInfo], instruction_data: &[u8]) -> ProgramResult {
-    let amount = if instruction_data.len() >= U64_BYTES {
-        // SAFETY: The minimum size of the instruction data is `U64_BYTES` bytes.
-        unsafe { u64::from_le_bytes(*(instruction_data.as_ptr() as *const [u8; U64_BYTES])) }
-    } else {
-        return Err(TokenError::InvalidInstruction.into());
-    };
+    let amount = unpack_amount(instruction_data)?;
 
     shared::transfer::process_transfer(accounts, amount, None)
 }

p-token/src/processor/transfer_checked.rs

@@ -1,25 +1,14 @@
 use {
-    super::{shared, U64_BYTES},
+    super::{shared, unpack_amount_and_decimals},
     pinocchio::{account_info::AccountInfo, ProgramResult},
-    spl_token_interface::error::TokenError,
 };
 
 #[inline(always)]
 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(U64_BYTES);
-        (
-            // SAFETY: The size of `amount` is `U64_BYTES` bytes.
-            unsafe { u64::from_le_bytes(*(amount.as_ptr() as *const [u8; U64_BYTES])) },
-            decimals.first().copied(),
-        )
-    } else {
-        return Err(TokenError::InvalidInstruction.into());
-    };
+    let (amount, decimals) = unpack_amount_and_decimals(instruction_data)?;
 
-    shared::transfer::process_transfer(accounts, amount, decimals)
+    shared::transfer::process_transfer(accounts, amount, Some(decimals))
 }