From 4a56a4b8402c568e0f37f4c37e935b2fbc2eaee1 Mon Sep 17 00:00:00 2001 From: Tyera Eulberg Date: Wed, 3 Aug 2022 23:25:07 -0700 Subject: [PATCH] Prep for spl-token v3.5.0 (#3426) * check that unpack is tolerant of small sizes (#3416) * 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 * Readd proptests without losing unit test, #3421 Co-authored-by: anatoly yakovenko Co-authored-by: Michael Vines --- token/program-2022/src/instruction.rs | 12 +++- token/program/src/instruction.rs | 83 ++++++++++++++------------- 2 files changed, 55 insertions(+), 40 deletions(-) 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); + } + } }