use std::env;
use std::io::{self, Read};
use std::panic;
extern "C" {
fn __afl_persistent_loop(counter: usize) -> isize;
fn __afl_manual_init();
static __afl_fuzz_len: *const u32;
static __afl_fuzz_ptr: *const u8;
}
#[doc(hidden)]
#[no_mangle]
pub static mut __afl_sharedmem_fuzzing: i32 = 1;
pub fn fuzz<F>(hook: bool, mut closure: F)
where
F: FnMut(&[u8]) + std::panic::RefUnwindSafe,
{
static PERSIST_MARKER: &str = "##SIG_AFL_PERSISTENT##\0";
static DEFERED_MARKER: &str = "##SIG_AFL_DEFER_FORKSRV##\0";
unsafe { std::ptr::read_volatile(&PERSIST_MARKER) }; unsafe { std::ptr::read_volatile(&DEFERED_MARKER) };
if hook {
let prev_hook = std::panic::take_hook();
std::panic::set_hook(Box::new(move |panic_info| {
prev_hook(panic_info);
std::process::abort();
}));
}
let mut input = vec![];
let loop_count = if let Ok(value) = env::var("AFL_FUZZER_LOOPCOUNT") {
value
.parse()
.expect("Failed to parse environment variable to a number")
} else {
usize::MAX
};
unsafe { __afl_manual_init() };
while unsafe { __afl_persistent_loop(loop_count) } != 0 {
let input_ref = if unsafe { __afl_fuzz_ptr.is_null() } {
let result = io::stdin().read_to_end(&mut input);
if result.is_err() {
return;
}
&input
} else {
unsafe {
let input_len = *__afl_fuzz_len as usize;
std::slice::from_raw_parts(__afl_fuzz_ptr, input_len)
}
};
let did_panic = std::panic::catch_unwind(panic::AssertUnwindSafe(|| {
closure(input_ref);
}))
.is_err();
if did_panic {
std::process::abort();
}
input.clear();
}
}
#[macro_export]
macro_rules! fuzz {
( $($x:tt)* ) => { $crate::__fuzz!(true, $($x)*) }
}
#[macro_export]
macro_rules! fuzz_nohook {
( $($x:tt)* ) => { $crate::__fuzz!(false, $($x)*) }
}
#[doc(hidden)]
#[macro_export]
macro_rules! __fuzz {
($hook:expr, |$buf:ident| $body:block) => {
$crate::fuzz($hook, |$buf| $body);
};
($hook:expr, |$buf:ident: &[u8]| $body:block) => {
$crate::fuzz($hook, |$buf| $body);
};
($hook:expr, |$buf:ident: $dty: ty| $body:block) => {
$crate::fuzz($hook, |$buf| {
let $buf: $dty = {
let mut data = ::arbitrary::Unstructured::new($buf);
if let Ok(d) = ::arbitrary::Arbitrary::arbitrary(&mut data).map_err(|_| "") {
d
} else {
return;
}
};
$body
});
};
}