From e65a2f9550d06131014e2383719caa2a5a9c42c8 Mon Sep 17 00:00:00 2001 From: Evan Richter Date: Sat, 22 May 2021 04:51:24 -0500 Subject: [PATCH] new Rand method: choose. fixes #117 (#118) The choose method takes an ExactSizeIterator and returns a randomly chosen item from it. Using this method prevents chosing items with an incorrect upper_bound on the index. Various macros help with defining and implementing repetitive mutation strategies. --- libafl/src/bolts/rands.rs | 23 ++ libafl/src/generators/mod.rs | 4 +- libafl/src/mutators/mutations.rs | 662 ++++++++----------------------- 3 files changed, 183 insertions(+), 506 deletions(-) diff --git a/libafl/src/bolts/rands.rs b/libafl/src/bolts/rands.rs index 4430b71d47..0058d6f941 100644 --- a/libafl/src/bolts/rands.rs +++ b/libafl/src/bolts/rands.rs @@ -49,6 +49,29 @@ pub trait Rand: Debug + Serialize + DeserializeOwned { debug_assert!(lower_bound_incl <= upper_bound_incl); lower_bound_incl + self.below(upper_bound_incl - lower_bound_incl + 1) } + + /// Choose an item at random from the given iterator, sampling uniformly. + /// + /// Note: the runtime cost is bound by the iterator's [`nth`][`Iterator::nth`] implementation + /// * For `Vec`, slice, array, this is O(1) + /// * For `HashMap`, `HashSet`, this is O(n) + fn choose(&mut self, from: I) -> T + where + I: IntoIterator, + E: ExactSizeIterator + Iterator, + { + // create iterator + let mut iter = from.into_iter(); + + // make sure there is something to choose from + debug_assert!(iter.len() > 0, "choosing from an empty iterator"); + + // pick a random, valid index + let index = self.below(iter.len() as u64) as usize; + + // return the item chosen + iter.nth(index).unwrap() + } } /// Has a Rand field, that can be used to get random values diff --git a/libafl/src/generators/mod.rs b/libafl/src/generators/mod.rs index e883500338..5ebfd55487 100644 --- a/libafl/src/generators/mod.rs +++ b/libafl/src/generators/mod.rs @@ -86,9 +86,7 @@ where size = 1; } let printables = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz \t\n!\"#$%&'()*+,-./:;<=>?@[\\]^_`{|}~".as_bytes(); - let random_bytes: Vec = (0..size) - .map(|_| printables[rand.below(printables.len() as u64) as usize]) - .collect(); + let random_bytes: Vec = (0..size).map(|_| *rand.choose(printables)).collect(); Ok(BytesInput::new(random_bytes)) } diff --git a/libafl/src/mutators/mutations.rs b/libafl/src/mutators/mutations.rs index 63a7acaaf5..8c4ae524c1 100644 --- a/libafl/src/mutators/mutations.rs +++ b/libafl/src/mutators/mutations.rs @@ -117,11 +117,9 @@ where if input.bytes().is_empty() { Ok(MutationResult::Skipped) } else { - let bit = state.rand_mut().below((input.bytes().len() << 3) as u64) as usize; - unsafe { - // Moar speed, no bound check - *input.bytes_mut().get_unchecked_mut(bit >> 3) ^= (128u8 >> (bit & 7)) as u8; - } + let bit = 1 << state.rand_mut().choose(0..8); + let byte = state.rand_mut().choose(input.bytes_mut()); + *byte ^= bit; Ok(MutationResult::Mutated) } } @@ -179,11 +177,7 @@ where if input.bytes().is_empty() { Ok(MutationResult::Skipped) } else { - let idx = state.rand_mut().below(input.bytes().len() as u64) as usize; - unsafe { - // Moar speed, no bound check - *input.bytes_mut().get_unchecked_mut(idx) ^= 0xff; - } + *state.rand_mut().choose(input.bytes_mut()) ^= 0xff; Ok(MutationResult::Mutated) } } @@ -241,12 +235,8 @@ where if input.bytes().is_empty() { Ok(MutationResult::Skipped) } else { - let idx = state.rand_mut().below(input.bytes().len() as u64) as usize; - unsafe { - // Moar speed, no bound check - let ptr = input.bytes_mut().get_unchecked_mut(idx); - *ptr = (*ptr).wrapping_add(1); - } + let byte = state.rand_mut().choose(input.bytes_mut()); + *byte = byte.wrapping_add(1); Ok(MutationResult::Mutated) } } @@ -304,12 +294,8 @@ where if input.bytes().is_empty() { Ok(MutationResult::Skipped) } else { - let idx = state.rand_mut().below(input.bytes().len() as u64) as usize; - unsafe { - // Moar speed, no bound check - let ptr = input.bytes_mut().get_unchecked_mut(idx); - *ptr = (*ptr).wrapping_sub(1); - } + let byte = state.rand_mut().choose(input.bytes_mut()); + *byte = byte.wrapping_sub(1); Ok(MutationResult::Mutated) } } @@ -367,11 +353,8 @@ where if input.bytes().is_empty() { Ok(MutationResult::Skipped) } else { - let idx = state.rand_mut().below(input.bytes().len() as u64) as usize; - unsafe { - // Moar speed, no bound check - *input.bytes_mut().get_unchecked_mut(idx) = !(*input.bytes().get_unchecked(idx)); - } + let byte = state.rand_mut().choose(input.bytes_mut()); + *byte = !*byte; Ok(MutationResult::Mutated) } } @@ -429,11 +412,8 @@ where if input.bytes().is_empty() { Ok(MutationResult::Skipped) } else { - let idx = state.rand_mut().below(input.bytes().len() as u64) as usize; - unsafe { - // Moar speed, no bound check - *input.bytes_mut().get_unchecked_mut(idx) = state.rand_mut().below(256) as u8; - } + let byte = state.rand_mut().choose(input.bytes_mut()); + *byte = state.rand_mut().next() as u8; Ok(MutationResult::Mutated) } } @@ -465,488 +445,168 @@ where } } -/// Byte add mutation for inputs with a bytes vector -#[derive(Default)] -pub struct ByteAddMutator -where - I: Input + HasBytesVec, - S: HasRand, - R: Rand, -{ - phantom: PhantomData<(I, R, S)>, -} +// Helper macro that defines the arithmetic addition/subtraction mutations where random slices +// within the input are treated as u8, u16, u32, or u64, then mutated in place. +macro_rules! add_mutator_impl { + ($name: ident, $size: ty) => { + /// Adds or subtracts a random value up to `ARITH_MAX` to a [`<$size>`] at a random place in the [`Vec`], in random byte order. + #[derive(Default)] + pub struct $name + where + I: Input + HasBytesVec, + S: HasRand, + R: Rand, + { + phantom: PhantomData<(I, R, S)>, + } -impl Mutator for ByteAddMutator -where - I: Input + HasBytesVec, - S: HasRand, - R: Rand, -{ - fn mutate( - &mut self, - state: &mut S, - input: &mut I, - _stage_idx: i32, - ) -> Result { - if input.bytes().is_empty() { - Ok(MutationResult::Skipped) - } else { - let idx = state.rand_mut().below(input.bytes().len() as u64) as usize; - unsafe { - // Moar speed, no bound check - let ptr = input.bytes_mut().get_unchecked_mut(idx) as *mut u8; - let num = 1 + state.rand_mut().below(ARITH_MAX) as u8; - match state.rand_mut().below(2) { - 0 => *ptr = (*ptr).wrapping_add(num), - _ => *ptr = (*ptr).wrapping_sub(num), - }; + impl Mutator for $name + where + I: Input + HasBytesVec, + S: HasRand, + R: Rand, + { + fn mutate( + &mut self, + state: &mut S, + input: &mut I, + _stage_idx: i32, + ) -> Result { + if input.bytes().len() < size_of::<$size>() { + Ok(MutationResult::Skipped) + } else { + // choose a random window of bytes (windows overlap) and convert to $size + let (index, bytes) = state + .rand_mut() + .choose(input.bytes().windows(size_of::<$size>()).enumerate()); + let val = <$size>::from_ne_bytes(bytes.try_into().unwrap()); + + // mutate + let num = 1 + state.rand_mut().below(ARITH_MAX) as $size; + let new_val = match state.rand_mut().below(4) { + 0 => val.wrapping_add(num), + 1 => val.wrapping_sub(num), + 2 => val.swap_bytes().wrapping_add(num).swap_bytes(), + _ => val.swap_bytes().wrapping_sub(num).swap_bytes(), + }; + + // set bytes to mutated value + let new_bytes = &mut input.bytes_mut()[index..index + size_of::<$size>()]; + new_bytes.copy_from_slice(&new_val.to_ne_bytes()); + Ok(MutationResult::Mutated) + } } - Ok(MutationResult::Mutated) } - } -} -impl Named for ByteAddMutator -where - I: Input + HasBytesVec, - S: HasRand, - R: Rand, -{ - fn name(&self) -> &str { - "ByteAddMutator" - } -} - -impl ByteAddMutator -where - I: Input + HasBytesVec, - S: HasRand, - R: Rand, -{ - /// Creates a new [`ByteAddMutator`]. - #[must_use] - pub fn new() -> Self { - Self { - phantom: PhantomData, - } - } -} - -/// Word add mutation for inputs with a bytes vector -/// adds or subtracts a random value up to [`ARITH_MAX`] to a [`u16`] at a random place in the [`Vec`]. -#[derive(Default)] -pub struct WordAddMutator -where - I: Input + HasBytesVec, - S: HasRand, - R: Rand, -{ - phantom: PhantomData<(I, R, S)>, -} - -impl Mutator for WordAddMutator -where - I: Input + HasBytesVec, - S: HasRand, - R: Rand, -{ - fn mutate( - &mut self, - state: &mut S, - input: &mut I, - _stage_idx: i32, - ) -> Result { - if input.bytes().len() < size_of::() { - Ok(MutationResult::Skipped) - } else { - let bytes = input.bytes_mut(); - let idx = state - .rand_mut() - .below((bytes.len() - size_of::() + 1) as u64) as usize; - let val = u16::from_ne_bytes(bytes[idx..idx + size_of::()].try_into().unwrap()); - let num = 1 + state.rand_mut().below(ARITH_MAX) as u16; - let new_bytes = match state.rand_mut().below(4) { - 0 => val.wrapping_add(num), - 1 => val.wrapping_sub(num), - 2 => val.swap_bytes().wrapping_add(num).swap_bytes(), - _ => val.swap_bytes().wrapping_sub(num).swap_bytes(), + impl Named for $name + where + I: Input + HasBytesVec, + S: HasRand, + R: Rand, + { + fn name(&self) -> &str { + stringify!($name) } - .to_ne_bytes(); - bytes[idx..idx + size_of::()].copy_from_slice(&new_bytes); - Ok(MutationResult::Mutated) } - } -} -impl Named for WordAddMutator -where - I: Input + HasBytesVec, - S: HasRand, - R: Rand, -{ - fn name(&self) -> &str { - "WordAddMutator" - } -} - -impl WordAddMutator -where - I: Input + HasBytesVec, - S: HasRand, - R: Rand, -{ - /// Creates a new [`WordAddMutator`]. - #[must_use] - pub fn new() -> Self { - Self { - phantom: PhantomData, - } - } -} - -/// Dword add mutation for inputs with a bytes vector. -/// Adds a random value up to `ARITH_MAX` to a [`u32`] at a random place in the [`Vec`], in random byte order. -#[derive(Default)] -pub struct DwordAddMutator -where - I: Input + HasBytesVec, - S: HasRand, - R: Rand, -{ - phantom: PhantomData<(I, R, S)>, -} - -impl Mutator for DwordAddMutator -where - I: Input + HasBytesVec, - S: HasRand, - R: Rand, -{ - fn mutate( - &mut self, - state: &mut S, - input: &mut I, - _stage_idx: i32, - ) -> Result { - if input.bytes().len() < size_of::() { - Ok(MutationResult::Skipped) - } else { - let bytes = input.bytes_mut(); - let idx = state - .rand_mut() - .below((bytes.len() - size_of::() + 1) as u64) as usize; - let val = u32::from_ne_bytes(bytes[idx..idx + size_of::()].try_into().unwrap()); - let num = 1 + state.rand_mut().below(ARITH_MAX) as u32; - let new_bytes = match state.rand_mut().below(4) { - 0 => val.wrapping_add(num), - 1 => val.wrapping_sub(num), - 2 => val.swap_bytes().wrapping_add(num).swap_bytes(), - _ => val.swap_bytes().wrapping_sub(num).swap_bytes(), + impl $name + where + I: Input + HasBytesVec, + S: HasRand, + R: Rand, + { + /// Creates a new [`$name`]. + #[must_use] + pub fn new() -> Self { + Self { + phantom: PhantomData, + } } - .to_ne_bytes(); - bytes[idx..idx + size_of::()].copy_from_slice(&new_bytes); - Ok(MutationResult::Mutated) } - } + }; } -impl Named for DwordAddMutator -where - I: Input + HasBytesVec, - S: HasRand, - R: Rand, -{ - fn name(&self) -> &str { - "DwordAddMutator" - } -} +add_mutator_impl!(ByteAddMutator, u8); +add_mutator_impl!(WordAddMutator, u16); +add_mutator_impl!(DwordAddMutator, u32); +add_mutator_impl!(QwordAddMutator, u64); -impl DwordAddMutator -where - I: Input + HasBytesVec, - S: HasRand, - R: Rand, -{ - /// Creates a new [`DwordAddMutator`]. - #[must_use] - pub fn new() -> Self { - Self { - phantom: PhantomData, +/////////////////////////// + +macro_rules! interesting_mutator_impl { + ($name: ident, $size: ty, $interesting: ident) => { + /// Inserts an interesting value at a random place in the input vector + #[derive(Default)] + pub struct $name + where + I: Input + HasBytesVec, + S: HasRand, + R: Rand, + { + phantom: PhantomData<(I, R, S)>, } - } -} -/// Qword add mutation for inputs with a bytes vector -#[derive(Default)] -pub struct QwordAddMutator -where - I: Input + HasBytesVec, - S: HasRand, - R: Rand, -{ - phantom: PhantomData<(I, R, S)>, -} - -impl Mutator for QwordAddMutator -where - I: Input + HasBytesVec, - S: HasRand, - R: Rand, -{ - fn mutate( - &mut self, - state: &mut S, - input: &mut I, - _stage_idx: i32, - ) -> Result { - if input.bytes().len() < 8 { - Ok(MutationResult::Skipped) - } else { - let bytes = input.bytes_mut(); - let idx = state - .rand_mut() - .below((bytes.len() - size_of::() + 1) as u64) as usize; - let val = u64::from_ne_bytes(bytes[idx..idx + size_of::()].try_into().unwrap()); - let num = 1 + state.rand_mut().below(ARITH_MAX) as u64; - let new_bytes = match state.rand_mut().below(4) { - 0 => val.wrapping_add(num), - 1 => val.wrapping_sub(num), - 2 => val.swap_bytes().wrapping_add(num).swap_bytes(), - _ => val.swap_bytes().wrapping_sub(num).swap_bytes(), + impl Mutator for $name + where + I: Input + HasBytesVec, + S: HasRand, + R: Rand, + { + #[allow(clippy::cast_sign_loss)] + fn mutate( + &mut self, + state: &mut S, + input: &mut I, + _stage_idx: i32, + ) -> Result { + if input.bytes().len() < size_of::<$size>() { + Ok(MutationResult::Skipped) + } else { + let bytes = input.bytes_mut(); + let upper_bound = (bytes.len() + 1 - size_of::<$size>()) as u64; + let idx = state.rand_mut().below(upper_bound) as usize; + let val = *state.rand_mut().choose(&$interesting) as $size; + let new_bytes = match state.rand_mut().choose(&[0, 1]) { + 0 => val.to_be_bytes(), + _ => val.to_le_bytes(), + }; + bytes[idx..idx + size_of::<$size>()].copy_from_slice(&new_bytes); + Ok(MutationResult::Mutated) + } } - .to_ne_bytes(); - bytes[idx..idx + size_of::()].copy_from_slice(&new_bytes); - Ok(MutationResult::Mutated) } - } -} -impl Named for QwordAddMutator -where - I: Input + HasBytesVec, - S: HasRand, - R: Rand, -{ - fn name(&self) -> &str { - "QwordAddMutator" - } -} - -impl QwordAddMutator -where - I: Input + HasBytesVec, - S: HasRand, - R: Rand, -{ - /// Creates a new [`QwordAddMutator`]. - #[must_use] - pub fn new() -> Self { - Self { - phantom: PhantomData, - } - } -} - -/// Byte interesting mutation for inputs with a bytes vector -#[derive(Default)] -pub struct ByteInterestingMutator -where - I: Input + HasBytesVec, - S: HasRand, - R: Rand, -{ - phantom: PhantomData<(I, R, S)>, -} - -impl Mutator for ByteInterestingMutator -where - I: Input + HasBytesVec, - S: HasRand, - R: Rand, -{ - #[allow(clippy::cast_sign_loss)] - fn mutate( - &mut self, - state: &mut S, - input: &mut I, - _stage_idx: i32, - ) -> Result { - if input.bytes().is_empty() { - Ok(MutationResult::Skipped) - } else { - let idx = state.rand_mut().below(input.bytes().len() as u64) as usize; - let val = - INTERESTING_8[state.rand_mut().below(INTERESTING_8.len() as u64) as usize] as u8; - unsafe { - // Moar speed, no bound check - *input.bytes_mut().get_unchecked_mut(idx) = val; + impl Named for $name + where + I: Input + HasBytesVec, + S: HasRand, + R: Rand, + { + fn name(&self) -> &str { + stringify!($name) } - Ok(MutationResult::Mutated) } - } -} -impl Named for ByteInterestingMutator -where - I: Input + HasBytesVec, - S: HasRand, - R: Rand, -{ - fn name(&self) -> &str { - "ByteInterestingMutator" - } -} - -impl ByteInterestingMutator -where - I: Input + HasBytesVec, - S: HasRand, - R: Rand, -{ - /// Creates a new [`ByteInterestingMutator`]. - #[must_use] - pub fn new() -> Self { - Self { - phantom: PhantomData, + impl $name + where + I: Input + HasBytesVec, + S: HasRand, + R: Rand, + { + /// Creates a new [`$name`]. + #[must_use] + pub fn new() -> Self { + Self { + phantom: PhantomData, + } + } } - } + }; } -/// Word interesting mutation for inputs with a bytes vector -#[derive(Default)] -pub struct WordInterestingMutator -where - I: Input + HasBytesVec, - S: HasRand, - R: Rand, -{ - phantom: PhantomData<(I, R, S)>, -} - -impl Mutator for WordInterestingMutator -where - I: Input + HasBytesVec, - S: HasRand, - R: Rand, -{ - #[allow(clippy::cast_sign_loss)] - fn mutate( - &mut self, - state: &mut S, - input: &mut I, - _stage_idx: i32, - ) -> Result { - if input.bytes().len() < 2 { - Ok(MutationResult::Skipped) - } else { - let bytes = input.bytes_mut(); - let idx = state.rand_mut().below(bytes.len() as u64 - 1) as usize; - let val = - INTERESTING_16[state.rand_mut().below(INTERESTING_8.len() as u64) as usize] as u16; - let new_bytes = if state.rand_mut().below(2) == 0 { - val.to_be_bytes() - } else { - val.to_le_bytes() - }; - bytes[idx..idx + size_of::()].copy_from_slice(&new_bytes); - Ok(MutationResult::Mutated) - } - } -} - -impl Named for WordInterestingMutator -where - I: Input + HasBytesVec, - S: HasRand, - R: Rand, -{ - fn name(&self) -> &str { - "WordInterestingMutator" - } -} - -impl WordInterestingMutator -where - I: Input + HasBytesVec, - S: HasRand, - R: Rand, -{ - /// Creates a new [`WordInterestingMutator`]. - #[must_use] - pub fn new() -> Self { - Self { - phantom: PhantomData, - } - } -} - -/// Dword interesting mutation for inputs with a bytes vector -#[derive(Default)] -pub struct DwordInterestingMutator -where - I: Input + HasBytesVec, - S: HasRand, - R: Rand, -{ - phantom: PhantomData<(I, R, S)>, -} - -impl Mutator for DwordInterestingMutator -where - I: Input + HasBytesVec, - S: HasRand, - R: Rand, -{ - #[allow(clippy::cast_sign_loss)] - fn mutate( - &mut self, - state: &mut S, - input: &mut I, - _stage_idx: i32, - ) -> Result { - if input.bytes().len() < 4 { - Ok(MutationResult::Skipped) - } else { - let bytes = input.bytes_mut(); - let idx = state.rand_mut().below(bytes.len() as u64 - 3) as usize; - let val = - INTERESTING_32[state.rand_mut().below(INTERESTING_8.len() as u64) as usize] as u32; - let new_bytes = if state.rand_mut().below(2) == 0 { - val.to_be_bytes() - } else { - val.to_le_bytes() - }; - bytes[idx..idx + new_bytes.len()].copy_from_slice(&new_bytes); - Ok(MutationResult::Mutated) - } - } -} - -impl Named for DwordInterestingMutator -where - I: Input + HasBytesVec, - S: HasRand, - R: Rand, -{ - fn name(&self) -> &str { - "DwordInterestingMutator" - } -} - -impl DwordInterestingMutator -where - I: Input + HasBytesVec, - S: HasRand, - R: Rand, -{ - /// Creates a new [`DwordInterestingMutator`]. - #[must_use] - pub fn new() -> Self { - Self { - phantom: PhantomData, - } - } -} +interesting_mutator_impl!(ByteInterestingMutator, u8, INTERESTING_8); +interesting_mutator_impl!(WordInterestingMutator, u16, INTERESTING_16); +interesting_mutator_impl!(DwordInterestingMutator, u32, INTERESTING_32); /// Bytes delete mutation for inputs with a bytes vector #[derive(Default)] @@ -1190,7 +850,7 @@ where } } - let val = state.rand_mut().below(256) as u8; + let val = state.rand_mut().next() as u8; input.bytes_mut().resize(size + len, 0); buffer_self_copy(input.bytes_mut(), off, off + len, size - off); @@ -1256,7 +916,7 @@ where let off = state.rand_mut().below(size as u64) as usize; let len = 1 + state.rand_mut().below(min(16, size - off) as u64) as usize; - let val = input.bytes()[state.rand_mut().below(size as u64) as usize]; + let val = *state.rand_mut().choose(input.bytes()); buffer_set(input.bytes_mut(), off, len, val); @@ -1320,7 +980,7 @@ where let off = state.rand_mut().below(size as u64) as usize; let len = 1 + state.rand_mut().below(min(16, size - off) as u64) as usize; - let val = state.rand_mut().below(256) as u8; + let val = state.rand_mut().next() as u8; buffer_set(input.bytes_mut(), off, len, val); @@ -1782,16 +1442,12 @@ where // Converts a hex u8 to its u8 value: 'A' -> 10 etc. fn from_hex(hex: u8) -> Result { - if (48..=57).contains(&hex) { - return Ok(hex - 48); + match hex { + 48..=57 => Ok(hex - 48), + 65..=70 => Ok(hex - 55), + 97..=102 => Ok(hex - 87), + _ => Err(Error::IllegalArgument("Invalid hex character".to_owned())), } - if (65..=70).contains(&hex) { - return Ok(hex - 55); - } - if (97..=102).contains(&hex) { - return Ok(hex - 87); - } - Err(Error::IllegalArgument("".to_owned())) } /// Decodes a dictionary token: 'foo\x41\\and\"bar' -> 'fooA\and"bar'