Browse Source

Relax instruction data length check

febo 5 months ago
parent
commit
16d3839722

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

@@ -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 => {

+ 7 - 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, MAX_FORMATTED_DIGITS, U64_BYTES},
     core::str::from_utf8_unchecked,
     pinocchio::{
         account_info::AccountInfo, program::set_return_data, program_error::ProgramError,
@@ -16,11 +16,12 @@ 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| TokenError::InvalidInstruction)?,
-    );
+    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 mint_info = accounts.first().ok_or(ProgramError::NotEnoughAccountKeys)?;
     check_account_owner(mint_info)?;

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

@@ -1,16 +1,17 @@
 use {
-    super::shared,
+    super::{shared, U64_BYTES},
     pinocchio::{account_info::AccountInfo, ProgramResult},
     spl_token_interface::error::TokenError,
 };
 
 #[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| TokenError::InvalidInstruction)?,
-    );
+    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());
+    };
 
     shared::approve::process_approve(accounts, amount, None)
 }

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

@@ -1,5 +1,5 @@
 use {
-    super::shared,
+    super::{shared, U64_BYTES},
     pinocchio::{account_info::AccountInfo, ProgramResult},
     spl_token_interface::error::TokenError,
 };
@@ -8,9 +8,10 @@ use {
 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>());
+        let (amount, decimals) = instruction_data.split_at(U64_BYTES);
         (
-            u64::from_le_bytes(amount.try_into().unwrap()),
+            // 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 {

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

@@ -1,16 +1,17 @@
 use {
-    super::shared,
+    super::{shared, U64_BYTES},
     pinocchio::{account_info::AccountInfo, ProgramResult},
     spl_token_interface::error::TokenError,
 };
 
 #[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| TokenError::InvalidInstruction)?,
-    );
+    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());
+    };
 
     shared::burn::process_burn(accounts, amount, None)
 }

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

@@ -1,5 +1,5 @@
 use {
-    super::shared,
+    super::{shared, U64_BYTES},
     pinocchio::{account_info::AccountInfo, ProgramResult},
     spl_token_interface::error::TokenError,
 };
@@ -8,9 +8,10 @@ use {
 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>());
+        let (amount, decimals) = instruction_data.split_at(U64_BYTES);
         (
-            u64::from_le_bytes(amount.try_into().unwrap()),
+            // 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 {

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

@@ -1,6 +1,10 @@
 use {
     super::shared,
-    pinocchio::{account_info::AccountInfo, pubkey::Pubkey, ProgramResult},
+    pinocchio::{
+        account_info::AccountInfo,
+        pubkey::{Pubkey, PUBKEY_BYTES},
+        ProgramResult,
+    },
     spl_token_interface::error::TokenError,
 };
 
@@ -9,9 +13,12 @@ pub fn process_initialize_account2(
     accounts: &[AccountInfo],
     instruction_data: &[u8],
 ) -> ProgramResult {
-    let owner: &Pubkey = instruction_data
-        .try_into()
-        .map_err(|_error| TokenError::InvalidInstruction)?;
+    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)
 }

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

@@ -1,6 +1,10 @@
 use {
     super::shared,
-    pinocchio::{account_info::AccountInfo, pubkey::Pubkey, ProgramResult},
+    pinocchio::{
+        account_info::AccountInfo,
+        pubkey::{Pubkey, PUBKEY_BYTES},
+        ProgramResult,
+    },
     spl_token_interface::error::TokenError,
 };
 
@@ -9,9 +13,12 @@ pub fn process_initialize_account3(
     accounts: &[AccountInfo],
     instruction_data: &[u8],
 ) -> ProgramResult {
-    let owner: &Pubkey = instruction_data
-        .try_into()
-        .map_err(|_error| TokenError::InvalidInstruction)?;
+    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 - 0
p-token/src/processor/initialize_multisig2.rs

@@ -11,5 +11,6 @@ pub fn process_initialize_multisig2(
     let m = instruction_data
         .first()
         .ok_or(TokenError::InvalidInstruction)?;
+
     shared::initialize_multisig::process_initialize_multisig(accounts, *m, false)
 }

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

@@ -1,16 +1,17 @@
 use {
-    super::shared,
+    super::{shared, U64_BYTES},
     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 = u64::from_le_bytes(
-        instruction_data
-            .try_into()
-            .map_err(|_error| TokenError::InvalidInstruction)?,
-    );
+    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());
+    };
 
     shared::mint_to::process_mint_to(accounts, amount, None)
 }

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

@@ -1,5 +1,5 @@
 use {
-    super::shared,
+    super::{shared, U64_BYTES},
     pinocchio::{account_info::AccountInfo, ProgramResult},
     spl_token_interface::error::TokenError,
 };
@@ -8,9 +8,10 @@ use {
 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>());
+        let (amount, decimals) = instruction_data.split_at(U64_BYTES);
         (
-            u64::from_le_bytes(amount.try_into().unwrap()),
+            // 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 {

+ 3 - 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

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

@@ -8,7 +8,7 @@ use {
 };
 
 #[inline(always)]
-pub fn process_revoke(accounts: &[AccountInfo], _instruction_data: &[u8]) -> ProgramResult {
+pub fn process_revoke(accounts: &[AccountInfo]) -> ProgramResult {
     let [source_account_info, owner_info, remaining @ ..] = accounts else {
         return Err(ProgramError::NotEnoughAccountKeys);
     };

+ 16 - 14
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(TokenError::InvalidInstruction.into()),
+    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.

+ 19 - 17
p-token/src/processor/shared/initialize_mint.rs

@@ -21,24 +21,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(TokenError::InvalidInstruction.into()),
+    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.

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

@@ -1,16 +1,17 @@
 use {
-    super::shared,
+    super::{shared, U64_BYTES},
     pinocchio::{account_info::AccountInfo, ProgramResult},
     spl_token_interface::error::TokenError,
 };
 
 #[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| TokenError::InvalidInstruction)?,
-    );
+    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());
+    };
 
     shared::transfer::process_transfer(accounts, amount, None)
 }

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

@@ -1,5 +1,5 @@
 use {
-    super::shared,
+    super::{shared, U64_BYTES},
     pinocchio::{account_info::AccountInfo, ProgramResult},
     spl_token_interface::error::TokenError,
 };
@@ -11,9 +11,10 @@ pub fn process_transfer_checked(
 ) -> 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>());
+        let (amount, decimals) = instruction_data.split_at(U64_BYTES);
         (
-            u64::from_le_bytes(amount.try_into().unwrap()),
+            // 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 {