diff --git a/docs/listings/baby_fuzzer/listing-04/src/main.rs b/docs/listings/baby_fuzzer/listing-04/src/main.rs index dd64597244..b7ab1b484b 100644 --- a/docs/listings/baby_fuzzer/listing-04/src/main.rs +++ b/docs/listings/baby_fuzzer/listing-04/src/main.rs @@ -77,7 +77,7 @@ fn main() { /* ANCHOR: generator */ // Generator of printable bytearrays of max size 32 - let mut generator = RandPrintablesGenerator::new(32); + let mut generator = RandPrintablesGenerator::new(32).unwrap(); // Generate 8 initial inputs state diff --git a/docs/listings/baby_fuzzer/listing-05/src/main.rs b/docs/listings/baby_fuzzer/listing-05/src/main.rs index 1377243fb7..6636c1cdf0 100644 --- a/docs/listings/baby_fuzzer/listing-05/src/main.rs +++ b/docs/listings/baby_fuzzer/listing-05/src/main.rs @@ -105,7 +105,7 @@ fn main() { /* ANCHOR_END: executor_with_observer */ // Generator of printable bytearrays of max size 32 - let mut generator = RandPrintablesGenerator::new(32); + let mut generator = RandPrintablesGenerator::new(32).unwrap(); // Generate 8 initial inputs state diff --git a/docs/listings/baby_fuzzer/listing-06/src/main.rs b/docs/listings/baby_fuzzer/listing-06/src/main.rs index 6faff74382..091c48b211 100644 --- a/docs/listings/baby_fuzzer/listing-06/src/main.rs +++ b/docs/listings/baby_fuzzer/listing-06/src/main.rs @@ -97,7 +97,7 @@ fn main() { .expect("Failed to create the Executor"); // Generator of printable bytearrays of max size 32 - let mut generator = RandPrintablesGenerator::new(32); + let mut generator = RandPrintablesGenerator::new(32).unwrap(); // Generate 8 initial inputs state diff --git a/fuzzers/baby/baby_fuzzer/src/main.rs b/fuzzers/baby/baby_fuzzer/src/main.rs index 36f7ca94bd..124af48abf 100644 --- a/fuzzers/baby/baby_fuzzer/src/main.rs +++ b/fuzzers/baby/baby_fuzzer/src/main.rs @@ -116,7 +116,7 @@ pub fn main() { .expect("Failed to create the Executor"); // Generator of printable bytearrays of max size 32 - let mut generator = RandPrintablesGenerator::new(32); + let mut generator = RandPrintablesGenerator::new(32).unwrap(); // Generate 8 initial inputs state diff --git a/fuzzers/baby/baby_fuzzer_minimizing/src/main.rs b/fuzzers/baby/baby_fuzzer_minimizing/src/main.rs index 30e920d380..375071a7a5 100644 --- a/fuzzers/baby/baby_fuzzer_minimizing/src/main.rs +++ b/fuzzers/baby/baby_fuzzer_minimizing/src/main.rs @@ -86,7 +86,7 @@ pub fn main() -> Result<(), Error> { .expect("Failed to create the Executor"); // Generator of printable bytearrays of max size 32 - let mut generator = RandPrintablesGenerator::new(32); + let mut generator = RandPrintablesGenerator::new(32).unwrap(); // Generate 8 initial inputs state diff --git a/fuzzers/baby/baby_fuzzer_swap_differential/src/main.rs b/fuzzers/baby/baby_fuzzer_swap_differential/src/main.rs index 471d6fee31..26a6a372c1 100644 --- a/fuzzers/baby/baby_fuzzer_swap_differential/src/main.rs +++ b/fuzzers/baby/baby_fuzzer_swap_differential/src/main.rs @@ -247,7 +247,7 @@ pub fn main() { ); // Generator of printable bytearrays of max size 32 - let mut generator = RandPrintablesGenerator::new(32); + let mut generator = RandPrintablesGenerator::new(32).unwrap(); // Generate 8 initial inputs state diff --git a/fuzzers/baby/backtrace_baby_fuzzers/c_code_with_fork_executor/src/main.rs b/fuzzers/baby/backtrace_baby_fuzzers/c_code_with_fork_executor/src/main.rs index ca1d87f5d1..97e3951f63 100644 --- a/fuzzers/baby/backtrace_baby_fuzzers/c_code_with_fork_executor/src/main.rs +++ b/fuzzers/baby/backtrace_baby_fuzzers/c_code_with_fork_executor/src/main.rs @@ -103,7 +103,7 @@ pub fn main() { .expect("Failed to create the Executor"); // Generator of printable bytearrays of max size 32 - let mut generator = RandPrintablesGenerator::new(32); + let mut generator = RandPrintablesGenerator::new(32).unwrap(); // Generate 8 initial inputs state diff --git a/fuzzers/baby/backtrace_baby_fuzzers/c_code_with_inprocess_executor/src/main.rs b/fuzzers/baby/backtrace_baby_fuzzers/c_code_with_inprocess_executor/src/main.rs index 9dbdb61dc9..84808fbe67 100644 --- a/fuzzers/baby/backtrace_baby_fuzzers/c_code_with_inprocess_executor/src/main.rs +++ b/fuzzers/baby/backtrace_baby_fuzzers/c_code_with_inprocess_executor/src/main.rs @@ -89,7 +89,7 @@ pub fn main() { .expect("Failed to create the Executor"); // Generator of printable bytearrays of max size 32 - let mut generator = RandPrintablesGenerator::new(32); + let mut generator = RandPrintablesGenerator::new(32).unwrap(); // Generate 8 initial inputs state diff --git a/fuzzers/baby/backtrace_baby_fuzzers/command_executor/src/main.rs b/fuzzers/baby/backtrace_baby_fuzzers/command_executor/src/main.rs index 65a3e282fe..6498751e3b 100644 --- a/fuzzers/baby/backtrace_baby_fuzzers/command_executor/src/main.rs +++ b/fuzzers/baby/backtrace_baby_fuzzers/command_executor/src/main.rs @@ -112,7 +112,7 @@ pub fn main() { let mut executor = MyExecutor { shmem_id }.into_executor(tuple_list!(observer, bt_observer)); // Generator of printable bytearrays of max size 32 - let mut generator = RandPrintablesGenerator::new(32); + let mut generator = RandPrintablesGenerator::new(32).unwrap(); // Generate 8 initial inputs state diff --git a/fuzzers/baby/backtrace_baby_fuzzers/forkserver_executor/src/main.rs b/fuzzers/baby/backtrace_baby_fuzzers/forkserver_executor/src/main.rs index e0322e930a..3cae2a0347 100644 --- a/fuzzers/baby/backtrace_baby_fuzzers/forkserver_executor/src/main.rs +++ b/fuzzers/baby/backtrace_baby_fuzzers/forkserver_executor/src/main.rs @@ -97,7 +97,7 @@ pub fn main() { .unwrap(); // Generator of printable bytearrays of max size 32 - let mut generator = RandPrintablesGenerator::new(3); + let mut generator = RandPrintablesGenerator::new(3).unwrap(); // Generate 8 initial inputs state diff --git a/fuzzers/baby/backtrace_baby_fuzzers/rust_code_with_fork_executor/src/main.rs b/fuzzers/baby/backtrace_baby_fuzzers/rust_code_with_fork_executor/src/main.rs index c88ca60952..704b3de42c 100644 --- a/fuzzers/baby/backtrace_baby_fuzzers/rust_code_with_fork_executor/src/main.rs +++ b/fuzzers/baby/backtrace_baby_fuzzers/rust_code_with_fork_executor/src/main.rs @@ -122,7 +122,7 @@ pub fn main() { .expect("Failed to create the Executor"); // Generator of printable bytearrays of max size 32 - let mut generator = RandPrintablesGenerator::new(32); + let mut generator = RandPrintablesGenerator::new(32).unwrap(); // Generate 8 initial inputs state diff --git a/fuzzers/baby/backtrace_baby_fuzzers/rust_code_with_inprocess_executor/src/main.rs b/fuzzers/baby/backtrace_baby_fuzzers/rust_code_with_inprocess_executor/src/main.rs index ed7ffe4e5d..b19c6751c1 100644 --- a/fuzzers/baby/backtrace_baby_fuzzers/rust_code_with_inprocess_executor/src/main.rs +++ b/fuzzers/baby/backtrace_baby_fuzzers/rust_code_with_inprocess_executor/src/main.rs @@ -113,7 +113,7 @@ pub fn main() { .expect("Failed to create the Executor"); // Generator of printable bytearrays of max size 32 - let mut generator = RandPrintablesGenerator::new(32); + let mut generator = RandPrintablesGenerator::new(32).unwrap(); // Generate 8 initial inputs state diff --git a/fuzzers/forkserver/baby_fuzzer_with_forkexecutor/src/main.rs b/fuzzers/forkserver/baby_fuzzer_with_forkexecutor/src/main.rs index 962f3e0d04..ab17fd2679 100644 --- a/fuzzers/forkserver/baby_fuzzer_with_forkexecutor/src/main.rs +++ b/fuzzers/forkserver/baby_fuzzer_with_forkexecutor/src/main.rs @@ -115,7 +115,7 @@ pub fn main() { .expect("Failed to create the Executor"); // Generator of printable bytearrays of max size 32 - let mut generator = RandPrintablesGenerator::new(32); + let mut generator = RandPrintablesGenerator::new(32).unwrap(); // Generate 8 initial inputs state diff --git a/fuzzers/forkserver/forkserver_libafl_cc/src/main.rs b/fuzzers/forkserver/forkserver_libafl_cc/src/main.rs index 4e4deb0680..a33c8893b0 100644 --- a/fuzzers/forkserver/forkserver_libafl_cc/src/main.rs +++ b/fuzzers/forkserver/forkserver_libafl_cc/src/main.rs @@ -206,7 +206,8 @@ pub fn main() { // Setup a mutational stage with a basic bytes mutator let mutator = - StdScheduledMutator::with_max_stack_pow(havoc_mutations().merge(tokens_mutations()), 6); + StdScheduledMutator::with_max_stack_pow(havoc_mutations().merge(tokens_mutations()), 6) + .unwrap(); let mut stages = tuple_list!(StdMutationalStage::new(mutator)); fuzzer diff --git a/fuzzers/forkserver/forkserver_simple/build.rs b/fuzzers/forkserver/forkserver_simple/build.rs index 22f71a6a29..dc8c8dd146 100644 --- a/fuzzers/forkserver/forkserver_simple/build.rs +++ b/fuzzers/forkserver/forkserver_simple/build.rs @@ -44,7 +44,7 @@ fn main() { let mut compile_command = Command::new(afl_cc_path); compile_command .args(["src/program.c", "-o"]) - .arg(&format!("{}/target/release/program", &cwd)); + .arg(format!("{cwd}/target/release/program")); if let Ok(llvm_config) = env::var("LLVM_CONFIG") { if !llvm_config.is_empty() { diff --git a/fuzzers/forkserver/forkserver_simple/src/main.rs b/fuzzers/forkserver/forkserver_simple/src/main.rs index c026197738..b846e1bfca 100644 --- a/fuzzers/forkserver/forkserver_simple/src/main.rs +++ b/fuzzers/forkserver/forkserver_simple/src/main.rs @@ -202,7 +202,8 @@ pub fn main() { // Setup a mutational stage with a basic bytes mutator let mutator = - StdScheduledMutator::with_max_stack_pow(havoc_mutations().merge(tokens_mutations()), 6); + StdScheduledMutator::with_max_stack_pow(havoc_mutations().merge(tokens_mutations()), 6) + .unwrap(); let mut stages = tuple_list!(StdMutationalStage::new(mutator)); fuzzer diff --git a/fuzzers/fuzz_anything/baby_fuzzer_wasm/src/lib.rs b/fuzzers/fuzz_anything/baby_fuzzer_wasm/src/lib.rs index b981565fd3..9dcfa4172d 100644 --- a/fuzzers/fuzz_anything/baby_fuzzer_wasm/src/lib.rs +++ b/fuzzers/fuzz_anything/baby_fuzzer_wasm/src/lib.rs @@ -126,7 +126,7 @@ pub fn fuzz() { .expect("Failed to create the Executor"); // Generator of printable bytearrays of max size 32 - let mut generator = RandPrintablesGenerator::new(32); + let mut generator = RandPrintablesGenerator::new(32).unwrap(); // Generate 8 initial inputs state diff --git a/fuzzers/fuzz_anything/baby_no_std/src/main.rs b/fuzzers/fuzz_anything/baby_no_std/src/main.rs index febcac17ff..2cfaa856b8 100644 --- a/fuzzers/fuzz_anything/baby_no_std/src/main.rs +++ b/fuzzers/fuzz_anything/baby_no_std/src/main.rs @@ -144,7 +144,7 @@ pub extern "C" fn main(_argc: isize, _argv: *const *const u8) -> isize { .expect("Failed to create the Executor"); // Generator of printable bytearrays of max size 32 - let mut generator = RandPrintablesGenerator::new(32); + let mut generator = RandPrintablesGenerator::new(32).unwrap(); // Generate 8 initial inputs state diff --git a/fuzzers/fuzz_anything/libafl_atheris/src/lib.rs b/fuzzers/fuzz_anything/libafl_atheris/src/lib.rs index da8b1d03fd..111647c758 100644 --- a/fuzzers/fuzz_anything/libafl_atheris/src/lib.rs +++ b/fuzzers/fuzz_anything/libafl_atheris/src/lib.rs @@ -240,7 +240,7 @@ pub extern "C" fn LLVMFuzzerRunDriver( if state.must_load_initial_inputs() { if input_dirs.is_empty() { // Generator of printable bytearrays of max size 32 - let mut generator = RandBytesGenerator::new(32); + let mut generator = RandBytesGenerator::new(32).unwrap(); // Generate 8 initial inputs state diff --git a/fuzzers/fuzz_anything/push_harness/src/main.rs b/fuzzers/fuzz_anything/push_harness/src/main.rs index 2614bad70f..f175bffe4b 100644 --- a/fuzzers/fuzz_anything/push_harness/src/main.rs +++ b/fuzzers/fuzz_anything/push_harness/src/main.rs @@ -1,6 +1,6 @@ //! [`Klo-routines`](https://github.com/andreafioraldi/klo-routines/) based fuzzer. //! The target loops and the harness pulls inputs out of `LibAFL`, instead of being called by `LibAFL`. -use std::path::PathBuf; +use std::{path::PathBuf, ptr::addr_of_mut}; use klo_routines::{yield_, KloRoutine}; use libafl::{ @@ -39,8 +39,12 @@ fn input_generator() { ExitKind::Ok }; + let signals_ptr = addr_of_mut!(SIGNALS); + let signals_len = unsafe { *signals_ptr }.len(); + // Create an observation channel using the signals map - let observer = unsafe { StdMapObserver::new("signals", &mut SIGNALS) }; + let observer = + unsafe { StdMapObserver::from_mut_ptr("signals", addr_of_mut!(SIGNALS) as _, signals_len) }; // Feedback to rate the interestingness of an input let mut feedback = MaxMapFeedback::new(&observer); @@ -89,7 +93,7 @@ fn input_generator() { .expect("Failed to create the Executor"); // Generator of printable bytearrays of max size 32 - let mut generator = RandPrintablesGenerator::new(32); + let mut generator = RandPrintablesGenerator::new(32).unwrap(); // Generate 8 initial inputs state diff --git a/fuzzers/inprocess/fuzzbench_text/src/lib.rs b/fuzzers/inprocess/fuzzbench_text/src/lib.rs index 8dc73143e4..1ee7ff987d 100644 --- a/fuzzers/inprocess/fuzzbench_text/src/lib.rs +++ b/fuzzers/inprocess/fuzzbench_text/src/lib.rs @@ -602,7 +602,8 @@ fn fuzz_text( GrimoireRandomDeleteMutator::new(), ), 3, - ); + ) + .unwrap(); let grimoire = StdMutationalStage::transforming(grimoire_mutator); // A minimization+queue policy to get testcasess from the corpus diff --git a/fuzzers/structure_aware/baby_fuzzer_custom_input/src/input.rs b/fuzzers/structure_aware/baby_fuzzer_custom_input/src/input.rs index 9a1dc569fa..904c96e321 100644 --- a/fuzzers/structure_aware/baby_fuzzer_custom_input/src/input.rs +++ b/fuzzers/structure_aware/baby_fuzzer_custom_input/src/input.rs @@ -60,23 +60,25 @@ impl CustomInput { } /// A generator for [`CustomInput`] used in this example -pub struct CustomInputGenerator { - pub max_len: usize, +pub struct CustomInputGenerator { + pub bytes_generator: RandBytesGenerator, } -impl CustomInputGenerator { +impl CustomInputGenerator { /// Creates a new [`CustomInputGenerator`] - pub fn new(max_len: usize) -> Self { - Self { max_len } + pub fn new(max_len: usize) -> Result { + Ok(Self { + bytes_generator: RandBytesGenerator::new(max_len)?, + }) } } -impl Generator for CustomInputGenerator +impl Generator for CustomInputGenerator where S: HasRand, { fn generate(&mut self, state: &mut S) -> Result { - let mut generator = RandBytesGenerator::new(self.max_len); + let generator = &mut self.bytes_generator; let byte_array = generator.generate(state).unwrap().target_bytes().into(); let optional_byte_array = state @@ -103,10 +105,10 @@ where S: HasRand, { /// Creates a new [`ToggleOptionalByteArrayMutator`] - pub fn new(length: usize) -> Self { - Self { - generator: RandBytesGenerator::new(length), - } + pub fn new(length: usize) -> Result { + Ok(Self { + generator: RandBytesGenerator::new(length)?, + }) } } diff --git a/fuzzers/structure_aware/baby_fuzzer_custom_input/src/main.rs b/fuzzers/structure_aware/baby_fuzzer_custom_input/src/main.rs index ec0c63f7d1..572afadaf5 100644 --- a/fuzzers/structure_aware/baby_fuzzer_custom_input/src/main.rs +++ b/fuzzers/structure_aware/baby_fuzzer_custom_input/src/main.rs @@ -2,7 +2,10 @@ mod input; #[cfg(windows)] use std::ptr::write_volatile; -use std::{path::PathBuf, ptr::write}; +use std::{ + path::PathBuf, + ptr::{addr_of_mut, write}, +}; use input::{ CustomInput, CustomInputGenerator, ToggleBooleanMutator, ToggleOptionalByteArrayMutator, @@ -37,8 +40,9 @@ use { }; /// Coverage map with explicit assignments due to the lack of instrumentation -static mut SIGNALS: [u8; 16] = [0; 16]; -static mut SIGNALS_PTR: *mut u8 = unsafe { SIGNALS.as_mut_ptr() }; +const SIGNALS_LEN: usize = 16; +static mut SIGNALS: [u8; SIGNALS_LEN] = [0; 16]; +static mut SIGNALS_PTR: *mut u8 = addr_of_mut!(SIGNALS) as _; /// Assign a signal to the signals map fn signals_set(idx: usize) { @@ -78,7 +82,7 @@ pub fn main() { }; // Create an observation channel using the signals map - let observer = unsafe { StdMapObserver::from_mut_ptr("signals", SIGNALS_PTR, SIGNALS.len()) }; + let observer = unsafe { StdMapObserver::from_mut_ptr("signals", SIGNALS_PTR, SIGNALS_LEN) }; // Feedback to rate the interestingness of an input let mut feedback = MaxMapFeedback::new(&observer); @@ -127,7 +131,8 @@ pub fn main() { .expect("Failed to create the Executor"); // Generator of printable bytearrays of max size 32 - let mut generator = CustomInputGenerator::new(1); + let mut generator = + CustomInputGenerator::new(1).expect("Failed to create our custom input generator"); // Generate 8 initial inputs state @@ -178,7 +183,9 @@ pub fn main() { // Then, mutators for the optional byte array, these return MutationResult::Skipped if the part is not present .merge(optional_mapped_mutators) // A custom mutator that sets the optional byte array to None if present, and generates a random byte array of length 1 if it is not - .prepend(ToggleOptionalByteArrayMutator::new(1)) + .prepend( + ToggleOptionalByteArrayMutator::new(1).expect("Failed to create bytearray mutator"), + ) // Finally, a custom mutator that toggles the boolean part of the input .prepend(ToggleBooleanMutator); diff --git a/fuzzers/structure_aware/baby_fuzzer_gramatron/src/main.rs b/fuzzers/structure_aware/baby_fuzzer_gramatron/src/main.rs index 357bd3b17a..6fdc0e6e5e 100644 --- a/fuzzers/structure_aware/baby_fuzzer_gramatron/src/main.rs +++ b/fuzzers/structure_aware/baby_fuzzer_gramatron/src/main.rs @@ -4,6 +4,7 @@ use std::{ fs, io::{BufReader, Read}, path::{Path, PathBuf}, + ptr::addr_of_mut, }; use libafl::{ @@ -27,8 +28,9 @@ use libafl::{ use libafl_bolts::{rands::StdRand, tuples::tuple_list}; /// Coverage map with explicit assignments due to the lack of instrumentation -static mut SIGNALS: [u8; 16] = [0; 16]; -static mut SIGNALS_PTR: *mut u8 = unsafe { SIGNALS.as_mut_ptr() }; +const SIGNALS_LEN: usize = 16; +static mut SIGNALS: [u8; SIGNALS_LEN] = [0; SIGNALS_LEN]; +static mut SIGNALS_PTR: *mut u8 = addr_of_mut!(SIGNALS) as _; /* /// Assign a signal to the signals map fn signals_set(idx: usize) { @@ -58,7 +60,7 @@ pub fn main() { }; // Create an observation channel using the signals map - let observer = unsafe { StdMapObserver::from_mut_ptr("signals", SIGNALS_PTR, SIGNALS.len()) }; + let observer = unsafe { StdMapObserver::from_mut_ptr("signals", SIGNALS_PTR, SIGNALS_LEN) }; // Feedback to rate the interestingness of an input let mut feedback = MaxMapFeedback::new(&observer); @@ -154,7 +156,8 @@ pub fn main() { GramatronRecursionMutator::new() ), 2, - ); + ) + .unwrap(); let mut stages = tuple_list!(StdMutationalStage::new(mutator)); fuzzer diff --git a/fuzzers/structure_aware/baby_fuzzer_grimoire/src/main.rs b/fuzzers/structure_aware/baby_fuzzer_grimoire/src/main.rs index 60162e5919..6f137c8ce0 100644 --- a/fuzzers/structure_aware/baby_fuzzer_grimoire/src/main.rs +++ b/fuzzers/structure_aware/baby_fuzzer_grimoire/src/main.rs @@ -139,7 +139,7 @@ pub fn main() { .expect("Failed to create the Executor"); // Setup a mutational stage with a basic bytes mutator - let mutator = StdScheduledMutator::with_max_stack_pow(havoc_mutations(), 2); + let mutator = StdScheduledMutator::with_max_stack_pow(havoc_mutations(), 2).unwrap(); let grimoire_mutator = StdScheduledMutator::with_max_stack_pow( tuple_list!( GrimoireExtensionMutator::new(), @@ -150,7 +150,8 @@ pub fn main() { GrimoireRandomDeleteMutator::new(), ), 3, - ); + ) + .unwrap(); let mut stages = tuple_list!( generalization, StdMutationalStage::new(mutator), diff --git a/fuzzers/structure_aware/baby_fuzzer_nautilus/src/main.rs b/fuzzers/structure_aware/baby_fuzzer_nautilus/src/main.rs index 75a2fcbe9d..eadc4393a7 100644 --- a/fuzzers/structure_aware/baby_fuzzer_nautilus/src/main.rs +++ b/fuzzers/structure_aware/baby_fuzzer_nautilus/src/main.rs @@ -158,7 +158,8 @@ pub fn main() { NautilusSpliceMutator::new(&context), ), 2, - ); + ) + .unwrap(); let mut stages = tuple_list!(StdMutationalStage::new(mutator)); fuzzer diff --git a/fuzzers/structure_aware/baby_fuzzer_tokens/src/main.rs b/fuzzers/structure_aware/baby_fuzzer_tokens/src/main.rs index a47f76d677..42f2f369e9 100644 --- a/fuzzers/structure_aware/baby_fuzzer_tokens/src/main.rs +++ b/fuzzers/structure_aware/baby_fuzzer_tokens/src/main.rs @@ -116,7 +116,7 @@ pub fn main() { .expect("Failed to create the Executor"); // Setup a mutational stage with a basic bytes mutator - let mutator = StdScheduledMutator::with_max_stack_pow(encoded_mutations(), 2); + let mutator = StdScheduledMutator::with_max_stack_pow(encoded_mutations(), 2).unwrap(); let mut stages = tuple_list!(StdMutationalStage::new(mutator)); println!("Decoder {:?} ...", &encoder_decoder); diff --git a/fuzzers/structure_aware/nautilus_sync/src/lib.rs b/fuzzers/structure_aware/nautilus_sync/src/lib.rs index 309f0e6b02..e0b7fb6a41 100644 --- a/fuzzers/structure_aware/nautilus_sync/src/lib.rs +++ b/fuzzers/structure_aware/nautilus_sync/src/lib.rs @@ -230,7 +230,8 @@ pub extern "C" fn libafl_main() { NautilusSpliceMutator::new(&context), ), 2, - ); + ) + .unwrap(); if let Some(conv) = event_converter.take() { let mut stages = tuple_list!( diff --git a/libafl/src/common/nautilus/grammartec/context.rs b/libafl/src/common/nautilus/grammartec/context.rs index 4e4d7e9453..c4a6ca72c3 100644 --- a/libafl/src/common/nautilus/grammartec/context.rs +++ b/libafl/src/common/nautilus/grammartec/context.rs @@ -1,4 +1,5 @@ use alloc::{borrow::ToOwned, string::String, vec::Vec}; +use core::num::NonZero; use hashbrown::HashMap; use libafl_bolts::rands::{Rand, RomuDuoJrRand}; @@ -263,7 +264,8 @@ impl Context { .iter() .take_while(move |r| self.rules_to_min_size[*r] <= max_len) .filter(move |r| { - self.rules_to_num_options[*r] > 1 || rand.below(100) <= p_include_short_rules + self.rules_to_num_options[*r] > 1 + || rand.below(NonZero::new(100).unwrap()) <= p_include_short_rules }) } diff --git a/libafl/src/common/nautilus/grammartec/mutator.rs b/libafl/src/common/nautilus/grammartec/mutator.rs index 851b2e9010..2ec854f79c 100644 --- a/libafl/src/common/nautilus/grammartec/mutator.rs +++ b/libafl/src/common/nautilus/grammartec/mutator.rs @@ -1,5 +1,5 @@ use alloc::vec::Vec; -use core::mem; +use core::{mem, num::NonZero}; use hashbrown::HashSet; use libafl_bolts::{rands::Rand, Error}; @@ -145,7 +145,10 @@ impl Mutator { where F: FnMut(&TreeMutation, &Context) -> Result<(), Error>, { - let n = NodeId::from(rand.below(tree.size())); + let Some(tree_size) = NonZero::new(tree.size()) else { + return Err(Error::illegal_argument("Empty tree in mut_splice")); + }; + let n = NodeId::from(rand.below(tree_size)); let old_rule_id = tree.get_rule_id(n); if let Some((repl_tree, repl_node)) = cks.get_alternative_to(rand, old_rule_id, ctx) { let repl = tree.mutate_replace_from_tree(n, repl_tree, repl_node); @@ -186,7 +189,10 @@ impl Mutator { where F: FnMut(&TreeMutation, &Context) -> Result<(), Error>, { - let n = NodeId::from(rand.below(tree.size())); + let Some(tree_size) = NonZero::new(tree.size()) else { + return Err(Error::illegal_argument("Empty tree in mut_random")); + }; + let n = NodeId::from(rand.below(tree_size)); let nterm = tree.get_rule(n, ctx).nonterm(); if ctx.check_if_nterm_has_multiple_possiblities(&nterm) { let len = ctx.get_random_len_for_nt(&nterm); diff --git a/libafl/src/common/nautilus/grammartec/recursion_info.rs b/libafl/src/common/nautilus/grammartec/recursion_info.rs index 230d6a5afb..49d7e5c99d 100644 --- a/libafl/src/common/nautilus/grammartec/recursion_info.rs +++ b/libafl/src/common/nautilus/grammartec/recursion_info.rs @@ -2,7 +2,10 @@ use alloc::vec::Vec; use std::fmt; use hashbrown::HashMap; -use libafl_bolts::rands::{loaded_dice::LoadedDiceSampler, Rand}; +use libafl_bolts::{ + rands::{loaded_dice::LoadedDiceSampler, Rand}, + Error, +}; use crate::common::nautilus::grammartec::{ context::Context, @@ -32,7 +35,8 @@ impl RecursionInfo { pub fn new(t: &Tree, n: NTermId, ctx: &Context) -> Option { let (recursive_parents, node_by_offset, depth_by_offset) = RecursionInfo::find_parents(t, n, ctx)?; - let sampler = RecursionInfo::build_sampler(&depth_by_offset); + let sampler = RecursionInfo::build_sampler(&depth_by_offset) + .expect("Sampler depth_by_offset invalid"); Some(Self { recursive_parents, sampler, @@ -79,7 +83,7 @@ impl RecursionInfo { } #[allow(clippy::cast_precision_loss)] - fn build_sampler(depths: &[usize]) -> LoadedDiceSampler { + fn build_sampler(depths: &[usize]) -> Result { let mut weights = depths.iter().map(|x| *x as f64).collect::>(); let norm: f64 = weights.iter().sum(); assert!(norm > 0.0); diff --git a/libafl/src/common/nautilus/regex_mutator/mod.rs b/libafl/src/common/nautilus/regex_mutator/mod.rs index 451d0504d6..e5885a0d86 100644 --- a/libafl/src/common/nautilus/regex_mutator/mod.rs +++ b/libafl/src/common/nautilus/regex_mutator/mod.rs @@ -1,4 +1,5 @@ use alloc::vec::Vec; +use core::num::NonZero; use libafl_bolts::rands::Rand; use regex_syntax::hir::{Class, ClassBytesRange, ClassUnicodeRange, Hir, Literal}; @@ -22,10 +23,12 @@ impl RegexScript { } pub fn get_mod(&mut self, rand: &mut R, val: usize) -> usize { - if self.remaining == 0 { + if self.remaining == 0 || val == 0 { 0 } else { - rand.below(val) + // # Safety + // This is checked above to be non-null. + rand.below(unsafe { NonZero::new(val).unwrap_unchecked() }) } } diff --git a/libafl/src/corpus/mod.rs b/libafl/src/corpus/mod.rs index 4bf07626c0..206e267c9b 100644 --- a/libafl/src/corpus/mod.rs +++ b/libafl/src/corpus/mod.rs @@ -63,23 +63,39 @@ impl From for usize { } } -/// Utility macro to call `Corpus::random_id`; fetches only enabled testcases +/// Utility macro to call `Corpus::random_id`; fetches only enabled [`Testcase`]`s` #[macro_export] macro_rules! random_corpus_id { ($corpus:expr, $rand:expr) => {{ let cnt = $corpus.count(); - let nth = $rand.below(cnt); + #[cfg(debug_assertions)] + let nth = $rand.below(core::num::NonZero::new(cnt).expect("Corpus may not be empty!")); + // # Safety + // This is a hot path. We try to be as fast as possible here. + // In debug this is checked (see above.) + // The worst that can happen is a wrong integer to get returned. + // In this case, the call below will fail. + #[cfg(not(debug_assertions))] + let nth = $rand.below(unsafe { core::num::NonZero::new(cnt).unwrap_unchecked() }); $corpus.nth(nth) }}; } -/// Utility macro to call `Corpus::random_id`; fetches both enabled and disabled testcases +/// Utility macro to call `Corpus::random_id`; fetches both enabled and disabled [`Testcase`]`s` /// Note: use `Corpus::get_from_all` as disabled entries are inaccessible from `Corpus::get` #[macro_export] macro_rules! random_corpus_id_with_disabled { ($corpus:expr, $rand:expr) => {{ let cnt = $corpus.count_all(); - let nth = $rand.below(cnt); + #[cfg(debug_assertions)] + let nth = $rand.below(core::num::NonZero::new(cnt).expect("Corpus may not be empty!")); + // # Safety + // This is a hot path. We try to be as fast as possible here. + // In debug this is checked (see above.) + // The worst that can happen is a wrong integer to get returned. + // In this case, the call below will fail. + #[cfg(not(debug_assertions))] + let nth = $rand.below(unsafe { core::num::NonZero::new(cnt).unwrap_unchecked() }); $corpus.nth_from_all(nth) }}; } diff --git a/libafl/src/events/mod.rs b/libafl/src/events/mod.rs index 08d7b99e36..abe1e9494a 100644 --- a/libafl/src/events/mod.rs +++ b/libafl/src/events/mod.rs @@ -34,7 +34,7 @@ pub use broker_hooks::*; #[cfg(feature = "std")] pub use launcher::*; #[cfg(all(unix, feature = "std"))] -use libafl_bolts::os::unix_signals::{siginfo_t, ucontext_t, Handler, Signal}; +use libafl_bolts::os::unix_signals::{siginfo_t, ucontext_t, Signal, SignalHandler}; #[cfg(all(unix, feature = "std"))] use libafl_bolts::os::CTRL_C_EXIT; use libafl_bolts::{ @@ -80,9 +80,12 @@ pub struct ShutdownSignalData {} /// Shutdown handler. `SigTerm`, `SigInterrupt`, `SigQuit` call this /// We can't handle SIGKILL in the signal handler, this means that you shouldn't kill your fuzzer with `kill -9` because then the shmem segments are never freed +/// +/// # Safety +/// This will exit the program #[cfg(all(unix, feature = "std"))] -impl Handler for ShutdownSignalData { - fn handle( +impl SignalHandler for ShutdownSignalData { + unsafe fn handle( &mut self, _signal: Signal, _info: &mut siginfo_t, diff --git a/libafl/src/executors/hooks/inprocess.rs b/libafl/src/executors/hooks/inprocess.rs index 7240d62f76..993fefdeec 100644 --- a/libafl/src/executors/hooks/inprocess.rs +++ b/libafl/src/executors/hooks/inprocess.rs @@ -526,5 +526,7 @@ pub unsafe fn inprocess_get_input<'a, I>() -> Option<&'a I> { /// Returns if we are executing in a crash/timeout handler #[must_use] pub fn inprocess_in_handler() -> bool { + // # Safety + // Safe because the state is set up and the handler is a single bool. Worst case we read an old value. unsafe { GLOBAL_STATE.in_handler } } diff --git a/libafl/src/executors/hooks/inprocess_fork.rs b/libafl/src/executors/hooks/inprocess_fork.rs index 8c1d63025e..8f86ad9569 100644 --- a/libafl/src/executors/hooks/inprocess_fork.rs +++ b/libafl/src/executors/hooks/inprocess_fork.rs @@ -10,7 +10,7 @@ use std::intrinsics::transmute; #[cfg(not(miri))] use libafl_bolts::os::unix_signals::setup_signal_handler; -use libafl_bolts::os::unix_signals::{ucontext_t, Handler, Signal}; +use libafl_bolts::os::unix_signals::{ucontext_t, Signal, SignalHandler}; use libc::siginfo_t; use crate::{ @@ -145,8 +145,13 @@ pub(crate) static mut FORK_EXECUTOR_GLOBAL_DATA: InProcessForkExecutorGlobalData timeout_handler: null(), }; -impl Handler for InProcessForkExecutorGlobalData { - fn handle(&mut self, signal: Signal, info: &mut siginfo_t, context: Option<&mut ucontext_t>) { +impl SignalHandler for InProcessForkExecutorGlobalData { + unsafe fn handle( + &mut self, + signal: Signal, + info: &mut siginfo_t, + context: Option<&mut ucontext_t>, + ) { match signal { Signal::SigUser2 | Signal::SigAlarm => unsafe { if !FORK_EXECUTOR_GLOBAL_DATA.timeout_handler.is_null() { diff --git a/libafl/src/executors/hooks/unix.rs b/libafl/src/executors/hooks/unix.rs index 7d9c1030c2..37ce65af9f 100644 --- a/libafl/src/executors/hooks/unix.rs +++ b/libafl/src/executors/hooks/unix.rs @@ -5,7 +5,7 @@ pub mod unix_signal_handler { use core::{mem::transmute, ptr::addr_of_mut}; use std::{io::Write, panic}; - use libafl_bolts::os::unix_signals::{ucontext_t, Handler, Signal}; + use libafl_bolts::os::unix_signals::{ucontext_t, Signal, SignalHandler}; use libc::siginfo_t; use crate::{ @@ -41,8 +41,10 @@ pub mod unix_signal_handler { }*/ #[cfg(unix)] - impl Handler for InProcessExecutorHandlerData { - fn handle( + impl SignalHandler for InProcessExecutorHandlerData { + /// # Safety + /// This will access global state. + unsafe fn handle( &mut self, signal: Signal, info: &mut siginfo_t, diff --git a/libafl/src/executors/hooks/windows.rs b/libafl/src/executors/hooks/windows.rs index 95b78252b3..eed08c5ae3 100644 --- a/libafl/src/executors/hooks/windows.rs +++ b/libafl/src/executors/hooks/windows.rs @@ -123,7 +123,8 @@ pub mod windows_exception_handler { use std::panic; use libafl_bolts::os::windows_exceptions::{ - ExceptionCode, Handler, CRASH_EXCEPTIONS, EXCEPTION_HANDLERS_SIZE, EXCEPTION_POINTERS, + ExceptionCode, ExceptionHandler, CRASH_EXCEPTIONS, EXCEPTION_HANDLERS_SIZE, + EXCEPTION_POINTERS, }; use windows::Win32::System::Threading::{ EnterCriticalSection, ExitProcess, LeaveCriticalSection, CRITICAL_SECTION, @@ -154,9 +155,15 @@ pub mod windows_exception_handler { ) { }*/ - impl Handler for InProcessExecutorHandlerData { + impl ExceptionHandler for InProcessExecutorHandlerData { + /// # Safety + /// Will dereference EXCEPTION_POINTERS and access `GLOBAL_STATE`. #[allow(clippy::not_unsafe_ptr_arg_deref)] - fn handle(&mut self, _code: ExceptionCode, exception_pointers: *mut EXCEPTION_POINTERS) { + unsafe fn handle( + &mut self, + _code: ExceptionCode, + exception_pointers: *mut EXCEPTION_POINTERS, + ) { unsafe { let data = addr_of_mut!(GLOBAL_STATE); let in_handler = (*data).set_in_handler(true); diff --git a/libafl/src/generators/gramatron.rs b/libafl/src/generators/gramatron.rs index 4320152e08..b308671886 100644 --- a/libafl/src/generators/gramatron.rs +++ b/libafl/src/generators/gramatron.rs @@ -1,6 +1,6 @@ //! Gramatron generator use alloc::{string::String, vec::Vec}; -use core::marker::PhantomData; +use core::{marker::PhantomData, num::NonZero}; use libafl_bolts::rands::Rand; use serde::{Deserialize, Serialize}; @@ -73,13 +73,21 @@ where .last() .map_or(self.automaton.init_state, |last| { let triggers = &self.automaton.pda[last.state]; - let idx = state.rand_mut().below(triggers.len()); + let idx = state.rand_mut().below( + NonZero::new(triggers.len()) + .expect("Triggers are empty in append_generated_terminals!"), + ); triggers[idx].dest }); while current_state != final_state { let triggers = &self.automaton.pda[current_state]; - let idx = state.rand_mut().below(triggers.len()); + let idx = + state + .rand_mut() + .below(NonZero::new(triggers.len()).expect( + "Automation.pda triggers are empty in append_generated_terminals!", + )); let trigger = &triggers[idx]; input .terminals_mut() diff --git a/libafl/src/generators/mod.rs b/libafl/src/generators/mod.rs index e7e47b5fea..00e34dd1e9 100644 --- a/libafl/src/generators/mod.rs +++ b/libafl/src/generators/mod.rs @@ -1,7 +1,10 @@ //! Generators may generate bytes or, in general, data, for inputs. use alloc::vec::Vec; -use core::marker::PhantomData; +use core::{ + marker::PhantomData, + num::{NonZero, NonZeroUsize}, +}; use libafl_bolts::rands::Rand; @@ -72,7 +75,7 @@ where #[derive(Clone, Debug)] /// Generates random bytes pub struct RandBytesGenerator { - max_size: usize, + max_size: NonZeroUsize, phantom: PhantomData, } @@ -86,16 +89,26 @@ where size = 1; } let random_bytes: Vec = (0..size) - .map(|_| state.rand_mut().below(256) as u8) + .map(|_| state.rand_mut().below(NonZero::new(256).unwrap()) as u8) .collect(); Ok(BytesInput::new(random_bytes)) } } impl RandBytesGenerator { + /// Returns a new [`RandBytesGenerator`], generating up to `max_size` random bytes. + /// + /// If you want to save one 0 check, use [`Self::from_nonzero`]. + pub fn new(max_size: usize) -> Result { + let Some(max_size) = NonZero::new(max_size) else { + return Err(Error::illegal_argument("The max_size may not be 0.")); + }; + Ok(Self::from_nonzero(max_size)) + } + /// Returns a new [`RandBytesGenerator`], generating up to `max_size` random bytes. #[must_use] - pub fn new(max_size: usize) -> Self { + pub fn from_nonzero(max_size: NonZeroUsize) -> Self { Self { max_size, phantom: PhantomData, @@ -106,7 +119,7 @@ impl RandBytesGenerator { #[derive(Clone, Debug)] /// Generates random printable characters pub struct RandPrintablesGenerator { - max_size: usize, + max_size: NonZeroUsize, phantom: PhantomData, } @@ -129,8 +142,18 @@ where impl RandPrintablesGenerator { /// Creates a new [`RandPrintablesGenerator`], generating up to `max_size` random printable characters. + /// + /// To skip the 0 `max_size` check, create this using [`Self::from_nonzero`] instead. + pub fn new(max_size: usize) -> Result { + let Some(max_size) = NonZero::new(max_size) else { + return Err(Error::illegal_argument("The max_size may not be 0.")); + }; + Ok(Self::from_nonzero(max_size)) + } + + /// Returns a new [`RandBytesGenerator`], generating up to `max_size` random bytes. #[must_use] - pub fn new(max_size: usize) -> Self { + pub fn from_nonzero(max_size: NonZeroUsize) -> Self { Self { max_size, phantom: PhantomData, diff --git a/libafl/src/mutators/encoded_mutations.rs b/libafl/src/mutators/encoded_mutations.rs index 6e07b02f74..b074d8068d 100644 --- a/libafl/src/mutators/encoded_mutations.rs +++ b/libafl/src/mutators/encoded_mutations.rs @@ -1,7 +1,10 @@ //! Mutations for [`EncodedInput`]s //! use alloc::{borrow::Cow, vec::Vec}; -use core::cmp::{max, min}; +use core::{ + cmp::{max, min}, + num::NonZero, +}; use libafl_bolts::{ rands::Rand, @@ -123,8 +126,8 @@ impl Mutator for EncodedAddMutator { Ok(MutationResult::Skipped) } else { let val = state.rand_mut().choose(input.codes_mut()).unwrap(); - let num = 1 + state.rand_mut().below(ARITH_MAX) as u32; - *val = match state.rand_mut().below(2) { + let num = 1 + state.rand_mut().below(NonZero::new(ARITH_MAX).unwrap()) as u32; + *val = match state.rand_mut().below(NonZero::new(2).unwrap()) { 0 => val.wrapping_add(num), _ => val.wrapping_sub(num), }; @@ -158,9 +161,16 @@ impl Mutator for EncodedDeleteMutator { if size <= 2 { return Ok(MutationResult::Skipped); } - - let off = state.rand_mut().below(size); - let len = state.rand_mut().below(size - off); + // # Safety + // The size is larger than 1 here (checked just above) + let off = state + .rand_mut() + .below(unsafe { NonZero::new(size).unwrap_unchecked() }); + // # Safety + // The size of the offset is below size, the value is never 0. + let len = state + .rand_mut() + .below(unsafe { NonZero::new(size - off).unwrap_unchecked() }); input.codes_mut().drain(off..off + len); Ok(MutationResult::Mutated) @@ -195,11 +205,17 @@ where fn mutate(&mut self, state: &mut S, input: &mut EncodedInput) -> Result { let max_size = state.max_size(); let size = input.codes().len(); - if size == 0 { + let Some(nonzero_size) = NonZero::new(size) else { return Ok(MutationResult::Skipped); - } - let off = state.rand_mut().below(size + 1); - let mut len = 1 + state.rand_mut().below(min(16, size)); + }; + + // # Safety + // The input.codes() len should never be close to an usize, so adding 1 will always result in a non-zero value. + // Worst case, we will get a wrong int value as return, not too bad. + let off = state + .rand_mut() + .below(unsafe { NonZero::new(size + 1).unwrap_unchecked() }); + let mut len = 1 + state.rand_mut().below(nonzero_size); if size + len > max_size { if max_size > size { @@ -209,10 +225,10 @@ where } } - let from = if size == len { - 0 + let from = if let Some(bound) = NonZero::new(size - len) { + state.rand_mut().below(bound) } else { - state.rand_mut().below(size - len) + 0 }; input.codes_mut().resize(size + len, 0); @@ -250,13 +266,17 @@ pub struct EncodedCopyMutator; impl Mutator for EncodedCopyMutator { fn mutate(&mut self, state: &mut S, input: &mut EncodedInput) -> Result { let size = input.codes().len(); - if size <= 1 { + let Some(size) = NonZero::new(size) else { return Ok(MutationResult::Skipped); - } + }; let from = state.rand_mut().below(size); let to = state.rand_mut().below(size); - let len = 1 + state.rand_mut().below(size - max(from, to)); + // # Safety + // Both from and to are smaller than size, so size minus any of these can never be 0. + let len = 1 + state + .rand_mut() + .below(unsafe { NonZero::new(size.get() - max(from, to)).unwrap_unchecked() }); unsafe { buffer_self_copy(input.codes_mut(), from, to, len); @@ -301,7 +321,12 @@ where } } + let Some(non_zero_size) = NonZero::new(size) else { + return Ok(MutationResult::Skipped); + }; + let other_size = { + // new scope to make the borrow checker happy let mut other_testcase = state.corpus().get_from_all(id)?.borrow_mut(); other_testcase.load_input(state.corpus())?.codes().len() }; @@ -310,10 +335,18 @@ where return Ok(MutationResult::Skipped); } + let Some(non_zero_other_size) = NonZero::new(other_size) else { + return Ok(MutationResult::Skipped); + }; + let max_size = state.max_size(); - let from = state.rand_mut().below(other_size); - let to = state.rand_mut().below(size); - let mut len = 1 + state.rand_mut().below(other_size - from); + let from = state.rand_mut().below(non_zero_other_size); + let to = state.rand_mut().below(non_zero_size); + // # Safety + // from is smaller than other_size, other_size is larger than 2, so the subtraction is larger than 0. + let mut len = 1 + state + .rand_mut() + .below(unsafe { NonZero::new(other_size - from).unwrap_unchecked() }); if size + len > max_size { if max_size > size { @@ -363,9 +396,6 @@ where { fn mutate(&mut self, state: &mut S, input: &mut EncodedInput) -> Result { let size = input.codes().len(); - if size == 0 { - return Ok(MutationResult::Skipped); - } let id = random_corpus_id_with_disabled!(state.corpus(), state.rand_mut()); // We don't want to use the testcase we're already using for splicing @@ -385,9 +415,22 @@ where return Ok(MutationResult::Skipped); } - let from = state.rand_mut().below(other_size); - let len = state.rand_mut().below(min(other_size - from, size)); - let to = state.rand_mut().below(size - len); + let Some(non_zero_other_size) = NonZero::new(other_size) else { + return Ok(MutationResult::Skipped); + }; + + let from = state.rand_mut().below(non_zero_other_size); + + let Some(non_zero_min_len) = NonZero::new(min(other_size - from, size)) else { + return Ok(MutationResult::Skipped); + }; + + let len = state.rand_mut().below(non_zero_min_len); + // # Safety + // size is non-zero, len is below min(size, ...), so the subtraction will always be positive. + let to = state + .rand_mut() + .below(unsafe { NonZero::new(size - len).unwrap_unchecked() }); let other_testcase = state.corpus().get_from_all(id)?.borrow_mut(); // no need to load the input again, it'll already be present at this point. diff --git a/libafl/src/mutators/gramatron.rs b/libafl/src/mutators/gramatron.rs index 0c963c60b0..cb74a975de 100644 --- a/libafl/src/mutators/gramatron.rs +++ b/libafl/src/mutators/gramatron.rs @@ -2,7 +2,7 @@ //! //! See the original gramatron repo [`Gramatron`](https://github.com/HexHive/Gramatron) for more details. use alloc::{borrow::Cow, vec::Vec}; -use core::cmp::max; +use core::{cmp::max, num::NonZero}; use hashbrown::HashMap; use libafl_bolts::{ @@ -42,7 +42,12 @@ where input: &mut GramatronInput, ) -> Result { if !input.terminals().is_empty() { - let size = state.rand_mut().below(input.terminals().len() + 1); + // # Safety + // We can assume that the count of terminals + 1 will never wrap around (otherwise it will break somewhere else). + // So len + 1 is always non-zero. + let size = state + .rand_mut() + .below(unsafe { NonZero::new(input.terminals().len() + 1).unwrap_unchecked() }); input.terminals_mut().truncate(size); } if self.generator.append_generated_terminals(input, state) > 0 { @@ -115,13 +120,13 @@ where state: &mut S, input: &mut GramatronInput, ) -> Result { - if input.terminals().is_empty() { + let Some(terminals_len) = NonZero::new(input.terminals().len()) else { return Ok(MutationResult::Skipped); - } + }; let id = random_corpus_id!(state.corpus(), state.rand_mut()); - let insert_at = state.rand_mut().below(input.terminals().len()); + let insert_at = state.rand_mut().below(terminals_len); let rand_num = state.rand_mut().next(); @@ -217,8 +222,12 @@ where let chosen = *state.rand_mut().choose(&self.states).unwrap(); let chosen_nums = self.counters.get(&chosen).unwrap().0; + let Some(non_zero_chosen_nums_minus_one) = NonZero::new(chosen_nums - 1) else { + return Ok(MutationResult::Skipped); + }; + #[allow(clippy::cast_sign_loss, clippy::pedantic)] - let mut first = state.rand_mut().below(chosen_nums - 1) as i64; + let mut first = state.rand_mut().below(non_zero_chosen_nums_minus_one) as i64; #[allow(clippy::cast_sign_loss, clippy::pedantic)] let mut second = state .rand_mut() @@ -250,7 +259,10 @@ where input.terminals_mut().truncate(idx_1); - for _ in 0..state.rand_mut().below(RECUR_THRESHOLD) { + for _ in 0..state + .rand_mut() + .below(NonZero::new(RECUR_THRESHOLD).unwrap()) + { input.terminals_mut().extend_from_slice(&self.feature); } diff --git a/libafl/src/mutators/grimoire.rs b/libafl/src/mutators/grimoire.rs index 0f647ad7b3..0394c9c1ea 100644 --- a/libafl/src/mutators/grimoire.rs +++ b/libafl/src/mutators/grimoire.rs @@ -2,7 +2,10 @@ //! See the original repo [`Grimoire`](https://github.com/RUB-SysSec/grimoire) for more details. use alloc::{borrow::Cow, vec::Vec}; -use core::cmp::{max, min}; +use core::{ + cmp::{max, min}, + num::NonZero, +}; use libafl_bolts::{ rands::{choose, fast_bound, Rand}, @@ -246,10 +249,11 @@ where let tokens_len = { let meta = state.metadata_map().get::(); if let Some(tokens) = meta { - if tokens.is_empty() { + if let Some(tokens_len) = NonZero::new(tokens.tokens().len()) { + tokens_len + } else { return Ok(MutationResult::Skipped); } - tokens.tokens().len() } else { return Ok(MutationResult::Skipped); } @@ -271,7 +275,10 @@ where let mut mutated = MutationResult::Skipped; let gen = generalised_meta.generalized_mut(); - let rand_idx = fast_bound(rand_idx, gen.len()); + let rand_idx = fast_bound( + rand_idx, + NonZero::new(gen.len()).ok_or_else(|| Error::empty("No Generalized Metadata found"))?, + ); 'first: for item in &mut gen[..rand_idx] { if let GeneralizedItem::Bytes(bytes) = item { @@ -363,8 +370,15 @@ where { self.gap_indices.push(i); } - let min_idx = self.gap_indices[state.rand_mut().below(self.gap_indices.len())]; - let max_idx = self.gap_indices[state.rand_mut().below(self.gap_indices.len())]; + + let Some(gap_indeces_len) = NonZero::new(self.gap_indices.len()) else { + return Err(Error::illegal_state( + "Gap indices may not be empty in grimoire mutator!", + )); + }; + + let min_idx = self.gap_indices[state.rand_mut().below(gap_indeces_len)]; + let max_idx = self.gap_indices[state.rand_mut().below(gap_indeces_len)]; let (min_idx, max_idx) = (min(min_idx, max_idx), max(min_idx, max_idx)); diff --git a/libafl/src/mutators/mopt_mutator.rs b/libafl/src/mutators/mopt_mutator.rs index cf7aa56aa5..607ab1e710 100644 --- a/libafl/src/mutators/mopt_mutator.rs +++ b/libafl/src/mutators/mopt_mutator.rs @@ -3,7 +3,10 @@ //! It uses a modified Particle Swarm Optimization algorithm to determine an optimal distribution of mutators. //! See and use alloc::{borrow::Cow, string::ToString, vec::Vec}; -use core::fmt::{self, Debug}; +use core::{ + fmt::{self, Debug}, + num::{NonZero, NonZeroUsize}, +}; use libafl_bolts::{ rands::{Rand, StdRand}, @@ -367,7 +370,7 @@ pub struct StdMOptMutator { mode: MOptMode, finds_before: usize, mutations: MT, - max_stack_pow: usize, + max_stack_pow: NonZeroUsize, } impl Mutator for StdMOptMutator @@ -511,6 +514,11 @@ impl StdMOptMutator { let rand_seed = state.rand_mut().next(); state.add_metadata::(MOpt::new(MT::LEN, swarm_num, rand_seed)?); } + let Some(max_stack_pow) = NonZero::new(max_stack_pow) else { + return Err(Error::illegal_argument( + "Got 0 as value for max_stack_pow in StdMOptMutator.", + )); + }; Ok(Self { name: Cow::from(format!("StdMOptMutator[{}]", mutations.names().join(","))), mode: MOptMode::Pilotfuzzing, diff --git a/libafl/src/mutators/multi.rs b/libafl/src/mutators/multi.rs index bebe923b4d..68be5211b8 100644 --- a/libafl/src/mutators/multi.rs +++ b/libafl/src/mutators/multi.rs @@ -1,6 +1,9 @@ //! Mutator definitions for [`MultipartInput`]s. See [`crate::inputs::multi`] for details. -use core::cmp::{min, Ordering}; +use core::{ + cmp::{min, Ordering}, + num::NonZero, +}; use libafl_bolts::{rands::Rand, Error}; @@ -41,13 +44,12 @@ where state: &mut S, input: &mut MultipartInput, ) -> Result { - if input.parts().is_empty() { - Ok(MutationResult::Skipped) - } else { - let selected = state.rand_mut().below(input.parts().len()); - let mutated = input.part_mut(selected).unwrap(); - self.mutate(state, mutated) - } + let Some(parts_len) = NonZero::new(input.parts().len()) else { + return Ok(MutationResult::Skipped); + }; + let selected = state.rand_mut().below(parts_len); + let mutated = input.part_mut(selected).unwrap(); + self.mutate(state, mutated) } fn post_exec(&mut self, state: &mut S, new_corpus_id: Option) -> Result<(), Error> { @@ -154,8 +156,17 @@ where .map(|(id, part)| (id, part.bytes().len())); if let Some((part_idx, size)) = maybe_size { - let target = state.rand_mut().below(size); - let range = rand_range(state, other_size, min(other_size, size - target)); + let Some(nonzero_size) = NonZero::new(size) else { + return Ok(MutationResult::Skipped); + }; + let target = state.rand_mut().below(nonzero_size); + // # Safety + // size is nonzero here (checked above), target is smaller than size + // -> the subtraction result is greater than 0. + // other_size is checked above to be larger than zero. + let range = rand_range(state, other_size, unsafe { + NonZero::new(min(other_size, size - target)).unwrap_unchecked() + }); let [part, chosen] = match part_idx.cmp(&choice) { Ordering::Less => input.parts_mut([part_idx, choice]), @@ -201,9 +212,18 @@ where .unwrap(); drop(other_testcase); let size = part.bytes().len(); + let Some(nonzero_size) = NonZero::new(size) else { + return Ok(MutationResult::Skipped); + }; - let target = state.rand_mut().below(size); - let range = rand_range(state, other_size, min(other_size, size - target)); + let target = state.rand_mut().below(nonzero_size); + // # Safety + // other_size is larger than 0, checked above. + // size is larger than 0. + // target is smaller than size -> the subtraction is larger than 0. + let range = rand_range(state, other_size, unsafe { + NonZero::new(min(other_size, size - target)).unwrap_unchecked() + }); let other_testcase = state.corpus().get(id)?.borrow_mut(); // No need to load the input again, it'll still be cached. @@ -265,8 +285,17 @@ where .map(|(id, part)| (id, part.bytes().len())); if let Some((part_idx, size)) = maybe_size { - let target = state.rand_mut().below(size); - let range = rand_range(state, other_size, min(other_size, size - target)); + let Some(nonzero_size) = NonZero::new(size) else { + return Ok(MutationResult::Skipped); + }; + + let target = state.rand_mut().below(nonzero_size); + // # Safety + // other_size is checked above. + // size is larger than than target and larger than 1. The subtraction result will always be positive. + let range = rand_range(state, other_size, unsafe { + NonZero::new(min(other_size, size - target)).unwrap_unchecked() + }); let [part, chosen] = match part_idx.cmp(&choice) { Ordering::Less => input.parts_mut([part_idx, choice]), @@ -306,9 +335,17 @@ where .unwrap(); drop(other_testcase); let size = part.bytes().len(); + let Some(nonzero_size) = NonZero::new(size) else { + return Ok(MutationResult::Skipped); + }; - let target = state.rand_mut().below(size); - let range = rand_range(state, other_size, min(other_size, size - target)); + let target = state.rand_mut().below(nonzero_size); + // # Safety + // other_size is checked above. + // size is larger than than target and larger than 1. The subtraction result will always be positive. + let range = rand_range(state, other_size, unsafe { + NonZero::new(min(other_size, size - target)).unwrap_unchecked() + }); let other_testcase = state.corpus().get(id)?.borrow_mut(); // No need to load the input again, it'll still be cached. diff --git a/libafl/src/mutators/mutations.rs b/libafl/src/mutators/mutations.rs index b0d31870d6..ee226cc738 100644 --- a/libafl/src/mutators/mutations.rs +++ b/libafl/src/mutators/mutations.rs @@ -4,7 +4,13 @@ use alloc::{ borrow::{Cow, ToOwned}, vec::Vec, }; -use core::{cmp::min, marker::PhantomData, mem::size_of, ops::Range}; +use core::{ + cmp::min, + marker::PhantomData, + mem::size_of, + num::{NonZero, NonZeroUsize}, + ops::Range, +}; use libafl_bolts::{rands::Rand, Named}; @@ -64,10 +70,13 @@ pub fn buffer_set(data: &mut [T], from: usize, len: usize, val: T) { /// /// This problem corresponds to: #[inline] -pub fn rand_range(state: &mut S, upper: usize, max_len: usize) -> Range { +pub fn rand_range(state: &mut S, upper: usize, max_len: NonZeroUsize) -> Range { let len = 1 + state.rand_mut().below(max_len); // sample from [1..upper + len] - let mut offset2 = 1 + state.rand_mut().below(upper + len - 1); + let Some(upper_len_minus1) = NonZero::new(upper + len - 1) else { + return 0..0; + }; + let mut offset2 = 1 + state.rand_mut().below(upper_len_minus1); let offset1 = offset2.saturating_sub(len); if offset2 > upper { offset2 = upper; @@ -305,7 +314,7 @@ where Ok(MutationResult::Skipped) } else { let byte = state.rand_mut().choose(input.bytes_mut()).unwrap(); - *byte ^= 1 + state.rand_mut().below(254) as u8; + *byte ^= 1 + state.rand_mut().below(NonZero::new(254).unwrap()) as u8; Ok(MutationResult::Mutated) } } @@ -356,8 +365,8 @@ macro_rules! add_mutator_impl { 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) { + let num = 1 + state.rand_mut().below(NonZero::new(ARITH_MAX).unwrap()) as $size; + let new_val = match state.rand_mut().below(NonZero::new(4).unwrap()) { 0 => val.wrapping_add(num), 1 => val.wrapping_sub(num), 2 => val.swap_bytes().wrapping_add(num).swap_bytes(), @@ -414,7 +423,11 @@ macro_rules! interesting_mutator_impl { } else { let bytes = input.bytes_mut(); let upper_bound = (bytes.len() + 1 - size_of::<$size>()); - let idx = state.rand_mut().below(upper_bound); + // # Safety + // the length is at least as large as the size here (checked above), and we add a 1 -> never zero. + let idx = state + .rand_mut() + .below(unsafe { NonZero::new(upper_bound).unwrap_unchecked() }); let val = *state.rand_mut().choose(&$interesting).unwrap() as $size; let new_bytes = match state.rand_mut().choose(&[0, 1]).unwrap() { 0 => val.to_be_bytes(), @@ -462,7 +475,11 @@ where return Ok(MutationResult::Skipped); } - let range = rand_range(state, size, size - 1); + // # Safety + // size - 1 is guaranteed to be larger than 0 because we abort on size <= 2 above. + let range = rand_range(state, size, unsafe { + NonZero::new(size - 1).unwrap_unchecked() + }); input.drain(range); @@ -501,7 +518,11 @@ where return Ok(MutationResult::Skipped); } - let range = rand_range(state, size, min(16, max_size - size)); + // # Safety + // max_size - size is larger than 0 because we check that size < max_size above + let range = rand_range(state, size, unsafe { + NonZero::new(min(16, max_size - size)).unwrap_unchecked() + }); input.resize(size + range.len(), 0); unsafe { @@ -548,8 +569,13 @@ where return Ok(MutationResult::Skipped); } - let mut amount = 1 + state.rand_mut().below(16); - let offset = state.rand_mut().below(size + 1); + let mut amount = 1 + state.rand_mut().below(NonZero::new(16).unwrap()); + // # Safety + // It's a safe assumption that size + 1 is never 0. + // If we wrap around we have _a lot_ of elements - and the code will break later anyway. + let offset = state + .rand_mut() + .below(unsafe { NonZero::new(size + 1).unwrap_unchecked() }); if size + amount > max_size { if max_size > size { @@ -559,7 +585,11 @@ where } } - let val = input.bytes()[state.rand_mut().below(size)]; + // # Safety + // size is larger than 0, checked above. + let val = input.bytes()[state + .rand_mut() + .below(unsafe { NonZero::new(size).unwrap_unchecked() })]; input.resize(size + amount, 0); unsafe { @@ -602,8 +632,12 @@ where return Ok(MutationResult::Skipped); } - let mut amount = 1 + state.rand_mut().below(16); - let offset = state.rand_mut().below(size + 1); + let mut amount = 1 + state.rand_mut().below(NonZero::new(16).unwrap()); + // # Safety + // size + 1 can never be 0 + let offset = state + .rand_mut() + .below(unsafe { NonZero::new(size.wrapping_add(1)).unwrap_unchecked() }); if size + amount > max_size { if max_size > size { @@ -654,7 +688,11 @@ where if size == 0 { return Ok(MutationResult::Skipped); } - let range = rand_range(state, size, min(size, 16)); + // # Safety + // Size is larger than 0, checked above (and 16 is also lager than 0 FWIW) + let range = rand_range(state, size, unsafe { + NonZero::new(min(size, 16)).unwrap_unchecked() + }); let val = *state.rand_mut().choose(input.bytes()).unwrap(); let quantity = range.len(); @@ -693,7 +731,11 @@ where if size == 0 { return Ok(MutationResult::Skipped); } - let range = rand_range(state, size, min(size, 16)); + // # Safety + // Size is larger than 0, checked above. 16 is larger than 0, according to my math teacher. + let range = rand_range(state, size, unsafe { + NonZero::new(min(size, 16)).unwrap_unchecked() + }); let val = state.rand_mut().next() as u8; let quantity = range.len(); @@ -733,8 +775,16 @@ where return Ok(MutationResult::Skipped); } - let target = state.rand_mut().below(size); - let range = rand_range(state, size, size - target); + // # Safety + // size is always larger than 0 here (checked above) + let target = state + .rand_mut() + .below(unsafe { NonZero::new(size).unwrap_unchecked() }); + // # Safety + // target is smaller than size (`below` is exclusive) -> The subtraction is always larger than 0 + let range = rand_range(state, size, unsafe { + NonZero::new(size - target).unwrap_unchecked() + }); unsafe { buffer_self_copy(input.bytes_mut(), range.start, target, range.len()); @@ -776,10 +826,20 @@ where return Ok(MutationResult::Skipped); } - let target = state.rand_mut().below(size); + // # Safety + // We checked that size is larger than 0 above. + let target = state + .rand_mut() + .below(unsafe { NonZero::new(size).unwrap_unchecked() }); // make sure that the sampled range is both in bounds and of an acceptable size let max_insert_len = min(size - target, state.max_size() - size); - let range = rand_range(state, size, min(16, max_insert_len)); + let max_insert_len = min(16, max_insert_len); + + let Some(max_insert_len) = NonZero::new(max_insert_len) else { + return Ok(MutationResult::Skipped); + }; + + let range = rand_range(state, size, max_insert_len); input.resize(size + range.len(), 0); self.tmp_buf.resize(range.len(), 0); @@ -837,11 +897,19 @@ where return Ok(MutationResult::Skipped); } - let first = rand_range(state, size, size); + // # Safety + // size is larger than 0, checked above. + let first = rand_range(state, size, unsafe { + NonZero::new(size).unwrap_unchecked() + }); if state.rand_mut().next() & 1 == 0 && first.start != 0 { // The second range comes before first. - let second = rand_range(state, first.start, first.start); + // # Safety + // first.start is larger than 0, checked above. + let second = rand_range(state, first.start, unsafe { + NonZero::new(first.start).unwrap_unchecked() + }); self.tmp_buf.resize(first.len(), 0); unsafe { // If range first is larger @@ -922,7 +990,17 @@ where Ok(MutationResult::Mutated) } else if first.end != size { // The first range comes before the second range - let mut second = rand_range(state, size - first.end, size - first.end); + debug_assert!( + first.end < size, + "First.end ({}) should never be larger than size ({})!", + first.end, + size + ); + // # Safety + // first.end is not equal to size, so subtracting them can never be 0. + let mut second = rand_range(state, size - first.end, unsafe { + NonZero::new(size - first.end).unwrap_unchecked() + }); second.start += first.end; second.end += first.end; @@ -1063,6 +1141,10 @@ where { fn mutate(&mut self, state: &mut S, input: &mut I) -> Result { let size = input.bytes().len(); + let Some(nonzero_size) = NonZero::new(size) else { + return Ok(MutationResult::Skipped); + }; + let max_size = state.max_size(); if size >= max_size { return Ok(MutationResult::Skipped); @@ -1085,8 +1167,13 @@ where return Ok(MutationResult::Skipped); } - let range = rand_range(state, other_size, min(other_size, max_size - size)); - let target = state.rand_mut().below(size); // TODO: fix bug if size is 0 + // # Safety + // other_size is checked above. + // size is smaller than max_size (also checked above) -> the subtraction result is larger than 0. + let range = rand_range(state, other_size, unsafe { + NonZero::new(min(other_size, max_size - size)).unwrap_unchecked() + }); + let target = state.rand_mut().below(nonzero_size); let other_testcase = state.corpus().get_from_all(id)?.borrow_mut(); // No need to load the input again, it'll still be cached. @@ -1167,8 +1254,17 @@ where return Ok(MutationResult::Skipped); } - let target = state.rand_mut().below(size); - let range = rand_range(state, other_size, min(other_size, size - target)); + // # Safety + // Size is > 0 here (checked above) + let target = state + .rand_mut() + .below(unsafe { NonZero::new(size).unwrap_unchecked() }); + // # Safety + // other_size is checked above. + // target is smaller than size (since below is exclusive) -> the subtraction result is larger than 0. + let range = rand_range(state, other_size, unsafe { + NonZero::new(min(other_size, size - target)).unwrap_unchecked() + }); let other_testcase = state.corpus().get_from_all(id)?.borrow_mut(); // No need to load the input again, it'll still be cached. @@ -1251,7 +1347,8 @@ where fn mutate(&mut self, state: &mut S, input: &mut I) -> Result { let size = input.bytes().len(); let max_size = state.max_size(); - if size >= max_size { + // TODO: fix bug if size is 0 (?) + if size >= max_size || size == 0 { return Ok(MutationResult::Skipped); } @@ -1274,8 +1371,17 @@ where return Ok(MutationResult::Skipped); } - let range = rand_range(state, other_size, min(other_size, max_size - size)); - let target = state.rand_mut().below(size); // TODO: fix bug if size is 0 + // # Safety + // other_size is checked to be larger than 0 + // max_size is checked to be larger than size, so the subtraction will always be positive and non-0 + let range = rand_range(state, other_size, unsafe { + NonZero::new(min(other_size, max_size - size)).unwrap_unchecked() + }); + // # Safety + // size is checked above to never be 0. + let target = state + .rand_mut() + .below(unsafe { NonZero::new(size).unwrap_unchecked() }); let other_testcase = state.corpus().get_from_all(id)?.borrow_mut(); // No need to load the input again, it'll still be cached. @@ -1353,8 +1459,17 @@ where return Ok(MutationResult::Skipped); } - let target = state.rand_mut().below(size); - let range = rand_range(state, other_size, min(other_size, size - target)); + // # Safety + // We checked for size == 0 above. + let target = state + .rand_mut() + .below(unsafe { NonZero::new(size).unwrap_unchecked() }); + // # Safety + // other_size is checked above to not be 0. + // size is larger than target since below is exclusive -> subtraction is always non-0. + let range = rand_range(state, other_size, unsafe { + NonZero::new(min(other_size, size - target)).unwrap_unchecked() + }); let other_testcase = state.corpus().get_from_all(id)?.borrow_mut(); // No need to load the input again, it'll still be cached. diff --git a/libafl/src/mutators/scheduled.rs b/libafl/src/mutators/scheduled.rs index 66729b2627..dfa3478069 100644 --- a/libafl/src/mutators/scheduled.rs +++ b/libafl/src/mutators/scheduled.rs @@ -3,6 +3,7 @@ use alloc::{borrow::Cow, vec::Vec}; use core::{ fmt::Debug, + num::{NonZero, NonZeroUsize}, ops::{Deref, DerefMut}, }; @@ -101,7 +102,7 @@ where pub struct StdScheduledMutator { name: Cow<'static, str>, mutations: MT, - max_stack_pow: usize, + max_stack_pow: NonZeroUsize, } impl Named for StdScheduledMutator { @@ -148,8 +149,13 @@ where /// Get the next mutation to apply fn schedule(&self, state: &mut S, _: &I) -> MutationId { - debug_assert!(self.mutations.len() != 0); - state.rand_mut().below(self.mutations.len()).into() + debug_assert_ne!(self.mutations.len(), 0); + // # Safety + // We check for empty mutations + state + .rand_mut() + .below(unsafe { NonZero::new(self.mutations.len()).unwrap_unchecked() }) + .into() } } @@ -165,20 +171,27 @@ where mutations.names().join(", ") )), mutations, - max_stack_pow: 7, + max_stack_pow: NonZero::new(7).unwrap(), } } /// Create a new [`StdScheduledMutator`] instance specifying mutations and the maximun number of iterations - pub fn with_max_stack_pow(mutations: MT, max_stack_pow: usize) -> Self { - StdScheduledMutator { + /// + /// # Errors + /// Will return [`Error::IllegalArgument`] for `max_stack_pow` of 0. + #[inline] + pub fn with_max_stack_pow(mutations: MT, max_stack_pow: usize) -> Result { + let Some(max_stack_pow) = NonZero::new(max_stack_pow) else { + return Err(Error::illegal_argument("Max stack pow may not be 0.")); + }; + Ok(Self { name: Cow::from(format!( "StdScheduledMutator[{}]", mutations.names().join(", ") )), mutations, max_stack_pow, - } + }) } } @@ -253,15 +266,17 @@ where { /// Compute the number of iterations used to apply stacked mutations fn iterations(&self, state: &mut S, _: &I) -> u64 { - 1 << (1 + state.rand_mut().below(6)) + 1 << (1 + state.rand_mut().below(NonZero::new(6).unwrap())) } /// Get the next mutation to apply fn schedule(&self, state: &mut S, _: &I) -> MutationId { debug_assert!(::LEN != 0); + // # Safety + // In debug we check the length. Worst case we end up with an illegal MutationId and fail later. state .rand_mut() - .below(::LEN) + .below(unsafe { NonZero::new(::LEN).unwrap_unchecked() }) .into() } diff --git a/libafl/src/mutators/token_mutations.rs b/libafl/src/mutators/token_mutations.rs index eb112fae31..724821fef8 100644 --- a/libafl/src/mutators/token_mutations.rs +++ b/libafl/src/mutators/token_mutations.rs @@ -6,6 +6,7 @@ use core::slice::from_raw_parts; use core::{ fmt::Debug, mem::size_of, + num::NonZero, ops::{Add, AddAssign, Deref}, slice::Iter, }; @@ -313,15 +314,18 @@ where let Some(meta) = state.metadata_map().get::() else { return Ok(MutationResult::Skipped); }; - if meta.tokens().is_empty() { + if let Some(tokens_len) = NonZero::new(meta.tokens().len()) { + tokens_len + } else { return Ok(MutationResult::Skipped); } - meta.tokens().len() }; let token_idx = state.rand_mut().below(tokens_len); let size = input.bytes().len(); - let off = state.rand_mut().below(size + 1); + let off = state + .rand_mut() + .below(NonZero::new(size.wrapping_add(1)).unwrap()); let meta = state.metadata_map().get::().unwrap(); let token = &meta.tokens()[token_idx]; @@ -372,22 +376,23 @@ where { fn mutate(&mut self, state: &mut S, input: &mut I) -> Result { let size = input.bytes().len(); - if size == 0 { + let Some(nonzero_size) = NonZero::new(size) else { return Ok(MutationResult::Skipped); - } + }; let tokens_len = { let Some(meta) = state.metadata_map().get::() else { return Ok(MutationResult::Skipped); }; - if meta.tokens().is_empty() { + if let Some(tokens_len) = NonZero::new(meta.tokens().len()) { + tokens_len + } else { return Ok(MutationResult::Skipped); } - meta.tokens().len() }; let token_idx = state.rand_mut().below(tokens_len); - let off = state.rand_mut().below(size); + let off = state.rand_mut().below(nonzero_size); let meta = state.metadata_map().get::().unwrap(); let token = &meta.tokens()[token_idx]; @@ -432,20 +437,22 @@ where #[allow(clippy::too_many_lines)] fn mutate(&mut self, state: &mut S, input: &mut I) -> Result { let size = input.bytes().len(); - if size == 0 { + let Some(size) = NonZero::new(size) else { return Ok(MutationResult::Skipped); - } + }; let cmps_len = { let Some(meta) = state.metadata_map().get::() else { return Ok(MutationResult::Skipped); }; log::trace!("meta: {:x?}", meta); - if meta.list.is_empty() { - return Ok(MutationResult::Skipped); - } meta.list.len() }; + + let Some(cmps_len) = NonZero::new(cmps_len) else { + return Ok(MutationResult::Skipped); + }; + let idx = state.rand_mut().below(cmps_len); let off = state.rand_mut().below(size); @@ -615,7 +622,11 @@ where S: HasRand, { let sz_log = SZ.ilog2() as usize; - let res = state.rand_mut().below_incl(sz_log); + // # Safety + // We add 1 so this can never be 0. + // On 32 bit systems this could overflow in theory but this is highly unlikely. + let sz_log_inclusive = unsafe { NonZero::new(sz_log + 1).unwrap_unchecked() }; + let res = state.rand_mut().below(sz_log_inclusive); 2_usize.pow(res as u32) } @@ -626,20 +637,16 @@ where { #[allow(clippy::too_many_lines)] fn mutate(&mut self, state: &mut S, input: &mut I) -> Result { - let size = input.bytes().len(); - if size == 0 { + let Some(size) = NonZero::new(input.bytes().len()) else { return Ok(MutationResult::Skipped); - } + }; + let Some(meta) = state.metadata_map().get::() else { + return Ok(MutationResult::Skipped); + }; + log::trace!("meta: {:x?}", meta); - let cmps_len = { - let Some(meta) = state.metadata_map().get::() else { - return Ok(MutationResult::Skipped); - }; - log::trace!("meta: {:x?}", meta); - if meta.list.is_empty() { - return Ok(MutationResult::Skipped); - } - meta.list.len() + let Some(cmps_len) = NonZero::new(meta.list.len()) else { + return Ok(MutationResult::Skipped); }; let idx = state.rand_mut().below(cmps_len); diff --git a/libafl/src/mutators/tuneable.rs b/libafl/src/mutators/tuneable.rs index 449c49c502..6db2bb4e2f 100644 --- a/libafl/src/mutators/tuneable.rs +++ b/libafl/src/mutators/tuneable.rs @@ -4,7 +4,10 @@ //! a specific mutator for a specified amount of iterations use alloc::{borrow::Cow, vec::Vec}; -use core::fmt::Debug; +use core::{ + fmt::Debug, + num::{NonZero, NonZeroUsize}, +}; use libafl_bolts::{ impl_serdeany, math::calculate_cumulative_distribution_in_place, rands::Rand, @@ -83,7 +86,7 @@ impl TuneableScheduledMutatorMetadata { pub struct TuneableScheduledMutator { name: Cow<'static, str>, mutations: MT, - max_stack_pow: usize, + max_stack_pow: NonZeroUsize, } impl Mutator for TuneableScheduledMutator @@ -153,7 +156,6 @@ where /// Get the next mutation to apply fn schedule(&self, state: &mut S, _: &I) -> MutationId { - debug_assert!(self.mutations.len() != 0); // Assumption: we can not reach this code path without previously adding this metadatum. let metadata = TuneableScheduledMutatorMetadata::get_mut(state).unwrap(); @@ -196,7 +198,10 @@ where } // fall back to random if no entries in either vec, the scheduling is not tuned. - state.rand_mut().below(self.mutations.len()).into() + state + .rand_mut() + .below(NonZero::new(self.mutations.len()).expect("No mutations provided!")) + .into() } } @@ -213,7 +218,7 @@ impl TuneableScheduledMutator { TuneableScheduledMutator { name: Cow::from(format!("TuneableMutator[{}]", mutations.names().join(", "))), mutations, - max_stack_pow: 7, + max_stack_pow: NonZero::new(7).unwrap(), } } } diff --git a/libafl/src/mutators/unicode/mod.rs b/libafl/src/mutators/unicode/mod.rs index 20921e0251..c5570f40ab 100644 --- a/libafl/src/mutators/unicode/mod.rs +++ b/libafl/src/mutators/unicode/mod.rs @@ -3,6 +3,7 @@ use alloc::{borrow::Cow, vec::Vec}; use core::{ cmp::{Ordering, Reverse}, + num::NonZero, ops::Range, }; @@ -68,7 +69,9 @@ fn choose_start( bytes: &[u8], meta: &UnicodeIdentificationMetadata, ) -> Option<(usize, usize)> { - let idx = rand.below(bytes.len()); + let bytes_len = NonZero::new(bytes.len())?; + + let idx = rand.below(bytes_len); let mut options = Vec::new(); for (start, range) in meta.ranges() { if idx @@ -82,12 +85,15 @@ fn choose_start( match options.len() { 0 => None, 1 => Some((options[0].0, options[0].1.len())), - _ => { + options_len => { + // # Safety + // options.len() is checked above. + let options_len_squared = + unsafe { NonZero::new(options_len * options_len).unwrap_unchecked() }; // bias towards longer strings options.sort_by_cached_key(|(_, entries)| entries.count_ones()); let selected = - libafl_bolts::math::integer_sqrt(rand.below(options.len() * options.len()) as u64) - as usize; + libafl_bolts::math::integer_sqrt(rand.below(options_len_squared) as u64) as usize; Some((options[selected].0, options[selected].1.len())) } } @@ -135,7 +141,8 @@ fn choose_category_range( string: &str, ) -> (Range, &'static [(u32, u32)]) { let chars = string.char_indices().collect::>(); - let idx = rand.below(chars.len()); + let chars_len = NonZero::new(chars.len()).expect("Got empty string in choose_category_range"); + let idx = rand.below(chars_len); let c = chars[idx].1; // figure out the categories for this char @@ -161,7 +168,8 @@ fn choose_category_range( .sum::(), ) }); - let options = categories.len() * categories.len(); + let options = NonZero::new(categories.len() * categories.len()) + .expect("Empty categories in choose_category_range"); let selected_idx = libafl_bolts::math::integer_sqrt(rand.below(options) as u64) as usize; let selected = categories[selected_idx]; @@ -179,7 +187,8 @@ fn choose_category_range( fn choose_subcategory_range(rand: &mut R, string: &str) -> (Range, (u32, u32)) { let chars = string.char_indices().collect::>(); - let idx = rand.below(chars.len()); + let idx = + rand.below(NonZero::new(chars.len()).expect("Empty string in choose_subcategory_range")); let c = chars[idx].1; // figure out the categories for this char @@ -198,7 +207,8 @@ fn choose_subcategory_range(rand: &mut R, string: &str) -> (Range char>( range: Range, char_gen: F, ) -> MutationResult { - let temp_range = rand_range(state, range.end - range.start, MAX_CHARS); + let temp_range = rand_range( + state, + range.end - range.start, + NonZero::new(MAX_CHARS).unwrap(), + ); let range = (range.start + temp_range.start)..(range.start + temp_range.end); let range = match core::str::from_utf8(&input.0.bytes()[range.clone()]) { Ok(_) => range, @@ -240,7 +254,7 @@ fn rand_replace_range char>( return MutationResult::Skipped; } - let replace_len = state.rand_mut().below(MAX_CHARS); + let replace_len = state.rand_mut().below(NonZero::new(MAX_CHARS).unwrap()); let orig_len = range.end - range.start; if input.0.len() - orig_len + replace_len > state.max_size() { return MutationResult::Skipped; @@ -305,7 +319,10 @@ where .map(|&(start, end)| end as usize - start as usize + 1) .sum(); let char_gen = |state: &mut S| loop { - let mut selected = state.rand_mut().below(options); + // Should this skip the mutation instead of expecting? + let mut selected = state.rand_mut().below( + NonZero::new(options).expect("Empty category in UnicodeCatgoryRandMutator"), + ); for &(min, max) in category { if let Some(next_selected) = selected.checked_sub(max as usize - min as usize + 1) @@ -360,6 +377,9 @@ where ); let options = subcategory.1 as usize - subcategory.0 as usize + 1; + let Some(options) = NonZero::new(options) else { + return Ok(MutationResult::Skipped); + }; let char_gen = |state: &mut S| loop { let selected = state.rand_mut().below(options); if let Some(new_c) = char::from_u32(selected as u32 + subcategory.0) { @@ -394,15 +414,14 @@ where return Ok(MutationResult::Skipped); } - let tokens_len = { - let Some(meta) = state.metadata_map().get::() else { - return Ok(MutationResult::Skipped); - }; - if meta.tokens().is_empty() { - return Ok(MutationResult::Skipped); - } - meta.tokens().len() + let Some(meta) = state.metadata_map().get::() else { + return Ok(MutationResult::Skipped); }; + + let Some(tokens_len) = NonZero::new(meta.tokens().len()) else { + return Ok(MutationResult::Skipped); + }; + let token_idx = state.rand_mut().below(tokens_len); let bytes = input.0.bytes(); @@ -454,15 +473,14 @@ where return Ok(MutationResult::Skipped); } - let tokens_len = { - let Some(meta) = state.metadata_map().get::() else { - return Ok(MutationResult::Skipped); - }; - if meta.tokens().is_empty() { - return Ok(MutationResult::Skipped); - } - meta.tokens().len() + let Some(meta) = state.metadata_map().get::() else { + return Ok(MutationResult::Skipped); }; + + let Some(tokens_len) = NonZero::new(meta.tokens().len()) else { + return Ok(MutationResult::Skipped); + }; + let token_idx = state.rand_mut().below(tokens_len); let bytes = input.0.bytes(); diff --git a/libafl/src/stages/colorization.rs b/libafl/src/stages/colorization.rs index d0b903c9b6..69a45b29f3 100644 --- a/libafl/src/stages/colorization.rs +++ b/libafl/src/stages/colorization.rs @@ -4,7 +4,7 @@ use alloc::{ collections::binary_heap::BinaryHeap, vec::Vec, }; -use core::{cmp::Ordering, fmt::Debug, marker::PhantomData, ops::Range}; +use core::{cmp::Ordering, fmt::Debug, marker::PhantomData, num::NonZero, ops::Range}; use libafl_bolts::{ rands::Rand, @@ -363,11 +363,11 @@ where let c = match bytes[idx] { 0x41..=0x46 => { // 'A' + 1 + rand('F' - 'A') - 0x41 + 1 + state.rand_mut().below(5) as u8 + 0x41 + 1 + state.rand_mut().below(NonZero::new(5).unwrap()) as u8 } 0x61..=0x66 => { // 'a' + 1 + rand('f' - 'a') - 0x61 + 1 + state.rand_mut().below(5) as u8 + 0x61 + 1 + state.rand_mut().below(NonZero::new(5).unwrap()) as u8 } 0x30 => { // '0' -> '1' @@ -379,35 +379,35 @@ where } 0x32..=0x39 => { // '2' + 1 + rand('9' - '2') - 0x32 + 1 + state.rand_mut().below(7) as u8 + 0x32 + 1 + state.rand_mut().below(NonZero::new(7).unwrap()) as u8 } 0x47..=0x5a => { // 'G' + 1 + rand('Z' - 'G') - 0x47 + 1 + state.rand_mut().below(19) as u8 + 0x47 + 1 + state.rand_mut().below(NonZero::new(19).unwrap()) as u8 } 0x67..=0x7a => { // 'g' + 1 + rand('z' - 'g') - 0x67 + 1 + state.rand_mut().below(19) as u8 + 0x67 + 1 + state.rand_mut().below(NonZero::new(19).unwrap()) as u8 } 0x21..=0x2a => { // '!' + 1 + rand('*' - '!'); - 0x21 + 1 + state.rand_mut().below(9) as u8 + 0x21 + 1 + state.rand_mut().below(NonZero::new(9).unwrap()) as u8 } 0x2c..=0x2e => { // ',' + 1 + rand('.' - ',') - 0x2c + 1 + state.rand_mut().below(2) as u8 + 0x2c + 1 + state.rand_mut().below(NonZero::new(2).unwrap()) as u8 } 0x3a..=0x40 => { // ':' + 1 + rand('@' - ':') - 0x3a + 1 + state.rand_mut().below(6) as u8 + 0x3a + 1 + state.rand_mut().below(NonZero::new(6).unwrap()) as u8 } 0x5b..=0x60 => { // '[' + 1 + rand('`' - '[') - 0x5b + 1 + state.rand_mut().below(5) as u8 + 0x5b + 1 + state.rand_mut().below(NonZero::new(5).unwrap()) as u8 } 0x7b..=0x7e => { // '{' + 1 + rand('~' - '{') - 0x7b + 1 + state.rand_mut().below(3) as u8 + 0x7b + 1 + state.rand_mut().below(NonZero::new(3).unwrap()) as u8 } 0x2b => { // '+' -> '/' diff --git a/libafl/src/stages/mutational.rs b/libafl/src/stages/mutational.rs index 311da95f8b..f1beffacf3 100644 --- a/libafl/src/stages/mutational.rs +++ b/libafl/src/stages/mutational.rs @@ -5,7 +5,10 @@ use alloc::{ borrow::{Cow, ToOwned}, string::ToString, }; -use core::marker::PhantomData; +use core::{ + marker::PhantomData, + num::{NonZero, NonZeroUsize}, +}; use libafl_bolts::{rands::Rand, Named}; @@ -162,7 +165,7 @@ pub struct StdMutationalStage { /// The mutator(s) to use mutator: M, /// The maximum amount of iterations we should do each round - max_iterations: usize, + max_iterations: NonZeroUsize, #[allow(clippy::type_complexity)] phantom: PhantomData<(E, EM, I, Z)>, } @@ -259,11 +262,13 @@ where { /// Creates a new default mutational stage pub fn new(mutator: M) -> Self { - Self::transforming_with_max_iterations(mutator, DEFAULT_MUTATIONAL_MAX_ITERATIONS) + // Safe to unwrap: DEFAULT_MUTATIONAL_MAX_ITERATIONS is never 0. + Self::transforming_with_max_iterations(mutator, DEFAULT_MUTATIONAL_MAX_ITERATIONS).unwrap() } /// Creates a new mutational stage with the given max iterations - pub fn with_max_iterations(mutator: M, max_iterations: usize) -> Self { + #[inline] + pub fn with_max_iterations(mutator: M, max_iterations: usize) -> Result { Self::transforming_with_max_iterations(mutator, max_iterations) } } @@ -278,25 +283,36 @@ where { /// Creates a new transforming mutational stage with the default max iterations pub fn transforming(mutator: M) -> Self { - Self::transforming_with_max_iterations(mutator, DEFAULT_MUTATIONAL_MAX_ITERATIONS) + // Safe to unwrap: DEFAULT_MUTATIONAL_MAX_ITERATIONS is never 0. + Self::transforming_with_max_iterations(mutator, DEFAULT_MUTATIONAL_MAX_ITERATIONS).unwrap() } /// Creates a new transforming mutational stage with the given max iterations - pub fn transforming_with_max_iterations(mutator: M, max_iterations: usize) -> Self { + /// + /// # Errors + /// Will return [`Error::IllegalArgument`] for `max_iterations` of 0. + #[inline] + pub fn transforming_with_max_iterations( + mutator: M, + max_iterations: usize, + ) -> Result { // unsafe but impossible that you create two threads both instantiating this instance + let Some(max_iterations) = NonZero::new(max_iterations) else { + return Err(Error::illegal_argument("0 max iterations is not allowed..")); + }; let stage_id = unsafe { let ret = MUTATIONAL_STAGE_ID; MUTATIONAL_STAGE_ID += 1; ret }; - Self { - name: Cow::Owned( - MUTATIONAL_STAGE_NAME.to_owned() + ":" + stage_id.to_string().as_str(), - ), + let name = + Cow::Owned(MUTATIONAL_STAGE_NAME.to_owned() + ":" + stage_id.to_string().as_str()); + Ok(Self { + name, mutator, max_iterations, phantom: PhantomData, - } + }) } } diff --git a/libafl/src/stages/push/mutational.rs b/libafl/src/stages/push/mutational.rs index be414321bf..b3900537a9 100644 --- a/libafl/src/stages/push/mutational.rs +++ b/libafl/src/stages/push/mutational.rs @@ -5,6 +5,7 @@ use alloc::rc::Rc; use core::{ cell::{Cell, RefCell}, fmt::Debug, + num::NonZero, }; use libafl_bolts::rands::Rand; @@ -71,7 +72,9 @@ where /// Gets the number of iterations as a random number #[allow(clippy::unused_self, clippy::unnecessary_wraps)] // TODO: we should put this function into a trait later fn iterations(&self, state: &mut Z::State, _corpus_id: CorpusId) -> Result { - Ok(1 + state.rand_mut().below(DEFAULT_MUTATIONAL_MAX_ITERATIONS)) + Ok(1 + state + .rand_mut() + .below(NonZero::new(DEFAULT_MUTATIONAL_MAX_ITERATIONS).unwrap())) } /// Sets the current corpus index diff --git a/libafl/src/stages/tuneable.rs b/libafl/src/stages/tuneable.rs index 622d143888..f577d44e14 100644 --- a/libafl/src/stages/tuneable.rs +++ b/libafl/src/stages/tuneable.rs @@ -1,7 +1,7 @@ //! A [`crate::stages::MutationalStage`] where the mutator iteration can be tuned at runtime use alloc::string::{String, ToString}; -use core::{marker::PhantomData, time::Duration}; +use core::{marker::PhantomData, num::NonZero, time::Duration}; use libafl_bolts::{current_time, impl_serdeany, rands::Rand}; use serde::{Deserialize, Serialize}; @@ -248,7 +248,9 @@ where fn iterations(&self, state: &mut Self::State) -> Result { Ok( // fall back to random - 1 + state.rand_mut().below(DEFAULT_MUTATIONAL_MAX_ITERATIONS), + 1 + state + .rand_mut() + .below(NonZero::new(DEFAULT_MUTATIONAL_MAX_ITERATIONS).unwrap()), ) } } diff --git a/libafl_bolts/src/llmp.rs b/libafl_bolts/src/llmp.rs index 29875669ef..d29763bdcd 100644 --- a/libafl_bolts/src/llmp.rs +++ b/libafl_bolts/src/llmp.rs @@ -97,7 +97,7 @@ use tuple_list::tuple_list; #[cfg(all(unix, not(miri)))] use crate::os::unix_signals::setup_signal_handler; #[cfg(unix)] -use crate::os::unix_signals::{siginfo_t, ucontext_t, Handler, Signal}; +use crate::os::unix_signals::{siginfo_t, ucontext_t, Signal, SignalHandler}; #[cfg(all(windows, feature = "std"))] use crate::os::windows_exceptions::{setup_ctrl_handler, CtrlHandler}; #[cfg(feature = "std")] @@ -2187,8 +2187,8 @@ pub struct LlmpShutdownSignalHandler { } #[cfg(unix)] -impl Handler for LlmpShutdownSignalHandler { - fn handle( +impl SignalHandler for LlmpShutdownSignalHandler { + unsafe fn handle( &mut self, _signal: Signal, _info: &mut siginfo_t, diff --git a/libafl_bolts/src/os/unix_signals.rs b/libafl_bolts/src/os/unix_signals.rs index a80af001f6..ea1fa23018 100644 --- a/libafl_bolts/src/os/unix_signals.rs +++ b/libafl_bolts/src/os/unix_signals.rs @@ -387,16 +387,25 @@ impl Display for Signal { /// A trait for `LibAFL` signal handling #[cfg(feature = "alloc")] -pub trait Handler { +pub trait SignalHandler { /// Handle a signal - fn handle(&mut self, signal: Signal, info: &mut siginfo_t, _context: Option<&mut ucontext_t>); + /// + /// # Safety + /// This is generally not safe to call. It should only be called through the signal it was registered for. + /// Signal handling is hard, don't mess with it :). + unsafe fn handle( + &mut self, + signal: Signal, + info: &mut siginfo_t, + _context: Option<&mut ucontext_t>, + ); /// Return a list of signals to handle fn signals(&self) -> Vec; } #[cfg(feature = "alloc")] struct HandlerHolder { - handler: UnsafeCell<*mut dyn Handler>, + handler: UnsafeCell<*mut dyn SignalHandler>, } #[cfg(feature = "alloc")] @@ -452,7 +461,9 @@ unsafe fn handle_signal(sig: c_int, info: *mut siginfo_t, void: *mut c_void) { /// The handler pointer will be dereferenced, and the data the pointer points to may therefore not move. /// A lot can go south in signal handling. Be sure you know what you are doing. #[cfg(feature = "alloc")] -pub unsafe fn setup_signal_handler(handler: *mut T) -> Result<(), Error> { +pub unsafe fn setup_signal_handler( + handler: *mut T, +) -> Result<(), Error> { // First, set up our own stack to be used during segfault handling. (and specify `SA_ONSTACK` in `sigaction`) if SIGNAL_STACK_PTR.is_null() { SIGNAL_STACK_PTR = malloc(SIGNAL_STACK_SIZE); @@ -478,7 +489,7 @@ pub unsafe fn setup_signal_handler(handler: *mut T) -> Res write_volatile( addr_of_mut!(SIGNAL_HANDLERS[sig as usize]), Some(HandlerHolder { - handler: UnsafeCell::new(handler as *mut dyn Handler), + handler: UnsafeCell::new(handler as *mut dyn SignalHandler), }), ); diff --git a/libafl_bolts/src/os/windows_exceptions.rs b/libafl_bolts/src/os/windows_exceptions.rs index 1cabf9e984..e187b91413 100644 --- a/libafl_bolts/src/os/windows_exceptions.rs +++ b/libafl_bolts/src/os/windows_exceptions.rs @@ -422,9 +422,13 @@ pub static EXCEPTION_CODES_MAPPING: [ExceptionCode; 79] = [ ]; #[cfg(feature = "alloc")] -pub trait Handler { +pub trait ExceptionHandler { /// Handle an exception - fn handle( + /// + /// # Safety + /// This is generally not safe to call. It should only be called through the signal it was registered for. + /// Signal handling is hard, don't mess with it :). + unsafe fn handle( &mut self, exception_code: ExceptionCode, exception_pointers: *mut EXCEPTION_POINTERS, @@ -434,7 +438,7 @@ pub trait Handler { } struct HandlerHolder { - handler: UnsafeCell<*mut dyn Handler>, + handler: UnsafeCell<*mut dyn ExceptionHandler>, } pub const EXCEPTION_HANDLERS_SIZE: usize = 96; @@ -528,7 +532,9 @@ unsafe extern "C" fn handle_signal(_signum: i32) { /// # Safety /// Exception handlers are usually ugly, handle with care! #[cfg(feature = "alloc")] -pub unsafe fn setup_exception_handler(handler: *mut T) -> Result<(), Error> { +pub unsafe fn setup_exception_handler( + handler: *mut T, +) -> Result<(), Error> { let exceptions = (*handler).exceptions(); let mut catch_assertions = false; for exception_code in exceptions { @@ -542,7 +548,7 @@ pub unsafe fn setup_exception_handler(handler: *mut T) -> write_volatile( addr_of_mut!(EXCEPTION_HANDLERS[index]), Some(HandlerHolder { - handler: UnsafeCell::new(handler as *mut dyn Handler), + handler: UnsafeCell::new(handler as *mut dyn ExceptionHandler), }), ); } @@ -550,7 +556,7 @@ pub unsafe fn setup_exception_handler(handler: *mut T) -> write_volatile( addr_of_mut!(EXCEPTION_HANDLERS[EXCEPTION_HANDLERS_SIZE - 1]), Some(HandlerHolder { - handler: UnsafeCell::new(handler as *mut dyn Handler), + handler: UnsafeCell::new(handler as *mut dyn ExceptionHandler), }), ); compiler_fence(Ordering::SeqCst); diff --git a/libafl_bolts/src/rands/loaded_dice.rs b/libafl_bolts/src/rands/loaded_dice.rs index 004ef562ff..474515e14f 100644 --- a/libafl_bolts/src/rands/loaded_dice.rs +++ b/libafl_bolts/src/rands/loaded_dice.rs @@ -11,7 +11,7 @@ Assume we want to sample from the following distribution: `p(0)=0.5, p(1)=0.3, p use libafl_bolts::rands::{StdRand, loaded_dice::LoadedDiceSampler}; fn main() { let mut rand = StdRand::new(); - let mut sampler = LoadedDiceSampler::new(&[0.5, 0.3, 0.1, 0.1]); + let mut sampler = LoadedDiceSampler::new(&[0.5, 0.3, 0.1, 0.1]).unwrap(); let iter: usize = 100; for i in (0..iter) { println!("{}", sampler.sample(&mut rand)); @@ -25,6 +25,7 @@ Original code by @eqv, see use alloc::vec::Vec; use super::Rand; +use crate::Error; /// Helper struct for [`LoadedDiceSampler`] #[derive(Clone, Debug, PartialEq)] @@ -53,15 +54,23 @@ pub struct LoadedDiceSampler { impl LoadedDiceSampler { /// Create a new [`LoadedDiceSampler`] with the given probabilities - #[must_use] - pub fn new(probs: &[f64]) -> Self { + pub fn new(probs: &[f64]) -> Result { + if probs.is_empty() { + return Err(Error::illegal_argument( + "Tried to construct LoadedDiceSampler with empty probs array", + )); + } let entries = LoadedDiceSampler::construct_table(probs); - Self { entries } + Ok(Self { entries }) } /// Get one sample according to the predefined probabilities. pub fn sample(&mut self, rand: &mut R) -> usize { - let index = rand.below(self.entries.len()); + let len = self.entries.len(); + debug_assert_ne!(len, 0, "Lenght should never be 0 here."); + // # SAFETY + // len can never be 0 here. + let index = rand.below(unsafe { len.try_into().unwrap_unchecked() }); let coin = rand.next_float(); let entry = &self.entries[index]; if coin > entry.prob_of_val { @@ -114,7 +123,7 @@ mod tests { let base = (0..len).map(|_| rng.next_float()).collect::>(); let sum: f64 = base.iter().sum(); let base = base.iter().map(|v| v / sum).collect::>(); - let mut sampler = LoadedDiceSampler::new(&base); + let mut sampler = LoadedDiceSampler::new(&base).unwrap(); let mut res: Vec = vec![0; len]; let iter: usize = 1000000; for _ in 0..iter { diff --git a/libafl_bolts/src/rands/mod.rs b/libafl_bolts/src/rands/mod.rs index 404655bb13..55cb542102 100644 --- a/libafl_bolts/src/rands/mod.rs +++ b/libafl_bolts/src/rands/mod.rs @@ -2,7 +2,11 @@ #[cfg(all(not(feature = "std"), target_has_atomic = "ptr"))] use core::sync::atomic::{AtomicUsize, Ordering}; -use core::{debug_assert, fmt::Debug}; +use core::{ + debug_assert, + fmt::Debug, + num::{NonZero, NonZeroUsize}, +}; use serde::{de::DeserializeOwned, Deserialize, Serialize}; @@ -70,12 +74,10 @@ where // create iterator let mut iter = from.into_iter(); - if iter.len() == 0 { - return None; - } + let len = NonZero::new(iter.len())?; // pick a random, valid index - let index = fast_bound(rand, iter.len()); + let index = fast_bound(rand, len); // return the item chosen Some(iter.nth(index).unwrap()) @@ -89,9 +91,8 @@ where /// See: [An optimal algorithm for bounded random integers](https://github.com/apple/swift/pull/39143). #[inline] #[must_use] -pub fn fast_bound(rand: u64, n: usize) -> usize { - debug_assert_ne!(n, 0); - let mul = u128::from(rand).wrapping_mul(u128::from(n as u64)); +pub fn fast_bound(rand: u64, n: NonZeroUsize) -> usize { + let mul = u128::from(rand).wrapping_mul(u128::from(n.get() as u64)); (mul >> 64) as usize } @@ -125,21 +126,31 @@ pub trait Rand: Debug + Serialize + DeserializeOwned { /// Gets a value below the given bound (exclusive) #[inline] - fn below(&mut self, upper_bound_excl: usize) -> usize { + fn below(&mut self, upper_bound_excl: NonZeroUsize) -> usize { fast_bound(self.next(), upper_bound_excl) } /// Gets a value below the given bound (inclusive) #[inline] fn below_incl(&mut self, upper_bound_incl: usize) -> usize { - self.below(upper_bound_incl + 1) + let Some(upper_bound) = NonZero::new(upper_bound_incl.wrapping_add(1)) else { + // The max value + 1 wrapped around to 0. We just do a "normal" random. + return self.next() as usize; + }; + self.below(upper_bound) } /// Gets a value between the given lower bound (inclusive) and upper bound (inclusive) #[inline] fn between(&mut self, lower_bound_incl: usize, upper_bound_incl: usize) -> usize { debug_assert!(lower_bound_incl <= upper_bound_incl); - lower_bound_incl + self.below(upper_bound_incl - lower_bound_incl + 1) + // # Safety + // We check that the upper_bound_incl <= lower_bound_incl above (alas only in debug), so the below is fine. + // Even if we encounter a 0 in release here, the worst-case scenario should be an invalid return value. + lower_bound_incl + + self.below(unsafe { + NonZero::new(upper_bound_incl - lower_bound_incl + 1).unwrap_unchecked() + }) } /// Convenient variant of [`choose`]. @@ -164,17 +175,19 @@ pub trait Rand: Debug + Serialize + DeserializeOwned { // when the Iterator is an ExactSizeIterator. This has a large performance impact on e.g. // seq_iter_choose_from_1000. if upper == Some(lower) { - return if lower == 0 { - None - } else { + return if let Some(lower) = NonZero::new(lower) { iter.nth(self.below(lower)) + } else { + None }; } // Continue until the iterator is exhausted loop { if lower > 1 { - let ix = self.below(lower + consumed); + // # Safety + // lower is > 1, we don't consume more than usize elements, so this should always be non-0. + let ix = self.below(unsafe { NonZero::new(lower + consumed).unwrap_unchecked() }); let skip = if ix < lower { result = iter.nth(ix); lower - (ix + 1) @@ -194,7 +207,9 @@ pub trait Rand: Debug + Serialize + DeserializeOwned { return result; } consumed += 1; - if self.below(consumed) == 0 { + // # SAFETY + // `consumed` can never be 0 here. We just increased it by 1 above. + if self.below(unsafe { NonZero::new(consumed).unwrap_unchecked() }) == 0 { result = elem; } } @@ -533,6 +548,8 @@ impl XkcdRand { #[cfg(test)] mod tests { + use core::num::NonZero; + use crate::rands::{ Rand, RomuDuoJrRand, RomuTrioRand, Sfc64Rand, StdRand, XorShift64Rand, Xoshiro256PlusPlusRand, @@ -540,8 +557,8 @@ mod tests { fn test_single_rand(rand: &mut R) { assert_ne!(rand.next(), rand.next()); - assert!(rand.below(100) < 100); - assert_eq!(rand.below(1), 0); + assert!(rand.below(NonZero::new(100).unwrap()) < 100); + assert_eq!(rand.below(NonZero::new(1).unwrap()), 0); assert_eq!(rand.between(10, 10), 10); assert!(rand.between(11, 20) > 10); } diff --git a/libafl_bolts/src/shmem.rs b/libafl_bolts/src/shmem.rs index a78ad04aef..5e968b47e6 100644 --- a/libafl_bolts/src/shmem.rs +++ b/libafl_bolts/src/shmem.rs @@ -847,6 +847,9 @@ pub mod unix_shmem { /// /// This function will return an error if the appropriate flags could not be extracted or set. pub fn persist_for_child_processes(&self) -> Result<&Self, Error> { + // # Safety + // No user-provided potentially unsafe parameters. + // FFI Calls. unsafe { let flags = fcntl(self.shm_fd, libc::F_GETFD); @@ -926,18 +929,25 @@ pub mod unix_shmem { type Target = [u8]; fn deref(&self) -> &[u8] { + // # Safety + // No user-provided potentially unsafe parameters. unsafe { slice::from_raw_parts(self.map, self.map_size) } } } impl DerefMut for MmapShMem { fn deref_mut(&mut self) -> &mut [u8] { + // # Safety + // No user-provided potentially unsafe parameters. unsafe { slice::from_raw_parts_mut(self.map, self.map_size) } } } impl Drop for MmapShMem { fn drop(&mut self) { + // # Safety + // No user-provided potentially unsafe parameters. + // Mutable borrow so no possible race. unsafe { assert!( !self.map.is_null(), diff --git a/libafl_frida/src/lib.rs b/libafl_frida/src/lib.rs index dd55a51fa2..7bb987e662 100644 --- a/libafl_frida/src/lib.rs +++ b/libafl_frida/src/lib.rs @@ -523,7 +523,8 @@ mod tests { ); let mutator = StdScheduledMutator::new(tuple_list!(BitFlipMutator::new())); - let mut stages = tuple_list!(StdMutationalStage::with_max_iterations(mutator, 1)); + let mut stages = + tuple_list!(StdMutationalStage::with_max_iterations(mutator, 1).unwrap()); log::info!("Starting fuzzing!"); fuzzer diff --git a/libafl_libfuzzer/runtime/src/lib.rs b/libafl_libfuzzer/runtime/src/lib.rs index c2fe8dc103..0aa224597c 100644 --- a/libafl_libfuzzer/runtime/src/lib.rs +++ b/libafl_libfuzzer/runtime/src/lib.rs @@ -374,7 +374,8 @@ macro_rules! fuzz_with { let custom_mutator = unsafe { LLVMCustomMutator::mutate_unchecked(StdScheduledMutator::new(havoc_mutations_no_crossover().merge(tokens_mutations()))) }; - let std_mutator_no_mutate = StdScheduledMutator::with_max_stack_pow(havoc_crossover(), 3); + // Safe to unwrap: stack pow is not 0. + let std_mutator_no_mutate = StdScheduledMutator::with_max_stack_pow(havoc_crossover(), 3).unwrap(); let cm_power: StdPowerMutationalStage<_, _, BytesInput, _, _> = StdPowerMutationalStage::new(custom_mutator); let cm_power = IfStage::new(|_, _, _, _| Ok(mutator_status.custom_mutation.into()), (cm_power, ())); @@ -385,11 +386,12 @@ macro_rules! fuzz_with { // a custom crossover is defined // while the scenario that a custom crossover is defined without a custom mutator is unlikely // we handle it here explicitly anyways + // Safe to unwrap: stack pow is not 0. let custom_crossover = unsafe { LLVMCustomMutator::crossover_unchecked(StdScheduledMutator::with_max_stack_pow( havoc_mutations_no_crossover().merge(tokens_mutations()), 3, - )) + ).unwrap()) }; let std_mutator_no_crossover = StdScheduledMutator::new(havoc_mutations_no_crossover().merge(tokens_mutations())); @@ -399,6 +401,7 @@ macro_rules! fuzz_with { let cc_std_power = IfStage::new(|_, _, _, _| Ok(mutator_status.std_no_crossover.into()), (cc_std_power, ())); + // Safe to unwrap: stack pow is not 0. let grimoire_mutator = StdScheduledMutator::with_max_stack_pow( tuple_list!( GrimoireExtensionMutator::new(), @@ -409,7 +412,7 @@ macro_rules! fuzz_with { GrimoireRandomDeleteMutator::new(), ), 3, - ); + ).unwrap(); let grimoire = IfStage::new(|_, _, _, _| Ok(grimoire.into()), (StdMutationalStage::transforming(grimoire_mutator), ())); // A minimization+queue policy to get testcasess from the corpus @@ -464,7 +467,7 @@ macro_rules! fuzz_with { } if state.corpus().count() < 1 { // Generator of bytearrays of max size 64 - let mut generator = RandBytesGenerator::from(RandBytesGenerator::new(64)); + let mut generator = RandBytesGenerator::from(RandBytesGenerator::new(64).unwrap()); // Generate 1024 initial inputs state diff --git a/libafl_sugar/src/forkserver.rs b/libafl_sugar/src/forkserver.rs index 6166ea5796..35a96b3abf 100644 --- a/libafl_sugar/src/forkserver.rs +++ b/libafl_sugar/src/forkserver.rs @@ -213,7 +213,7 @@ impl<'a> ForkserverBytesCoverageSugar<'a> { if state.must_load_initial_inputs() { if self.input_dirs.is_empty() { // Generator of printable bytearrays of max size 32 - let mut generator = RandBytesGenerator::new(32); + let mut generator = RandBytesGenerator::new(32).unwrap(); // Generate 8 initial inputs state diff --git a/libafl_sugar/src/inmemory.rs b/libafl_sugar/src/inmemory.rs index cd8728f8b3..9643741c62 100644 --- a/libafl_sugar/src/inmemory.rs +++ b/libafl_sugar/src/inmemory.rs @@ -229,7 +229,7 @@ where if state.must_load_initial_inputs() { if self.input_dirs.is_empty() { // Generator of printable bytearrays of max size 32 - let mut generator = RandBytesGenerator::new(32); + let mut generator = RandBytesGenerator::new(32).unwrap(); // Generate 8 initial inputs state diff --git a/libafl_sugar/src/qemu.rs b/libafl_sugar/src/qemu.rs index f92f3e7d6d..e97bed43ec 100644 --- a/libafl_sugar/src/qemu.rs +++ b/libafl_sugar/src/qemu.rs @@ -259,7 +259,7 @@ where if state.must_load_initial_inputs() { if self.input_dirs.is_empty() { // Generator of printable bytearrays of max size 32 - let mut generator = RandBytesGenerator::new(32); + let mut generator = RandBytesGenerator::new(32).unwrap(); // Generate 8 initial inputs state @@ -372,7 +372,7 @@ where if state.must_load_initial_inputs() { if self.input_dirs.is_empty() { // Generator of printable bytearrays of max size 32 - let mut generator = RandBytesGenerator::new(32); + let mut generator = RandBytesGenerator::new(32).unwrap(); // Generate 8 initial inputs state diff --git a/utils/libafl_benches/benches/hash_speeds.rs b/utils/libafl_benches/benches/hash_speeds.rs index eb29562336..a929a6a51e 100644 --- a/utils/libafl_benches/benches/hash_speeds.rs +++ b/utils/libafl_benches/benches/hash_speeds.rs @@ -1,6 +1,9 @@ //! Compare the speed of rust hash implementations -use std::hash::{BuildHasher, Hasher}; +use std::{ + hash::{BuildHasher, Hasher}, + num::NonZero, +}; use criterion::{black_box, criterion_group, criterion_main, Criterion}; use libafl_bolts::rands::{Rand, StdRand}; @@ -11,7 +14,7 @@ fn criterion_benchmark(c: &mut Criterion) { let mut rand = StdRand::with_seed(0); let mut bench_vec: Vec = vec![]; for _ in 0..2 << 16 { - bench_vec.push(rand.below(256) as u8); + bench_vec.push(rand.below(NonZero::new(256).unwrap()) as u8); } c.bench_function("xxh3", |b| {