diff --git a/token/program-2022/src/instruction.rs b/token/program-2022/src/instruction.rs index e948a9991f3..116b16097ec 100644 --- a/token/program-2022/src/instruction.rs +++ b/token/program-2022/src/instruction.rs @@ -1787,7 +1787,7 @@ pub(crate) fn encode_instruction, D: Pod>( #[cfg(test)] mod test { - use super::*; + use {super::*, proptest::prelude::*}; #[test] fn test_instruction_packing() { @@ -2284,4 +2284,14 @@ mod test { ui_amount, )); } + + proptest! { + #![proptest_config(ProptestConfig::with_cases(1024))] + #[test] + fn test_instruction_unpack_proptest( + data in prop::collection::vec(any::(), 0..255) + ) { + let _no_panic = TokenInstruction::unpack(&data); + } + } } diff --git a/token/program/src/instruction.rs b/token/program/src/instruction.rs index c69c7bbe2eb..9dccdfba42c 100644 --- a/token/program/src/instruction.rs +++ b/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,47 +521,19 @@ impl<'a> TokenInstruction<'a> { 10 => Self::FreezeAccount, 11 => Self::ThawAccount, 12 => { - 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 => { - 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 => { - 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 => { - 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 => { @@ -588,12 +562,7 @@ impl<'a> TokenInstruction<'a> { 21 => Self::GetAccountDataSize, 22 => Self::InitializeImmutableOwner, 23 => { - 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 => { @@ -745,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 @@ -1447,7 +1431,7 @@ pub fn is_valid_signer_index(index: usize) -> bool { #[cfg(test)] mod test { - use super::*; + use {super::*, proptest::prelude::*}; #[test] fn test_instruction_packing() { @@ -1689,4 +1673,25 @@ mod test { let unpacked = TokenInstruction::unpack(&expect).unwrap(); assert_eq!(unpacked, check); } + + #[test] + fn test_instruction_unpack_panic() { + for i in 0..255u8 { + for j in 1..10 { + let mut data = vec![0; j]; + data[0] = i; + let _no_panic = TokenInstruction::unpack(&data); + } + } + } + + proptest! { + #![proptest_config(ProptestConfig::with_cases(1024))] + #[test] + fn test_instruction_unpack_proptest( + data in prop::collection::vec(any::(), 0..255) + ) { + let _no_panic = TokenInstruction::unpack(&data); + } + } }