Refactor unpack and make test more robust (#3417)

* Refactor hasty fix to match token-2022

* Make test exhaustive

* cargo fmt

Co-authored-by: Michael Vines <[email protected]>

Author: Tyera Eulberg

token/program/src/instruction.rs

@@ -15,6 +15,8 @@ use std::mem::size_of;
 pub const MIN_SIGNERS: usize = 1;
 /// Maximum number of multisignature signers (max N)
 pub const MAX_SIGNERS: usize = 11;
+/// Serialized length of a u64, for unpacking
+const U64_BYTES: usize = 8;
 
 /// Instructions supported by the token program.
 #[repr(C)]
@@ -519,59 +521,19 @@ impl<'a> TokenInstruction<'a> {
             10 => Self::FreezeAccount,
             11 => Self::ThawAccount,
             12 => {
-                if rest.len() < 8 {
-                    return Err(TokenError::InvalidInstruction.into());
-                }
-                let (amount, rest) = rest.split_at(8);
-                let amount = amount
-                    .try_into()
-                    .ok()
-                    .map(u64::from_le_bytes)
-                    .ok_or(InvalidInstruction)?;
-                let (&decimals, _rest) = rest.split_first().ok_or(InvalidInstruction)?;
-
+                let (amount, decimals, _rest) = Self::unpack_amount_decimals(rest)?;
                 Self::TransferChecked { amount, decimals }
             }
             13 => {
-                if rest.len() < 8 {
-                    return Err(TokenError::InvalidInstruction.into());
-                }
-                let (amount, rest) = rest.split_at(8);
-                let amount = amount
-                    .try_into()
-                    .ok()
-                    .map(u64::from_le_bytes)
-                    .ok_or(InvalidInstruction)?;
-                let (&decimals, _rest) = rest.split_first().ok_or(InvalidInstruction)?;
-
+                let (amount, decimals, _rest) = Self::unpack_amount_decimals(rest)?;
                 Self::ApproveChecked { amount, decimals }
             }
             14 => {
-                if rest.len() < 8 {
-                    return Err(TokenError::InvalidInstruction.into());
-                }
-                let (amount, rest) = rest.split_at(8);
-                let amount = amount
-                    .try_into()
-                    .ok()
-                    .map(u64::from_le_bytes)
-                    .ok_or(InvalidInstruction)?;
-                let (&decimals, _rest) = rest.split_first().ok_or(InvalidInstruction)?;
-
+                let (amount, decimals, _rest) = Self::unpack_amount_decimals(rest)?;
                 Self::MintToChecked { amount, decimals }
             }
             15 => {
-                if rest.len() < 8 {
-                    return Err(TokenError::InvalidInstruction.into());
-                }
-                let (amount, rest) = rest.split_at(8);
-                let amount = amount
-                    .try_into()
-                    .ok()
-                    .map(u64::from_le_bytes)
-                    .ok_or(InvalidInstruction)?;
-                let (&decimals, _rest) = rest.split_first().ok_or(InvalidInstruction)?;
-
+                let (amount, decimals, _rest) = Self::unpack_amount_decimals(rest)?;
                 Self::BurnChecked { amount, decimals }
             }
             16 => {
@@ -600,15 +562,7 @@ impl<'a> TokenInstruction<'a> {
             21 => Self::GetAccountDataSize,
             22 => Self::InitializeImmutableOwner,
             23 => {
-                if rest.len() < 8 {
-                    return Err(TokenError::InvalidInstruction.into());
-                }
-                let (amount, _rest) = rest.split_at(8);
-                let amount = amount
-                    .try_into()
-                    .ok()
-                    .map(u64::from_le_bytes)
-                    .ok_or(InvalidInstruction)?;
+                let (amount, _rest) = Self::unpack_u64(rest)?;
                 Self::AmountToUiAmount { amount }
             }
             24 => {
@@ -760,6 +714,21 @@ impl<'a> TokenInstruction<'a> {
             COption::None => buf.push(0),
         }
     }
+
+    fn unpack_u64(input: &[u8]) -> Result<(u64, &[u8]), ProgramError> {
+        let value = input
+            .get(..U64_BYTES)
+            .and_then(|slice| slice.try_into().ok())
+            .map(u64::from_le_bytes)
+            .ok_or(TokenError::InvalidInstruction)?;
+        Ok((value, &input[U64_BYTES..]))
+    }
+
+    fn unpack_amount_decimals(input: &[u8]) -> Result<(u64, u8, &[u8]), ProgramError> {
+        let (amount, rest) = Self::unpack_u64(input)?;
+        let (&decimals, rest) = rest.split_first().ok_or(TokenError::InvalidInstruction)?;
+        Ok((amount, decimals, rest))
+    }
 }
 
 /// Specifies the authority type for SetAuthority instructions
@@ -1708,8 +1677,11 @@ mod test {
     #[test]
     fn test_instruction_unpack_panic() {
         for i in 0..255u8 {
-            let expect = Vec::from([i, 1, 0, 0, 0, 0, 0, 0, 0, 2]);
-            _ = TokenInstruction::unpack(&expect[0..2]);
+            for j in 1..10 {
+                let mut data = vec![0; j];
+                data[0] = i;
+                let _no_panic = TokenInstruction::unpack(&data);
+            }
         }
     }
 }