From f7bb4d756740b5acf3a41031600e02204f994cd5 Mon Sep 17 00:00:00 2001 From: Evian-Zhang Date: Mon, 21 Apr 2025 20:07:46 +0800 Subject: [PATCH 1/2] Make unicornafl more Rust-friendly and sync with newest unicorn-engine Rust crate --- examples/sample.rs | 39 +++--- src/executor.rs | 302 +++++++++++++++++++++++---------------------- src/lib.rs | 247 +++++++++++++++++++++++++++++++----- src/target.rs | 148 ++++++---------------- 4 files changed, 420 insertions(+), 316 deletions(-) diff --git a/examples/sample.rs b/examples/sample.rs index 959d3268..c251640c 100644 --- a/examples/sample.rs +++ b/examples/sample.rs @@ -1,18 +1,14 @@ -use std::ffi::{c_uchar, c_void, CString}; +use std::path::PathBuf; -use unicorn_engine::{ffi::uc_handle, Arch, Mode, Permission, RegisterX86, Unicorn}; -use unicornafl::target::child_fuzz; +use unicorn_engine::{Arch, Mode, Prot, RegisterX86, Unicorn}; +use unicornafl::{afl_fuzz, executor::UnicornFuzzData}; -extern "C" fn place_input_cb( - uc: uc_handle, - input: *const c_uchar, - input_len: usize, +fn place_input_cb<'a, D: 'a>( + uc: &mut Unicorn<'a, UnicornFuzzData>, + input: &[u8], _persistent_round: u64, - _data: *mut c_void, ) -> bool { - let mut uc = unsafe { Unicorn::from_handle(uc) }.expect("fail to create inner"); let mut buf = [0; 8]; - let input = unsafe { std::slice::from_raw_parts(input, input_len) }; if input.len() < 8 { // decline the input return false; @@ -28,30 +24,23 @@ extern "C" fn place_input_cb( fn main() { let input_file = std::env::args().into_iter().skip(1).nth(0); - let mut uc = Unicorn::new(Arch::X86, Mode::MODE_64).expect("fail to open uc"); + let mut uc = Unicorn::new_with_data(Arch::X86, Mode::MODE_64, UnicornFuzzData::default()) + .expect("fail to open uc"); // ks.asm("mov rax, rdx; cmp rax, 0x114514; je die; xor rax, rax; die: mov rax, [rax]; xor rax, rax") let code = b"\x48\x89\xd0\x48\x3d\x14\x45\x11\x00\x74\x03\x48\x31\xc0\x48\x8b\x00\x48\x31\xc0"; - uc.mem_map(0x1000, 0x4000, Permission::all()) - .expect("fail to map"); + uc.mem_map(0x1000, 0x4000, Prot::ALL).expect("fail to map"); uc.mem_write(0x1000, code).expect("fail to write code"); let pc = 0x1000; uc.reg_write(RegisterX86::RIP, pc) .expect("fail to write pc"); - let input_file = input_file.map(|t| CString::new(t).expect("fail to CString")); - child_fuzz( - uc.get_handle(), - input_file - .as_ref() - .map(|t| t.as_ptr()) - .unwrap_or(std::ptr::null()), - 1, // This is not too effective but enough here for testing + let input_file = input_file.map(|t| PathBuf::from(t)); + afl_fuzz( + uc, + input_file, place_input_cb, - None, vec![0x100b, 0x1011], - None, false, - true, - std::ptr::null_mut(), + 1, ) .expect("fail to fuzz?") } diff --git a/src/executor.rs b/src/executor.rs index e9afdeb1..bd021ccf 100644 --- a/src/executor.rs +++ b/src/executor.rs @@ -1,7 +1,6 @@ use std::{ hash::{Hash, Hasher}, ops::Deref, - os::raw::c_void, }; use libafl::{ @@ -17,14 +16,9 @@ use libafl_bolts::{ use libafl_targets::EDGES_MAP_PTR; use log::{trace, warn}; use serde::{Deserialize, Serialize}; -use unicorn_engine::{ - ffi::{uc_ctl, uc_handle, uc_hook, uc_hook_add, uc_hook_del}, - uc_error, ControlType, -}; +use unicorn_engine::{uc_error, TcgOpCode, TcgOpFlag, UcHookId, Unicorn}; -use crate::{ - hash::afl_hash_ip, uc_afl_cb_place_input_t, uc_afl_cb_validate_crash_t, uc_afl_fuzz_cb_t, -}; +use crate::hash::afl_hash_ip; #[derive(Serialize, Deserialize, Debug, Clone)] pub struct UnsafeSliceInput<'a> { @@ -75,6 +69,51 @@ struct HookState { map_size: u32, } +fn get_afl_map_size() -> u32 { + std::env::var("AFL_MAP_SIZE") + .ok() + .map(|sz| u32::from_str_radix(&sz, 10).ok()) + .flatten() + .unwrap_or(1 << 16) // MAP_SIZE +} + +/// Data persisted during fuzzing. You can use `uc.get_data()`[Unicorn::get_data] +/// and `uc.get_data_mut()`[Unicorn::get_data_mut] to access this data in callbacks +/// and hooks during fuzzing. +/// +/// You can create a default fuzz data by [`UnicornFuzzData::default()`] if you don't +/// want any custom data. +#[derive(Debug)] +pub struct UnicornFuzzData { + hook_state: HookState, + /// User-defined data. + pub user_data: D, +} + +impl UnicornFuzzData { + pub(crate) fn map_size(&self) -> u32 { + self.hook_state.map_size + } +} + +impl Default for UnicornFuzzData<()> { + fn default() -> Self { + Self::new(()) + } +} + +impl UnicornFuzzData { + pub fn new(user_data: D) -> Self { + Self { + hook_state: HookState { + prev_loc: 0, + map_size: get_afl_map_size(), + }, + user_data, + } + } +} + unsafe fn update_coverage(idx: usize) { unsafe { let loc = EDGES_MAP_PTR.byte_add(idx); @@ -88,9 +127,12 @@ unsafe fn update_with_prev(loc: u32, prev: u32) { update_coverage(idx as usize); } -#[no_mangle] -extern "C" fn hook_code_coverage(_uc: uc_handle, address: u64, _size: u32, data: *mut c_void) { - let state: &mut HookState = unsafe { (data as *mut HookState).as_mut().unwrap() }; +fn hook_code_coverage<'a, D: 'a>( + uc: &mut Unicorn<'a, UnicornFuzzData>, + address: u64, + _size: u32, +) { + let state = &mut uc.get_data_mut().hook_state; let cur_loc = afl_hash_ip(address) & (state.map_size - 1); unsafe { update_with_prev(cur_loc, state.prev_loc) }; @@ -139,16 +181,14 @@ fn hook_sub_impl_64(cur_loc: u32, prev_loc: u32, arg1: u64, arg2: u64) { } } -#[no_mangle] -extern "C" fn hook_opcode_cmpcov( - _uc: uc_handle, +fn hook_opcode_cmpcov<'a, D: 'a>( + uc: &mut Unicorn<'a, UnicornFuzzData>, address: u64, arg1: u64, arg2: u64, - size: u32, - data: *mut c_void, + size: usize, ) { - let state: &mut HookState = unsafe { (data as *mut HookState).as_mut().unwrap() }; + let state = &mut uc.get_data_mut().hook_state; let mut cur_loc = afl_hash_ip(address) & (state.map_size - 1); if size >= 64 { @@ -169,124 +209,102 @@ extern "C" fn hook_opcode_cmpcov( } } -#[derive(Debug)] -pub struct UnicornAflExecutor { - uc: uc_handle, - place_input_cb: uc_afl_cb_place_input_t, - validate_crash_cb: Option, - fuzz_callback: uc_afl_fuzz_cb_t, +/// Executor for unicorn +pub struct UnicornAflExecutor<'a, D, FI, FV, FC> +where + D: 'a, + FI: FnMut(&mut Unicorn<'a, UnicornFuzzData>, &[u8], u64) -> bool + 'a, + FV: FnMut(&mut Unicorn<'a, UnicornFuzzData>, Result<(), uc_error>, &[u8], u64) -> bool + 'a, + FC: FnMut(&mut Unicorn<'a, UnicornFuzzData>) -> Result<(), uc_error> + 'a, +{ + uc: Unicorn<'a, UnicornFuzzData>, + /// Place the generated input into unicorn's memory. + /// + /// Return false if the generated input is not acceptable + place_input_cb: FI, + /// Return true if the crash is valid after validation + validate_crash_cb: FV, + /// The real procedure to kick unicorn engine start + fuzz_callback: FC, + /// Whether the `validate_crash_cb` is invoked everytime regardless of + /// the execution result. + /// + /// If false, only execution failure will lead to the callback. always_validate: bool, - data: *mut c_void, - _state: Box, - block_hook: uc_hook, - cmp_hook: uc_hook, - sub_hook: uc_hook, + /// Stored for deleting hook when dropping + block_hook: UcHookId, + /// Stored for deleting hook when dropping + sub_hook: UcHookId, dumb_ob: tuple_list_type!(ValueObserver<'static, bool>), } -impl UnicornAflExecutor { +impl<'a, D, FI, FV, FC> UnicornAflExecutor<'a, D, FI, FV, FC> +where + D: 'a, + FI: FnMut(&mut Unicorn<'a, UnicornFuzzData>, &[u8], u64) -> bool + 'a, + FV: FnMut(&mut Unicorn<'a, UnicornFuzzData>, Result<(), uc_error>, &[u8], u64) -> bool + 'a, + FC: FnMut(&mut Unicorn<'a, UnicornFuzzData>) -> Result<(), uc_error> + 'a, +{ + /// Create a new executor pub fn new( - uc: uc_handle, - place_input_cb: uc_afl_cb_place_input_t, - validate_crash_cb: Option, - fuzz_callback: uc_afl_fuzz_cb_t, + mut uc: Unicorn<'a, UnicornFuzzData>, + place_input_cb: FI, + validate_crash_cb: FV, + fuzz_callback: FC, always_validate: bool, exits: Vec, - map_size: u32, - data: *mut c_void, ) -> Result { - let mut block_hook = std::ptr::null_mut(); - let mut cmp_hook = std::ptr::null_mut(); - let mut sub_hook = std::ptr::null_mut(); - let mut state = Box::new(HookState { - prev_loc: 0, - map_size, - }); - unsafe { + if !exits.is_empty() { // Enable exits if requested - if exits.len() > 0 { - let ret = uc_ctl( - uc, - ControlType::UC_CTL_IO_WRITE as u32 | ControlType::UC_CTL_UC_USE_EXITS as u32, - 1u32, - ); - if ret != uc_error::OK { - warn!("Fail to enable exits due to {:?}", ret); - return Err(ret); - } - let ret = uc_ctl( - uc, - ControlType::UC_CTL_IO_WRITE as u32 | ControlType::UC_CTL_UC_EXITS as u32, - exits.as_ptr(), - exits.len(), - ); - if ret != uc_error::OK { - warn!("Fail to write exits due to {:?}", ret); - return Err(ret); - } - } + uc.ctl_exits_enable().inspect_err(|ret| { + warn!("Fail to enable exits due to {ret:?}"); + })?; + uc.ctl_set_exits(&exits).inspect_err(|ret| { + warn!("Fail to write exits due to {ret:?}"); + })?; + } - let ret = uc_hook_add( - uc, - &mut block_hook, - unicorn_engine::HookType::BLOCK, - hook_code_coverage as _, - state.as_mut() as *mut HookState as _, - 1, - 0, - ); - if ret != uc_error::OK { - warn!("Fail to add block hooks due to {:?}", ret); - return Err(ret); - } - let ret = uc_hook_add( - uc, - &mut cmp_hook, - unicorn_engine::HookType::TCG_OPCODE, - hook_opcode_cmpcov as _, - state.as_mut() as *mut HookState as _, - 1, - 0, - unicorn_engine::TcgOp::SUB, - unicorn_engine::TcgOpFlag::CMP, - ); - if ret != uc_error::OK { - warn!("Fail to add cmp hooks due to {:?}", ret); - return Err(ret); - } - let ret = uc_hook_add( - uc, - &mut sub_hook, - unicorn_engine::HookType::TCG_OPCODE, - hook_opcode_cmpcov as _, - state.as_mut() as *mut HookState as _, + let block_hook = uc + .add_block_hook(1, 0, |uc, address, size| { + hook_code_coverage(uc, address, size); + }) + .inspect_err(|ret| { + warn!("Fail to add block hooks due to {ret:?}"); + })?; + let sub_hook = uc + .add_tcg_hook( + TcgOpCode::SUB, + TcgOpFlag::CMP | TcgOpFlag::DIRECT, 1, 0, - unicorn_engine::TcgOp::SUB, - unicorn_engine::TcgOpFlag::DIRECT, - ); - if ret != uc_error::OK { - warn!("Fail to add sub hooks due to {:?}", ret); - return Err(ret); - } - } + |uc, address, arg1, arg2, size| { + hook_opcode_cmpcov(uc, address, arg1, arg2, size); + }, + ) + .inspect_err(|ret| { + warn!("Fail to add cmp and sub hooks due to {ret:?}"); + })?; + Ok(Self { uc, place_input_cb, validate_crash_cb, fuzz_callback, always_validate, - data, - _state: state, block_hook, - cmp_hook, sub_hook, dumb_ob: tuple_list!(ValueObserver::new("dumb_ob", OwnedRef::Owned(false.into()))), }) } } -impl HasObservers for UnicornAflExecutor { +impl<'a, D, FI, FV, FC> HasObservers for UnicornAflExecutor<'a, D, FI, FV, FC> +where + D: 'a, + FI: FnMut(&mut Unicorn<'a, UnicornFuzzData>, &[u8], u64) -> bool + 'a, + FV: FnMut(&mut Unicorn<'a, UnicornFuzzData>, Result<(), uc_error>, &[u8], u64) -> bool + 'a, + FC: FnMut(&mut Unicorn<'a, UnicornFuzzData>) -> Result<(), uc_error> + 'a, +{ type Observers = tuple_list_type!(ValueObserver<'static, bool>); fn observers(&self) -> libafl_bolts::tuples::RefIndexable<&Self::Observers, Self::Observers> { RefIndexable::from(&self.dumb_ob) @@ -299,28 +317,31 @@ impl HasObservers for UnicornAflExecutor { } } -impl Drop for UnicornAflExecutor { +impl<'a, D, FI, FV, FC> Drop for UnicornAflExecutor<'a, D, FI, FV, FC> +where + D: 'a, + FI: FnMut(&mut Unicorn<'a, UnicornFuzzData>, &[u8], u64) -> bool + 'a, + FV: FnMut(&mut Unicorn<'a, UnicornFuzzData>, Result<(), uc_error>, &[u8], u64) -> bool + 'a, + FC: FnMut(&mut Unicorn<'a, UnicornFuzzData>) -> Result<(), uc_error> + 'a, +{ fn drop(&mut self) { - unsafe { - let ret = uc_hook_del(self.uc, self.block_hook); - if ret != uc_error::OK { - warn!("Fail to uninstall block hook due to {:?}", ret) - } - let ret = uc_hook_del(self.uc, self.cmp_hook); - if ret != uc_error::OK { - warn!("Fail to uninstall cmp tcg opcode hook due to {:?}", ret); - } - let ret = uc_hook_del(self.uc, self.sub_hook); - if ret != uc_error::OK { - warn!("Fail to uninstall sub tcg opcode hook due to {:?}", ret); - } + if let Err(ret) = self.uc.remove_hook(self.block_hook) { + warn!("Fail to uninstall block hook due to {ret:?}"); + } + if let Err(ret) = self.uc.remove_hook(self.sub_hook) { + warn!("Fail to uninstall cmp and sub tcg opcode hook due to {ret:?}"); } } } -impl<'a, EM, S, Z> Executor, S, Z> for UnicornAflExecutor +impl<'a, 'b, EM, S, Z, D, FI, FV, FC> Executor, S, Z> + for UnicornAflExecutor<'b, D, FI, FV, FC> where S: HasExecutions, + D: 'b, + FI: FnMut(&mut Unicorn<'b, UnicornFuzzData>, &[u8], u64) -> bool + 'b, + FV: FnMut(&mut Unicorn<'b, UnicornFuzzData>, Result<(), uc_error>, &[u8], u64) -> bool + 'b, + FC: FnMut(&mut Unicorn<'b, UnicornFuzzData>) -> Result<(), uc_error> + 'b, { fn run_target( &mut self, @@ -329,36 +350,19 @@ where _mgr: &mut EM, input: &UnsafeSliceInput<'a>, ) -> Result { - let accepted = (self.place_input_cb)( - self.uc, - input.as_ref().as_ptr(), - input.len(), - *state.executions(), - self.data, - ); + let accepted = (self.place_input_cb)(&mut self.uc, input.as_ref(), *state.executions()); if !accepted { trace!("Input not accepted"); return Ok(ExitKind::Ok); } - let err = (self.fuzz_callback)(self.uc, self.data); - - trace!("Child returns: {:?}", err); - - if err != uc_error::OK || self.always_validate { - if let Some(validate_cb) = self.validate_crash_cb { - if (validate_cb)( - self.uc, - err, - input.as_ref().as_ptr(), - input.len(), - *state.executions(), - self.data, - ) { - return Ok(ExitKind::Crash); - } - } else if err != uc_error::OK { + let err = (self.fuzz_callback)(&mut self.uc); + + trace!("Child returns: {err:?}"); + + if err.is_err() || self.always_validate { + if (self.validate_crash_cb)(&mut self.uc, err, input.as_ref(), *state.executions()) { return Ok(ExitKind::Crash); } } diff --git a/src/lib.rs b/src/lib.rs index b2d6826c..c140f0d0 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,10 +1,11 @@ use std::{ - ffi::c_uchar, + ffi::{c_uchar, CStr}, os::raw::{c_char, c_void}, + path::PathBuf, }; -use target::child_fuzz; -use unicorn_engine::{ffi::uc_handle, uc_error}; +use executor::UnicornFuzzData; +use unicorn_engine::{uc_error, unicorn_const::uc_engine, Unicorn}; pub mod executor; pub mod harness; @@ -41,7 +42,7 @@ impl From for uc_afl_ret { #[allow(non_camel_case_types)] pub type uc_afl_cb_place_input_t = extern "C" fn( - uc: uc_handle, + uc: *mut uc_engine, input: *const c_uchar, input_len: usize, persistent_round: u64, @@ -50,7 +51,7 @@ pub type uc_afl_cb_place_input_t = extern "C" fn( #[allow(non_camel_case_types)] pub type uc_afl_cb_validate_crash_t = extern "C" fn( - uc: uc_handle, + uc: *mut uc_engine, unicorn_result: uc_error, input: *const c_uchar, input_len: usize, @@ -59,12 +60,76 @@ pub type uc_afl_cb_validate_crash_t = extern "C" fn( ) -> bool; #[allow(non_camel_case_types)] -pub type uc_afl_fuzz_cb_t = extern "C" fn(uc: uc_handle, data: *mut c_void) -> uc_error; +pub type uc_afl_fuzz_cb_t = extern "C" fn(uc: *mut uc_engine, data: *mut c_void) -> uc_error; +/// Customized afl fuzz routine entrypoint for Rust user. +/// +/// If you want to use default crash validation callback, pass +/// [`dummy_uc_validate_crash_callback`][target::dummy_uc_validate_crash_callback]. +/// +/// If you want to use default fuzz callback, pass +/// [`dummy_uc_fuzz_callback`][target::dummy_uc_fuzz_callback]. +/// +/// `exits` means instruction addresses that stop the execution. You can pass +/// an empty vec here if there is not explicit exit. +/// +/// If `always_validate` is true, then `validate_crash_cb` is invoked everytime +/// regardless of the result of execution; Otherwise, only failed execution will +/// invoke such callback. +/// +/// `persistent_iters` is the number of persistent execution rounds. +pub fn afl_fuzz_custom<'a, D: 'a>( + uc: Unicorn<'a, UnicornFuzzData>, + input_file: Option, + place_input_cb: impl FnMut(&mut Unicorn<'a, UnicornFuzzData>, &[u8], u64) -> bool + 'a, + validate_crash_cb: impl FnMut(&mut Unicorn<'a, UnicornFuzzData>, Result<(), uc_error>, &[u8], u64) -> bool + + 'a, + fuzz_callback: impl FnMut(&mut Unicorn<'a, UnicornFuzzData>) -> Result<(), uc_error> + 'a, + exits: Vec, + always_validate: bool, + persistent_iters: u32, +) -> Result<(), uc_afl_ret> { + target::child_fuzz( + uc, + input_file, + persistent_iters, + place_input_cb, + validate_crash_cb, + fuzz_callback, + exits, + always_validate, + true, + ) +} + +/// Simplified afl fuzz routine entrypoint for Rust user. +/// +/// If you want to manually validate crash or kick fuzzing, call [`afl_fuzz_custom`]. +pub fn afl_fuzz<'a, D: 'a>( + uc: Unicorn<'a, UnicornFuzzData>, + input_file: Option, + place_input_cb: impl FnMut(&mut Unicorn<'a, UnicornFuzzData>, &[u8], u64) -> bool + 'a, + exits: Vec, + always_validate: bool, + persistent_iters: u32, +) -> Result<(), uc_afl_ret> { + afl_fuzz_custom( + uc, + input_file, + place_input_cb, + target::dummy_uc_validate_crash_callback, + target::dummy_uc_fuzz_callback, + exits, + always_validate, + persistent_iters, + ) +} + +/// Fuzzing entrypoint for FFI #[no_mangle] #[allow(non_camel_case_types)] pub extern "C" fn uc_afl_fuzz( - uc: uc_handle, + uc_handle: *mut uc_engine, input_file: *const c_char, place_input_callback: uc_afl_cb_place_input_t, exits: *const u64, @@ -74,31 +139,25 @@ pub extern "C" fn uc_afl_fuzz( persistent_iters: u32, data: *mut c_void, ) -> uc_afl_ret { - match child_fuzz( - uc, + uc_afl_fuzz_internal( + uc_handle, input_file, - persistent_iters, place_input_callback, - validate_crash_callback, - if exits.is_null() { - vec![] - } else { - unsafe { std::slice::from_raw_parts(exits, exit_count) }.to_vec() - }, + exits, + exit_count, None, + validate_crash_callback, always_validate, - true, + persistent_iters, data, - ) { - Ok(_) => uc_afl_ret::UC_AFL_RET_OK, - Err(e) => e, - } + ) } +/// Custom fuzzing entrypoint for FFI #[no_mangle] #[allow(non_camel_case_types)] pub extern "C" fn uc_afl_fuzz_custom( - uc: uc_handle, + uc_handle: *mut uc_engine, input_file: *const c_char, place_input_callback: uc_afl_cb_place_input_t, fuzz_callback: uc_afl_fuzz_cb_t, @@ -107,19 +166,147 @@ pub extern "C" fn uc_afl_fuzz_custom( persistent_iters: u32, data: *mut c_void, ) -> uc_afl_ret { - match child_fuzz( - uc, + uc_afl_fuzz_internal( + uc_handle, input_file, - persistent_iters, place_input_callback, - validate_crash_callback, - vec![], + std::ptr::null(), + 0, Some(fuzz_callback), + validate_crash_callback, always_validate, - true, + persistent_iters, data, - ) { + ) +} + +// In the implementation, there is a lot of manually created closures. +// This is due to the fact that two closure have different types even if +// their signature is the same. As a result, we must split the invocation +// to avoid checking the emptyness inside every round. +fn uc_afl_fuzz_internal( + uc_handle: *mut uc_engine, + input_file: *const c_char, + place_input_callback: uc_afl_cb_place_input_t, + exits: *const u64, + exit_count: usize, + fuzz_callback: Option, + validate_crash_callback: Option, + always_validate: bool, + persistent_iters: u32, + data: *mut c_void, +) -> uc_afl_ret { + let fuzz_data = UnicornFuzzData::new(data); + let uc = match unsafe { Unicorn::from_handle_with_data(uc_handle, fuzz_data) } { + Ok(uc) => uc, + Err(err) => { + return err.into(); + } + }; + + let place_input_cb = move |uc: &mut Unicorn<'_, UnicornFuzzData<*mut c_void>>, + input: &[u8], + persistent_round: u64| { + let handle = uc.get_handle(); + let data = uc.get_data_mut().user_data; + (place_input_callback)(handle, input.as_ptr(), input.len(), persistent_round, data) + }; + let validate_crash_cb = validate_crash_callback.map(|validate_crash_callback| { + move |uc: &mut Unicorn<'_, UnicornFuzzData<*mut c_void>>, + unicorn_result: Result<(), uc_error>, + input: &[u8], + persistent_round: u64| { + let handle = uc.get_handle(); + let data = uc.get_data_mut().user_data; + let unicorn_result = if let Err(err) = unicorn_result { + err + } else { + uc_error::OK + }; + (validate_crash_callback)( + handle, + unicorn_result, + input.as_ptr(), + input.len(), + persistent_round, + data, + ) + } + }); + let fuzz_cb = fuzz_callback.map(|fuzz_callback| { + move |uc: &mut Unicorn<'_, UnicornFuzzData<*mut c_void>>| { + let handle = uc.get_handle(); + let data = uc.get_data_mut().user_data; + let unicorn_result = fuzz_callback(handle, data); + if unicorn_result == uc_error::OK { + Ok(()) + } else { + Err(unicorn_result) + } + } + }); + + let input_file = if input_file.is_null() { + None + } else { + // legacy usage + let Ok(input_file_str) = unsafe { CStr::from_ptr(input_file) }.to_str() else { + return uc_afl_ret::UC_AFL_RET_FFI; + }; + Some(PathBuf::from(input_file_str)) + }; + + let exits = if exits.is_null() { + vec![] + } else { + unsafe { std::slice::from_raw_parts(exits, exit_count) }.to_vec() + }; + + let res = match (validate_crash_cb, fuzz_cb) { + (Some(validate_crash_cb), Some(fuzz_cb)) => afl_fuzz_custom( + uc, + input_file, + place_input_cb, + validate_crash_cb, + fuzz_cb, + exits, + always_validate, + persistent_iters, + ), + (Some(validate_crash_cb), None) => afl_fuzz_custom( + uc, + input_file, + place_input_cb, + validate_crash_cb, + target::dummy_uc_fuzz_callback, + exits, + always_validate, + persistent_iters, + ), + (None, Some(fuzz_cb)) => afl_fuzz_custom( + uc, + input_file, + place_input_cb, + target::dummy_uc_validate_crash_callback, + fuzz_cb, + exits, + always_validate, + persistent_iters, + ), + (None, None) => afl_fuzz_custom( + uc, + input_file, + place_input_cb, + target::dummy_uc_validate_crash_callback, + target::dummy_uc_fuzz_callback, + exits, + always_validate, + persistent_iters, + ), + }; + + match res { Ok(_) => uc_afl_ret::UC_AFL_RET_OK, - Err(e) => e, + Err(err) => err, } } diff --git a/src/target.rs b/src/target.rs index c12a737a..7d38d35d 100644 --- a/src/target.rs +++ b/src/target.rs @@ -1,8 +1,4 @@ -use std::{ - ffi::{c_char, CStr}, - os::raw::c_void, - path::PathBuf, -}; +use std::path::PathBuf; use libafl::{ corpus::{Corpus, InMemoryCorpus, Testcase}, @@ -20,110 +16,53 @@ use libafl_bolts::{ }; use libafl_targets::{EDGES_MAP_SIZE, SHM_FUZZING}; use log::{debug, trace, warn}; -use unicorn_engine::{ - ffi::{uc_ctl, uc_emu_start, uc_handle, uc_reg_read}, - uc_error, Arch, ControlType, Mode, RegisterARM, RegisterARM64, RegisterM68K, RegisterMIPS, - RegisterPPC, RegisterRISCV, RegisterS390X, RegisterSPARC, RegisterTRICORE, RegisterX86, -}; +use unicorn_engine::{uc_error, Arch, RegisterARM, Unicorn}; use crate::{ - executor::{UnicornAflExecutor, UnsafeSliceInput}, + executor::{UnicornAflExecutor, UnicornFuzzData, UnsafeSliceInput}, harness::LegacyHarnessStage, - uc_afl_cb_place_input_t, uc_afl_cb_validate_crash_t, uc_afl_fuzz_cb_t, uc_afl_ret, + uc_afl_ret, }; -fn get_afl_map_size() -> u32 { - std::env::var("AFL_MAP_SIZE") - .ok() - .map(|sz| u32::from_str_radix(&sz, 10).ok()) - .flatten() - .unwrap_or(1 << 16) // MAP_SIZE -} - -#[no_mangle] -extern "C" fn dummy_uc_fuzz_callback(uc: uc_handle, _data: *mut c_void) -> uc_error { - let mut arch = 0i32; - let ret = unsafe { - uc_ctl( - uc, - ControlType::UC_CTL_UC_ARCH as u32 | ControlType::UC_CTL_IO_READ as u32, - &mut arch, - ) - }; - if ret != uc_error::OK { - return ret; - } +/// Dummy fuzz callback if user don't specify their own callback +pub fn dummy_uc_fuzz_callback<'a, D: 'a>( + uc: &mut Unicorn<'a, UnicornFuzzData>, +) -> Result<(), uc_error> { + let arch = uc.get_arch(); - let mut mode = 0i32; - let ret = unsafe { - uc_ctl( - uc, - ControlType::UC_CTL_UC_MODE as u32 | ControlType::UC_CTL_IO_READ as u32, - &mut mode, - ) - }; - if ret != uc_error::OK { - return ret; - } - - let mut pc = 0u64; - let ret = if arch == Arch::X86 as i32 { - if mode == Mode::MODE_32.bits() { - unsafe { uc_reg_read(uc, RegisterX86::EIP as i32, &mut pc as *mut u64 as _) } - } else if mode == Mode::MODE_16.bits() { - unsafe { uc_reg_read(uc, RegisterX86::IP as i32, &mut pc as *mut u64 as _) } - } else { - unsafe { uc_reg_read(uc, RegisterX86::RIP as i32, &mut pc as *mut u64 as _) } - } - } else if arch == Arch::ARM as i32 { - let mut cpsr = 0u64; - let ret = unsafe { uc_reg_read(uc, RegisterARM::CPSR as i32, &mut cpsr as *mut u64 as _) }; - if ret != uc_error::OK { - return ret; - } - let ret = unsafe { uc_reg_read(uc, RegisterARM::PC as i32, &mut pc as *mut u64 as _) }; + let mut pc = uc.pc_read()?; + if arch == Arch::ARM { + let cpsr = uc.reg_read(RegisterARM::CPSR)?; if cpsr & 0x20 == 1 { pc |= 1; } - ret - } else if arch == Arch::RISCV as i32 { - unsafe { uc_reg_read(uc, RegisterRISCV::PC as i32, &mut pc as *mut u64 as _) } - } else if arch == Arch::MIPS as i32 { - unsafe { uc_reg_read(uc, RegisterMIPS::PC as i32, &mut pc as *mut u64 as _) } - } else if arch == Arch::PPC as i32 { - unsafe { uc_reg_read(uc, RegisterPPC::PC as i32, &mut pc as *mut u64 as _) } - } else if arch == Arch::SPARC as i32 { - unsafe { uc_reg_read(uc, RegisterSPARC::PC as i32, &mut pc as *mut u64 as _) } - } else if arch == Arch::M68K as i32 { - unsafe { uc_reg_read(uc, RegisterM68K::PC as i32, &mut pc as *mut u64 as _) } - } else if arch == Arch::S390X as i32 { - unsafe { uc_reg_read(uc, RegisterS390X::PC as i32, &mut pc as *mut u64 as _) } - } else if arch == Arch::ARM64 as i32 { - unsafe { uc_reg_read(uc, RegisterARM64::PC as i32, &mut pc as *mut u64 as _) } - } else if arch == Arch::TRICORE as i32 { - unsafe { uc_reg_read(uc, RegisterTRICORE::PC as i32, &mut pc as *mut u64 as _) } - } else { - uc_error::ARCH - }; - - if ret != uc_error::OK { - return ret; } - unsafe { uc_emu_start(uc, pc, 0, 0, 0) } + uc.emu_start(pc, 0, 0, 0) +} + +/// Dummy crash validation callback if user don't specify their own callback +pub fn dummy_uc_validate_crash_callback<'a, D: 'a>( + _uc: &mut Unicorn<'a, UnicornFuzzData>, + unicorn_result: Result<(), uc_error>, + _input: &[u8], + _persistent_round: u64, +) -> bool { + unicorn_result.is_err() } -pub fn child_fuzz( - uc: uc_handle, - input_file: *const c_char, +/// Internal entrypoint for fuzzing +pub fn child_fuzz<'a, D: 'a>( + uc: Unicorn<'a, UnicornFuzzData>, + input_file: Option, iters: u32, - place_input_cb: uc_afl_cb_place_input_t, - validate_crash_cb: Option, + place_input_cb: impl FnMut(&mut Unicorn<'a, UnicornFuzzData>, &[u8], u64) -> bool + 'a, + validate_crash_cb: impl FnMut(&mut Unicorn<'a, UnicornFuzzData>, Result<(), uc_error>, &[u8], u64) -> bool + + 'a, + fuzz_callback: impl FnMut(&mut Unicorn<'a, UnicornFuzzData>) -> Result<(), uc_error> + 'a, exits: Vec, - fuzz_callback: Option, always_validate: bool, run_once_if_no_afl_present: bool, - data: *mut c_void, ) -> Result<(), uc_afl_ret> { // Enable logging #[cfg(feature = "env_logger")] @@ -132,14 +71,14 @@ pub fn child_fuzz( let has_afl = libafl_targets::map_input_shared_memory() && libafl_targets::map_shared_memory(); trace!("AFL detected: {}", has_afl); - if !input_file.is_null() && has_afl { + if !input_file.is_none() && has_afl { warn!("Shared memory fuzzing is enabled and the input file is ignored!"); } - if input_file.is_null() && !has_afl { + if input_file.is_none() && !has_afl { warn!("No input file is provided. We will run harness with zero inputs."); } if has_afl || run_once_if_no_afl_present { - let map_size = get_afl_map_size(); + let map_size = uc.get_data().map_size(); unsafe { EDGES_MAP_SIZE = map_size as usize; SHM_FUZZING = 1; @@ -152,11 +91,9 @@ pub fn child_fuzz( uc, place_input_cb, validate_crash_cb, - fuzz_callback.unwrap_or(dummy_uc_fuzz_callback), + fuzz_callback, always_validate, exits, - map_size as u32, - data, )?; let mut fb = BoolValueFeedback::new(&Handle::new("dumb_ob".into())); @@ -178,20 +115,7 @@ pub fn child_fuzz( })); let sched = QueueScheduler::new(); let iters = if run_once_if_no_afl_present { 1 } else { iters }; - let input_file = if has_afl { - None - } else { - if input_file.is_null() { - None - } else { - // legact usage - Some(PathBuf::from( - unsafe { CStr::from_ptr(input_file) } - .to_str() - .map_err(|_| uc_afl_ret::UC_AFL_RET_FFI)?, - )) - } - }; + let input_file = if has_afl { None } else { input_file }; let stage = LegacyHarnessStage::new(iters as usize, map_size, input_file); let mut stages = tuple_list!(stage); let mut fuzzer = StdFuzzer::new(sched, fb, sol); From 1769b6c60b6ab7771f77451dae5099ea55020e33 Mon Sep 17 00:00:00 2001 From: mio Date: Tue, 22 Apr 2025 14:59:16 +0800 Subject: [PATCH 2/2] Add TODO --- Readme.md | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/Readme.md b/Readme.md index 4472d360..1ae9a931 100644 --- a/Readme.md +++ b/Readme.md @@ -3,4 +3,10 @@ The project builds a bridge between AFL++ and unicorn engine. You can fuzz unicorn targets using python, rust, and C. -Check out [the examples](https://github.com/AFLplusplus/AFLplusplus/tree/stable/unicorn_mode/samples) in AFLplusplus/unicorn_mode \ No newline at end of file +Check out [the examples](https://github.com/AFLplusplus/AFLplusplus/tree/stable/unicorn_mode/samples) in AFLplusplus/unicorn_mode + +TODO: + +- [] Python/C bindings +- [] cmplog support +- [] libafl upstream \ No newline at end of file