WIP: Implement ASAN instrumentation using frida (#45)

* frida_asan: Implemented initial asan runtime library

* frida_asan: Switch to hashbrown

* Implemented GOT-based hooking to isolate the hooking of the memory functions. Implemented initial ASAN instrumentation

* WIP: Shadowing all used memory. Currently tracking pages using a BTreeSet. Slow AF!

* Add SigTrap to unix_signals and inprocess

* Working frida-asan, almost no speed degradation.

Currently the shadow check is reversed, so it checks only that the shadow is not 0.
We need to implement sub-8-byte checking.

* Format

* Cleanup and formatting

* Sub-qword and 16-byte checks implemented; Fixed unaligned access to QWORD

* Pass the ucontext_t to signal handlers. Initial regdump on crash

* Fix typo

* Make the context argument a mut ref

* Add missing files; Implement initial reporting

* Refactor out gothook; Move safety checkers to dynasm

* Get rid of const assembly blobs no longer needed

* Move to a handler function instead of using SIGTRAP.

This bloats the transformed code, but doesn't seem to have a major impact on performance.

Also, implemented pretty backtraces and assembly output.

* Formatting

* Get rid of all the pinning crap I wasted my day on, We don't need it

* windows fixes

* ashmem

* ashmem_service: server side ready

* ashmem_service: client side ready. Ready for integration

* ashmem_service: changes to UnixShMem to make it 'threadable'

* ashmem_service: format

* ashmem_service: Undo changes to UnixShMem, make the thread own the AshmemService instead; Fix protocol bug

* ashmem_service: working ashmem service. Fix merge issues

* use the newly released capston e 0.8.0; Fix a nasty bug where the afl_area an pc_pointer were reversed. Changed Vectors to Boxed [u8]

* Implement type detection for reporting; Implement double-free/unallocated free checking

* fmt

* Cleanup code a little

* frida-asan: This is an omnibus commit. Should probably have been a bunch of small commits, but I don't have the time/patience.

 - Implemented DrCov support in order to debug a failing harness. This is actually
   generic and should be moved out of libafl_frida.
 - Implemented LIBAFL_FRIDA_OPTIONS env var to pass options to the frida helper,
   to dynamically enable/disable asan and drcov.
 - Implemented memory reuse - after each test case the used pages are recycled and
   can be reused in the next test case.
 - Implemented and tested vectorized instruction instrumentation.
 - Implemented not instrumenting atomic load/store instructions. The cost of
   trying to emulate their behaviour is too high at the moment.
 - Implemented probing of shadow bit to determine the best match for the current
   system.
 - Implemented shadow memory pre-mapping where it is available. We probe for this
   too.
 - Implemented ability to specify a list of modules to instrument on the command
   line. This allows fine-grained control of which modules are instrumented for
   coverage/asan/drcov.
 - Implemented unpoisoning of the Input target_bytes in a pre_exec hook.
 - Added support for zero-sized allocations. We return 0x10 bytes at the moment.
 - Added all known operator new/delete functions to hooks.
 - Added workaround for frida_gum_allocate_near bug.
 - Cleaned up reporting, added reporting for different error types.

* frida-asan: Implement leak detection

* Fix merge issues

* Rebased on dev to get llmp/shmem changes; Clippy fixes

* Add FridaOptions struct

* Add the Custom ExitKind; Get rid of Clone/PartialEq on ExitKind

* Make it possible to recover from an ASAN error

* Add SIGTRAP to crashing signals

* Add back (conditional) crashing on Asan errors.

* Fix too-large immediates in add instruction

* Implement RcShMemProvider, finally fix the EOP bug

* Clear ASAN_ERRORS before each test

* Fix warnings; Fix review issues

* Cleanup prints

* Add timeout to Frida mode

* Make allocation-/free-site backtraces optional

* CPU Context and backtrace (on android/aarch64 atm) on crash

* Make stalker conditional

* Add metadata to solution, and write metadata files

* Add addresses to backtrace; Add reporting of ASAN stack errors; Fix ASAN reporting bugs

* Remove meaningless backtrace on crash

* Fix the x0, x1 load in report

* use upstream color-backtrace

* use __builtin_thread_pointer instead of custom asm

* Don't unwrap ASAN_ERRORS if it isn't some

* Fix bug where we weren't clearing the drcov basicblocks after each run

* Fix bug where we were dropping an ashmem too soon

* Fix OwnedPtr instead of CPtr

* Fix gettls for all archs

* cfg guards for target arch, disabling Frida-ASAN/-DrCov if not on aarch64

* Cargo fmt

* Only panic in options when asan/drcov are turned on; Merge fixes

* gothook only supported on unix

* Fix gettls on msvc

* Another attempt to fix MSVC gettls

* Fix backtrace use

* nostd fixes; warning fixes

* formatting

* Migrate FridaEdgeCoverageHelper into libafl_frida, and rename to FridaInstrumentationHelper

* Clean up uses

* Move DrCovWriter to libafl_targets

* Refactor DrCovWriter to get a vec of DrCovBasicBlocks; formatting

* Update to newer backtrace which supports android with gimli

* windows fixes

Co-authored-by: Dominik Maier <domenukk@gmail.com>
Co-authored-by: andreafioraldi <andreafioraldi@gmail.com>
This commit is contained in:
s1341 2021-04-28 11:12:49 +03:00 committed by GitHub
parent b8b01baf59
commit 5c856cccc8
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
29 changed files with 3496 additions and 579 deletions

View File

@ -10,6 +10,7 @@ members = [
"libafl_derive", "libafl_derive",
"libafl_cc", "libafl_cc",
"libafl_targets", "libafl_targets",
"libafl_frida",
] ]
exclude = [ exclude = [
"fuzzers/libfuzzer_libpng", "fuzzers/libfuzzer_libpng",

View File

@ -22,9 +22,14 @@ num_cpus = "1.0"
which = "4.1" which = "4.1"
[target.'cfg(unix)'.dependencies] [target.'cfg(unix)'.dependencies]
libafl = { path = "../../libafl/" } libafl = { path = "../../libafl/", features = [ "std" ] } #, "llmp_small_maps", "llmp_debug"]}
frida-gum = { version = "0.3.2", optional = true, features = ["auto-download", "event-sink", "invocation-listener"] } capstone = "0.8.0"
frida-gum-sys = { version = "0.2.2", optional = true, features = ["auto-download", "event-sink", "invocation-listener"] } frida-gum = { version = "0.4", optional = true, features = [ "auto-download", "event-sink", "invocation-listener"] }
frida-gum-sys = { version = "0.2.4", optional = true, features = [ "auto-download", "event-sink", "invocation-listener"] }
libafl_frida = { path = "../../libafl_frida", version = "0.1.0" }
lazy_static = "1.4.0" lazy_static = "1.4.0"
libc = "0.2" libc = "0.2"
libloading = "0.7.0" libloading = "0.7.0"
num-traits = "0.2.14"
rangemap = "0.1.10"
seahash = "4.1.0"

View File

@ -119,6 +119,8 @@ fn main() {
//.arg("HAS_DUMMY_CRASH=1") //.arg("HAS_DUMMY_CRASH=1")
.arg("-fPIC") .arg("-fPIC")
.arg("-shared") .arg("-shared")
.arg("-O3")
//.arg("-fomit-frame-pointer")
.arg(if env::var("CARGO_CFG_TARGET_OS").unwrap() == "android" { .arg(if env::var("CARGO_CFG_TARGET_OS").unwrap() == "android" {
"-static-libstdc++" "-static-libstdc++"
} else { } else {

View File

@ -17,6 +17,7 @@
#include <stddef.h> #include <stddef.h>
#include <stdint.h> #include <stdint.h>
#include <string.h> #include <string.h>
#include <stdlib.h>
#include <vector> #include <vector>
@ -83,6 +84,25 @@ extern "C" int afl_libfuzzer_init() {
return 0; return 0;
} }
static char * allocation = NULL;
__attribute__((noinline))
void func3( char * alloc) {
printf("func3\n");
if (random() % 5 == 0) {
alloc[0xff] = 0xde;
}
}
__attribute__((noinline))
void func2() {
allocation = (char*)malloc(0xff);
printf("func2\n");
func3(allocation);
}
__attribute__((noinline))
void func1() {
printf("func1\n");
func2();
}
// Entry point for LibFuzzer. // Entry point for LibFuzzer.
// Roughly follows the libpng book example: // Roughly follows the libpng book example:
// http://www.libpng.org/pub/png/book/chapter13.html // http://www.libpng.org/pub/png/book/chapter13.html
@ -91,6 +111,9 @@ extern "C" int LLVMFuzzerTestOneInput(const uint8_t* data, size_t size) {
return 0; return 0;
} }
func1();
std::vector<unsigned char> v(data, data + size); std::vector<unsigned char> v(data, data + size);
if (png_sig_cmp(v.data(), 0, kPngHeaderSize)) { if (png_sig_cmp(v.data(), 0, kPngHeaderSize)) {
// not a PNG. // not a PNG.

View File

@ -4,12 +4,14 @@
use libafl::{ use libafl::{
bolts::tuples::{tuple_list, Named}, bolts::tuples::{tuple_list, Named},
corpus::{ corpus::{
Corpus, InMemoryCorpus, IndexesLenTimeMinimizerCorpusScheduler, OnDiskCorpus, ondisk::OnDiskMetadataFormat, Corpus, InMemoryCorpus,
QueueCorpusScheduler, IndexesLenTimeMinimizerCorpusScheduler, OnDiskCorpus, QueueCorpusScheduler,
}, },
events::{setup_restarting_mgr_std, EventManager}, events::{setup_restarting_mgr_std, EventManager},
executors::{inprocess::InProcessExecutor, Executor, ExitKind, HasObservers}, executors::{
feedbacks::{CrashFeedback, MaxMapFeedback}, inprocess::InProcessExecutor, timeout::TimeoutExecutor, Executor, ExitKind, HasObservers,
},
feedbacks::{CrashFeedback, MaxMapFeedback, TimeoutFeedback},
fuzzer::{Fuzzer, StdFuzzer}, fuzzer::{Fuzzer, StdFuzzer},
inputs::{HasTargetBytes, Input}, inputs::{HasTargetBytes, Input},
mutators::scheduled::{havoc_mutations, StdScheduledMutator}, mutators::scheduled::{havoc_mutations, StdScheduledMutator},
@ -22,229 +24,38 @@ use libafl::{
Error, Error,
}; };
use core::cell::RefCell;
#[cfg(target_arch = "x86_64")]
use frida_gum::instruction_writer::X86Register;
#[cfg(target_arch = "aarch64")]
use frida_gum::instruction_writer::{Aarch64Register, IndexMode};
use frida_gum::{ use frida_gum::{
instruction_writer::InstructionWriter, stalker::{NoneEventSink, Stalker},
stalker::{NoneEventSink, Stalker, Transformer}, Gum, NativePointer,
}; };
use frida_gum::{Gum, MemoryRange, Module, NativePointer, PageProtection};
use std::{env, ffi::c_void, path::PathBuf};
/// An helper that feeds FridaInProcessExecutor with user-supplied instrumentation use std::{env, ffi::c_void, marker::PhantomData, path::PathBuf, time::Duration};
pub trait FridaHelper<'a> {
fn transformer(&self) -> &Transformer<'a>;
}
const MAP_SIZE: usize = 64 * 1024; use libafl_frida::{
asan_rt::{AsanErrorsFeedback, AsanErrorsObserver, ASAN_ERRORS},
helper::{FridaHelper, FridaInstrumentationHelper, MAP_SIZE},
FridaOptions,
};
/// An helper that feeds FridaInProcessExecutor with edge-coverage instrumentation struct FridaInProcessExecutor<'a, 'b, 'c, FH, H, I, OT>
struct FridaEdgeCoverageHelper<'a> {
map: [u8; MAP_SIZE],
previous_pc: RefCell<u64>,
base_address: u64,
size: usize,
current_log_impl: u64,
/// Transformer that has to be passed to FridaInProcessExecutor
transformer: Option<Transformer<'a>>,
}
impl<'a> FridaHelper<'a> for FridaEdgeCoverageHelper<'a> {
fn transformer(&self) -> &Transformer<'a> {
self.transformer.as_ref().unwrap()
}
}
/// Helper function to get the size of a module's CODE section from frida
pub fn get_module_size(module_name: &str) -> usize {
let mut code_size = 0;
let code_size_ref = &mut code_size;
Module::enumerate_ranges(module_name, PageProtection::ReadExecute, move |details| {
*code_size_ref = details.memory_range().size() as usize;
true
});
code_size
}
/// A minimal maybe_log implementation. We insert this into the transformed instruction stream
/// every time we need a copy that is within a direct branch of the start of the transformed basic
/// block.
#[cfg(target_arch = "x86_64")]
const MAYBE_LOG_CODE: [u8; 47] = [
0x9c, /* pushfq */
0x50, /* push rax */
0x51, /* push rcx */
0x52, /* push rdx */
0x48, 0x8d, 0x05, 0x24, 0x00, 0x00, 0x00, /* lea rax, sym._afl_area_ptr_ptr */
0x48, 0x8b, 0x00, /* mov rax, qword [rax] */
0x48, 0x8d, 0x0d, 0x22, 0x00, 0x00, 0x00, /* lea rcx, sym.previous_pc */
0x48, 0x8b, 0x11, /* mov rdx, qword [rcx] */
0x48, 0x8b, 0x12, /* mov rdx, qword [rdx] */
0x48, 0x31, 0xfa, /* xor rdx, rdi */
0xfe, 0x04, 0x10, /* inc byte [rax + rdx] */
0x48, 0xd1, 0xef, /* shr rdi, 1 */
0x48, 0x8b, 0x01, /* mov rax, qword [rcx] */
0x48, 0x89, 0x38, /* mov qword [rax], rdi */
0x5a, /* pop rdx */
0x59, /* pop rcx */
0x58, /* pop rax */
0x9d, /* popfq */
0xc3, /* ret */
/* Read-only data goes here: */
/* uint8_t* afl_area_ptr */
/* uint64_t* afl_prev_loc_ptr */
];
#[cfg(target_arch = "aarch64")]
const MAYBE_LOG_CODE: [u8; 56] = [
// __afl_area_ptr[current_pc ^ previous_pc]++;
// previous_pc = current_pc >> 1;
0xE1, 0x0B, 0xBF, 0xA9, // stp x1, x2, [sp, -0x10]!
0xE3, 0x13, 0xBF, 0xA9, // stp x3, x4, [sp, -0x10]!
// x0 = current_pc
0x81, 0x01, 0x00, 0x58, // ldr x1, #0x30, =__afl_area_ptr
0xa2, 0x01, 0x00, 0x58, // ldr x2, #0x38, =&previous_pc
0x44, 0x00, 0x40, 0xf9, // ldr x4, [x2] (=previous_pc)
// __afl_area_ptr[current_pc ^ previous_pc]++;
0x84, 0x00, 0x00, 0xca, // eor x4, x4, x0
0x23, 0x68, 0x64, 0xf8, // ldr x3, [x1, x4]
0x63, 0x04, 0x00, 0x91, // add x3, x3, #1
0x23, 0x68, 0x24, 0xf8, // str x3, [x1, x2]
// previous_pc = current_pc >> 1;
0xe0, 0x07, 0x40, 0x8b, // add x0, xzr, x0, LSR #1
0x40, 0x00, 0x00, 0xf9, // str x0, [x2]
0xE3, 0x13, 0xc1, 0xA8, // ldp x3, x4, [sp], #0x10
0xE1, 0x0B, 0xc1, 0xA8, // ldp x1, x2, [sp], #0x10
0xC0, 0x03, 0x5F, 0xD6, // ret
// &afl_area_ptr
// &afl_prev_loc_ptr
];
/// The implementation of the FridaEdgeCoverageHelper
impl<'a> FridaEdgeCoverageHelper<'a> {
/// Constructor function to create a new FridaEdgeCoverageHelper, given a module_name.
pub fn new(gum: &'a Gum, module_name: &str) -> Self {
let mut helper = Self {
map: [0u8; MAP_SIZE],
previous_pc: RefCell::new(0x0),
base_address: Module::find_base_address(module_name).0 as u64,
size: get_module_size(module_name),
current_log_impl: 0,
transformer: None,
};
let transformer = Transformer::from_callback(gum, |basic_block, _output| {
let mut first = true;
for instruction in basic_block {
if first {
first = false;
let address = unsafe { (*instruction.instr()).address };
if address >= helper.base_address
&& address <= helper.base_address + helper.size as u64
{
let writer = _output.writer();
if helper.current_log_impl == 0
|| !writer.can_branch_directly_to(helper.current_log_impl)
|| !writer.can_branch_directly_between(
writer.pc() + 128,
helper.current_log_impl,
)
{
let after_log_impl = writer.code_offset() + 1;
#[cfg(target_arch = "x86_64")]
writer.put_jmp_near_label(after_log_impl);
#[cfg(target_arch = "aarch64")]
writer.put_b_label(after_log_impl);
helper.current_log_impl = writer.pc();
writer.put_bytes(&MAYBE_LOG_CODE);
let prev_loc_pointer = helper.previous_pc.as_ptr() as *mut _ as usize;
let map_pointer = helper.map.as_ptr() as usize;
writer.put_bytes(&prev_loc_pointer.to_ne_bytes());
writer.put_bytes(&map_pointer.to_ne_bytes());
writer.put_label(after_log_impl);
}
#[cfg(target_arch = "x86_64")]
{
println!("here");
writer.put_lea_reg_reg_offset(
X86Register::Rsp,
X86Register::Rsp,
-(frida_gum_sys::GUM_RED_ZONE_SIZE as i32),
);
writer.put_push_reg(X86Register::Rdi);
writer.put_mov_reg_address(
X86Register::Rdi,
((address >> 4) ^ (address << 8)) & (MAP_SIZE - 1) as u64,
);
writer.put_call_address(helper.current_log_impl);
writer.put_pop_reg(X86Register::Rdi);
writer.put_lea_reg_reg_offset(
X86Register::Rsp,
X86Register::Rsp,
frida_gum_sys::GUM_RED_ZONE_SIZE as i32,
);
}
#[cfg(target_arch = "aarch64")]
{
writer.put_stp_reg_reg_reg_offset(
Aarch64Register::Lr,
Aarch64Register::X0,
Aarch64Register::Sp,
-(16 + frida_gum_sys::GUM_RED_ZONE_SIZE as i32) as i64,
IndexMode::PreAdjust,
);
writer.put_ldr_reg_u64(
Aarch64Register::X0,
((address >> 4) ^ (address << 8)) & (MAP_SIZE - 1) as u64,
);
writer.put_bl_imm(helper.current_log_impl);
writer.put_ldp_reg_reg_reg_offset(
Aarch64Register::Lr,
Aarch64Register::X0,
Aarch64Register::Sp,
16 + frida_gum_sys::GUM_RED_ZONE_SIZE as i64,
IndexMode::PostAdjust,
);
}
}
}
instruction.keep()
}
});
helper.transformer = Some(transformer);
helper
}
}
struct FridaInProcessExecutor<'a, FH, H, I, OT>
where where
FH: FridaHelper<'a>, FH: FridaHelper<'b>,
H: FnMut(&[u8]) -> ExitKind, H: FnMut(&[u8]) -> ExitKind,
I: Input + HasTargetBytes, I: Input + HasTargetBytes,
OT: ObserversTuple, OT: ObserversTuple,
{ {
base: InProcessExecutor<'a, H, I, OT>, base: TimeoutExecutor<InProcessExecutor<'a, H, I, OT>, I, OT>,
/// Frida's dynamic rewriting engine /// Frida's dynamic rewriting engine
stalker: Stalker<'a>, stalker: Stalker<'a>,
/// User provided callback for instrumentation /// User provided callback for instrumentation
helper: &'a FH, helper: &'c mut FH,
followed: bool, followed: bool,
_phantom: PhantomData<&'b u8>,
} }
impl<'a, FH, H, I, OT> Executor<I> for FridaInProcessExecutor<'a, FH, H, I, OT> impl<'a, 'b, 'c, FH, H, I, OT> Executor<I> for FridaInProcessExecutor<'a, 'b, 'c, FH, H, I, OT>
where where
FH: FridaHelper<'a>, FH: FridaHelper<'b>,
H: FnMut(&[u8]) -> ExitKind, H: FnMut(&[u8]) -> ExitKind,
I: Input + HasTargetBytes, I: Input + HasTargetBytes,
OT: ObserversTuple, OT: ObserversTuple,
@ -255,22 +66,34 @@ where
where where
EM: EventManager<I, S>, EM: EventManager<I, S>,
{ {
if !self.followed { if self.helper.stalker_enabled() {
self.followed = true; if !self.followed {
self.stalker self.followed = true;
.follow_me::<NoneEventSink>(self.helper.transformer(), None); self.stalker
} else { .follow_me::<NoneEventSink>(self.helper.transformer(), None);
self.stalker.activate(NativePointer( } else {
self.base.harness_mut() as *mut _ as *mut c_void self.stalker.activate(NativePointer(
)) self.base.inner().harness_mut() as *mut _ as *mut c_void
))
}
} }
self.helper.pre_exec(input);
self.base.pre_exec(state, event_mgr, input) self.base.pre_exec(state, event_mgr, input)
} }
/// Instruct the target about the input and run /// Instruct the target about the input and run
#[inline] #[inline]
fn run_target(&mut self, input: &I) -> Result<ExitKind, Error> { fn run_target(&mut self, input: &I) -> Result<ExitKind, Error> {
self.base.run_target(input) let res = self.base.run_target(input);
if unsafe { ASAN_ERRORS.is_some() && !ASAN_ERRORS.as_ref().unwrap().is_empty() } {
println!("Crashing target as it had ASAN errors");
unsafe {
libc::raise(libc::SIGABRT);
}
}
res
} }
/// Called right after execution finished. /// Called right after execution finished.
@ -284,14 +107,17 @@ where
where where
EM: EventManager<I, S>, EM: EventManager<I, S>,
{ {
self.stalker.deactivate(); if self.helper.stalker_enabled() {
self.stalker.deactivate();
}
self.helper.post_exec(input);
self.base.post_exec(state, event_mgr, input) self.base.post_exec(state, event_mgr, input)
} }
} }
impl<'a, FH, H, I, OT> HasObservers<OT> for FridaInProcessExecutor<'a, FH, H, I, OT> impl<'a, 'b, 'c, FH, H, I, OT> HasObservers<OT> for FridaInProcessExecutor<'a, 'b, 'c, FH, H, I, OT>
where where
FH: FridaHelper<'a>, FH: FridaHelper<'b>,
H: FnMut(&[u8]) -> ExitKind, H: FnMut(&[u8]) -> ExitKind,
I: Input + HasTargetBytes, I: Input + HasTargetBytes,
OT: ObserversTuple, OT: ObserversTuple,
@ -307,9 +133,9 @@ where
} }
} }
impl<'a, FH, H, I, OT> Named for FridaInProcessExecutor<'a, FH, H, I, OT> impl<'a, 'b, 'c, FH, H, I, OT> Named for FridaInProcessExecutor<'a, 'b, 'c, FH, H, I, OT>
where where
FH: FridaHelper<'a>, FH: FridaHelper<'b>,
H: FnMut(&[u8]) -> ExitKind, H: FnMut(&[u8]) -> ExitKind,
I: Input + HasTargetBytes, I: Input + HasTargetBytes,
OT: ObserversTuple, OT: ObserversTuple,
@ -319,31 +145,37 @@ where
} }
} }
impl<'a, FH, H, I, OT> FridaInProcessExecutor<'a, FH, H, I, OT> impl<'a, 'b, 'c, FH, H, I, OT> FridaInProcessExecutor<'a, 'b, 'c, FH, H, I, OT>
where where
FH: FridaHelper<'a>, FH: FridaHelper<'b>,
H: FnMut(&[u8]) -> ExitKind, H: FnMut(&[u8]) -> ExitKind,
I: Input + HasTargetBytes, I: Input + HasTargetBytes,
OT: ObserversTuple, OT: ObserversTuple,
{ {
pub fn new(gum: &'a Gum, base: InProcessExecutor<'a, H, I, OT>, helper: &'a FH) -> Self { pub fn new(
let mut stalker = Stalker::new(gum); gum: &'a Gum,
base: InProcessExecutor<'a, H, I, OT>,
helper: &'c mut FH,
timeout: Duration,
) -> Self {
let stalker = Stalker::new(gum);
// Let's exclude the main module and libc.so at least: // Let's exclude the main module and libc.so at least:
stalker.exclude(&MemoryRange::new( //stalker.exclude(&MemoryRange::new(
Module::find_base_address(&env::args().next().unwrap()), //Module::find_base_address(&env::args().next().unwrap()),
get_module_size(&env::args().next().unwrap()), //get_module_size(&env::args().next().unwrap()),
)); //));
stalker.exclude(&MemoryRange::new( //stalker.exclude(&MemoryRange::new(
Module::find_base_address("libc.so"), //Module::find_base_address("libc.so"),
get_module_size("libc.so"), //get_module_size("libc.so"),
)); //));
Self { Self {
base, base: TimeoutExecutor::new(base, timeout),
stalker, stalker,
helper, helper,
followed: false, followed: false,
_phantom: PhantomData,
} }
} }
} }
@ -362,6 +194,11 @@ pub fn main() {
fuzz( fuzz(
&env::args().nth(1).expect("no module specified"), &env::args().nth(1).expect("no module specified"),
&env::args().nth(2).expect("no symbol specified"), &env::args().nth(2).expect("no symbol specified"),
env::args()
.nth(3)
.expect("no modules to instrument specified")
.split(":")
.collect(),
vec![PathBuf::from("./corpus")], vec![PathBuf::from("./corpus")],
PathBuf::from("./crashes"), PathBuf::from("./crashes"),
1337, 1337,
@ -387,6 +224,7 @@ fn fuzz(
unsafe fn fuzz( unsafe fn fuzz(
module_name: &str, module_name: &str,
symbol_name: &str, symbol_name: &str,
modules_to_instrument: Vec<&str>,
corpus_dirs: Vec<PathBuf>, corpus_dirs: Vec<PathBuf>,
objective_dir: PathBuf, objective_dir: PathBuf,
broker_port: u16, broker_port: u16,
@ -395,37 +233,43 @@ unsafe fn fuzz(
let stats = SimpleStats::new(|s| println!("{}", s)); let stats = SimpleStats::new(|s| println!("{}", s));
// The restarting state will spawn the same process again as child, then restarted it each time it crashes. // The restarting state will spawn the same process again as child, then restarted it each time it crashes.
let (state, mut restarting_mgr) = let (state, mut restarting_mgr) = match setup_restarting_mgr_std(stats, broker_port) {
match setup_restarting_mgr_std(stats, broker_port) { Ok(res) => res,
Ok(res) => res, Err(err) => match err {
Err(err) => match err { Error::ShuttingDown => {
Error::ShuttingDown => { return Ok(());
return Ok(()); }
} _ => {
_ => { panic!("Failed to setup the restarter: {}", err);
panic!("Failed to setup the restarter: {}", err); }
} },
}, };
};
let gum = Gum::obtain(); let gum = Gum::obtain();
let lib = libloading::Library::new(module_name).unwrap(); let lib = libloading::Library::new(module_name).unwrap();
let target_func: libloading::Symbol<unsafe extern "C" fn(data: *const u8, size: usize) -> i32> = let target_func: libloading::Symbol<unsafe extern "C" fn(data: *const u8, size: usize) -> i32> =
lib.get(symbol_name.as_bytes()).unwrap(); lib.get(symbol_name.as_bytes()).unwrap();
let mut frida_helper = FridaEdgeCoverageHelper::new(&gum, module_name);
// Create an observation channel using the coverage map
let edges_observer = HitcountsMapObserver::new(StdMapObserver::new_from_ptr(
"edges",
frida_helper.map.as_mut_ptr(),
MAP_SIZE,
));
let mut frida_harness = move |buf: &[u8]| { let mut frida_harness = move |buf: &[u8]| {
(target_func)(buf.as_ptr(), buf.len()); (target_func)(buf.as_ptr(), buf.len());
ExitKind::Ok ExitKind::Ok
}; };
let mut frida_helper = FridaInstrumentationHelper::new(
&gum,
FridaOptions::parse_env_options(),
module_name,
&modules_to_instrument,
);
// Create an observation channel using the coverage map
let edges_observer = HitcountsMapObserver::new(StdMapObserver::new_from_ptr(
"edges",
frida_helper.map_ptr(),
MAP_SIZE,
));
// If not restarting, create a State from scratch // If not restarting, create a State from scratch
let mut state = state.unwrap_or_else(|| { let mut state = state.unwrap_or_else(|| {
State::new( State::new(
@ -441,9 +285,14 @@ unsafe fn fuzz(
)), )),
// Corpus in which we store solutions (crashes in this example), // Corpus in which we store solutions (crashes in this example),
// on disk so the user can get them after stopping the fuzzer // on disk so the user can get them after stopping the fuzzer
OnDiskCorpus::new(objective_dir).unwrap(), OnDiskCorpus::new_save_meta(objective_dir, Some(OnDiskMetadataFormat::JsonPretty))
.unwrap(),
// Feedbacks to recognize an input as solution // Feedbacks to recognize an input as solution
tuple_list!(CrashFeedback::new()), tuple_list!(
CrashFeedback::new(),
TimeoutFeedback::new(),
AsanErrorsFeedback::new()
),
) )
}); });
@ -474,21 +323,22 @@ unsafe fn fuzz(
InProcessExecutor::new( InProcessExecutor::new(
"in-process(edges)", "in-process(edges)",
&mut frida_harness, &mut frida_harness,
tuple_list!(edges_observer), tuple_list!(edges_observer, AsanErrorsObserver::new(&ASAN_ERRORS)),
&mut state, &mut state,
&mut restarting_mgr, &mut restarting_mgr,
)?, )?,
&frida_helper, &mut frida_helper,
Duration::new(10, 0),
); );
// Let's exclude the main module and libc.so at least: // Let's exclude the main module and libc.so at least:
executor.stalker.exclude(&MemoryRange::new( //executor.stalker.exclude(&MemoryRange::new(
Module::find_base_address(&env::args().next().unwrap()), //Module::find_base_address(&env::args().next().unwrap()),
get_module_size(&env::args().next().unwrap()), //get_module_size(&env::args().next().unwrap()),
)); //));
executor.stalker.exclude(&MemoryRange::new( //executor.stalker.exclude(&MemoryRange::new(
Module::find_base_address("libc.so"), //Module::find_base_address("libc.so"),
get_module_size("libc.so"), //get_module_size("libc.so"),
)); //));
// In case the corpus is empty (on first run), reset // In case the corpus is empty (on first run), reset
if state.corpus().count() < 1 { if state.corpus().count() < 1 {
@ -501,6 +351,7 @@ unsafe fn fuzz(
println!("We imported {} inputs from disk.", state.corpus().count()); println!("We imported {} inputs from disk.", state.corpus().count());
} }
//executor.helper.register_thread();
fuzzer.fuzz_loop(&mut state, &mut executor, &mut restarting_mgr, &scheduler)?; fuzzer.fuzz_loop(&mut state, &mut executor, &mut restarting_mgr, &scheduler)?;
// Never reached // Never reached

View File

@ -4,12 +4,12 @@
use std::{env, path::PathBuf}; use std::{env, path::PathBuf};
use libafl::{ use libafl::{
bolts::{shmem::StdShMem, tuples::tuple_list}, bolts::tuples::tuple_list,
corpus::{ corpus::{
Corpus, InMemoryCorpus, IndexesLenTimeMinimizerCorpusScheduler, OnDiskCorpus, Corpus, InMemoryCorpus, IndexesLenTimeMinimizerCorpusScheduler, OnDiskCorpus,
QueueCorpusScheduler, QueueCorpusScheduler,
}, },
events::setup_restarting_mgr, events::setup_restarting_mgr_std,
executors::{inprocess::InProcessExecutor, ExitKind}, executors::{inprocess::InProcessExecutor, ExitKind},
feedbacks::{CrashFeedback, MaxMapFeedback, TimeFeedback}, feedbacks::{CrashFeedback, MaxMapFeedback, TimeFeedback},
fuzzer::{Fuzzer, StdFuzzer}, fuzzer::{Fuzzer, StdFuzzer},
@ -25,7 +25,6 @@ use libafl::{
use libafl_targets::{libfuzzer_initialize, libfuzzer_test_one_input, EDGES_MAP, MAX_EDGES_NUM}; use libafl_targets::{libfuzzer_initialize, libfuzzer_test_one_input, EDGES_MAP, MAX_EDGES_NUM};
/// The main fn, no_mangle as it is a C main
pub fn main() { pub fn main() {
// Registry the metadata types used in this fuzzer // Registry the metadata types used in this fuzzer
// Needed only on no_std // Needed only on no_std
@ -49,18 +48,17 @@ fn fuzz(corpus_dirs: Vec<PathBuf>, objective_dir: PathBuf, broker_port: u16) ->
let stats = SimpleStats::new(|s| println!("{}", s)); let stats = SimpleStats::new(|s| println!("{}", s));
// The restarting state will spawn the same process again as child, then restarted it each time it crashes. // The restarting state will spawn the same process again as child, then restarted it each time it crashes.
let (state, mut restarting_mgr) = let (state, mut restarting_mgr) = match setup_restarting_mgr_std(stats, broker_port) {
match setup_restarting_mgr::<_, _, StdShMem, _>(stats, broker_port) { Ok(res) => res,
Ok(res) => res, Err(err) => match err {
Err(err) => match err { Error::ShuttingDown => {
Error::ShuttingDown => { return Ok(());
return Ok(()); }
} _ => {
_ => { panic!("Failed to setup the restarter: {}", err);
panic!("Failed to setup the restarter: {}", err); }
} },
}, };
};
// Create an observation channel using the coverage map // Create an observation channel using the coverage map
// We don't use the hitcounts (see the Cargo.toml, we use pcguard_edges) // We don't use the hitcounts (see the Cargo.toml, we use pcguard_edges)

View File

@ -39,7 +39,7 @@ std = [] # print, sharedmap, ... support
anymap_debug = ["serde_json"] # uses serde_json to Debug the anymap trait. Disable for smaller footprint. anymap_debug = ["serde_json"] # uses serde_json to Debug the anymap trait. Disable for smaller footprint.
derive = ["libafl_derive"] # provide derive(SerdeAny) macro. derive = ["libafl_derive"] # provide derive(SerdeAny) macro.
llmp_small_maps = [] # reduces initial map size for llmp llmp_small_maps = [] # reduces initial map size for llmp
llmp_debug = [] # Enables debug output for LLMP llmp_debug = ["backtrace"] # Enables debug output for LLMP
[[example]] [[example]]
name = "llmp_test" name = "llmp_test"
@ -59,13 +59,20 @@ ctor = "*"
libafl_derive = { version = "*", optional = true, path = "../libafl_derive" } libafl_derive = { version = "*", optional = true, path = "../libafl_derive" }
serde_json = { version = "1.0", optional = true, default-features = false, features = ["alloc"] } # an easy way to debug print SerdeAnyMap serde_json = { version = "1.0", optional = true, default-features = false, features = ["alloc"] } # an easy way to debug print SerdeAnyMap
num_enum = "0.5.1" num_enum = "0.5.1"
spin = "0.9.0"
backtrace = "0.3" # for llmp_debug [target.'cfg(target_os = "android")'.dependencies]
backtrace = { version = "0.3", optional = true, default-features = false, features = ["std", "libbacktrace"] } # for llmp_debug
[target.'cfg(not(target_os = "android"))'.dependencies]
backtrace = { version = "0.3", optional = true } # for llmp_debug
[target.'cfg(unix)'.dependencies] [target.'cfg(unix)'.dependencies]
libc = "0.2" # For (*nix) libc libc = "0.2" # For (*nix) libc
nix = "0.20.0" nix = "0.20.0"
uds = "0.2.3" uds = "0.2.3"
lock_api = "0.4.3"
regex = "1.4.5"
[target.'cfg(windows)'.dependencies] [target.'cfg(windows)'.dependencies]
windows = "0.4.0" windows = "0.4.0"

View File

@ -25,8 +25,8 @@ const _TAG_1MEG_V1: Tag = 0xB1111161;
#[cfg(all(unix, feature = "std"))] #[cfg(all(unix, feature = "std"))]
fn adder_loop(port: u16) -> ! { fn adder_loop(port: u16) -> ! {
let shmem_provider = Rc::new(RefCell::new(StdShMemProvider::new())); let shmem_provider = StdShMemProvider::new().unwrap();
let mut client = llmp::LlmpClient::create_attach_to_tcp(&shmem_provider, port).unwrap(); let mut client = llmp::LlmpClient::create_attach_to_tcp(shmem_provider, port).unwrap();
let mut last_result: u32 = 0; let mut last_result: u32 = 0;
let mut current_result: u32 = 0; let mut current_result: u32 = 0;
loop { loop {
@ -68,11 +68,8 @@ fn adder_loop(port: u16) -> ! {
#[cfg(all(unix, feature = "std"))] #[cfg(all(unix, feature = "std"))]
fn large_msg_loop(port: u16) -> ! { fn large_msg_loop(port: u16) -> ! {
let mut client = llmp::LlmpClient::create_attach_to_tcp( let mut client =
&Rc::new(RefCell::new(StdShMemProvider::new())), llmp::LlmpClient::create_attach_to_tcp(StdShMemProvider::new().unwrap(), port).unwrap();
port,
)
.unwrap();
let meg_buf = [1u8; 1 << 20]; let meg_buf = [1u8; 1 << 20];
@ -133,8 +130,7 @@ fn main() {
match mode.as_str() { match mode.as_str() {
"broker" => { "broker" => {
let mut broker = let mut broker = llmp::LlmpBroker::new(StdShMemProvider::new().unwrap()).unwrap();
llmp::LlmpBroker::new(&Rc::new(RefCell::new(StdShMemProvider::new()))).unwrap();
broker broker
.launch_listener(llmp::Listener::Tcp( .launch_listener(llmp::Listener::Tcp(
std::net::TcpListener::bind(format!("127.0.0.1:{}", port)).unwrap(), std::net::TcpListener::bind(format!("127.0.0.1:{}", port)).unwrap(),
@ -143,11 +139,9 @@ fn main() {
broker.loop_forever(&mut broker_message_hook, Some(Duration::from_millis(5))) broker.loop_forever(&mut broker_message_hook, Some(Duration::from_millis(5)))
} }
"ctr" => { "ctr" => {
let mut client = llmp::LlmpClient::create_attach_to_tcp( let mut client =
&Rc::new(RefCell::new(StdShMemProvider::new())), llmp::LlmpClient::create_attach_to_tcp(StdShMemProvider::new().unwrap(), port)
port, .unwrap();
)
.unwrap();
let mut counter: u32 = 0; let mut counter: u32 = 0;
loop { loop {
counter = counter.wrapping_add(1); counter = counter.wrapping_add(1);

View File

@ -52,9 +52,8 @@ Then register some clientloops using llmp_broker_register_threaded_clientloop
*/ */
use alloc::{rc::Rc, string::String, vec::Vec}; use alloc::{string::String, vec::Vec};
use core::{ use core::{
cell::RefCell,
cmp::max, cmp::max,
fmt::Debug, fmt::Debug,
mem::size_of, mem::size_of,
@ -75,11 +74,13 @@ use std::{
use backtrace::Backtrace; use backtrace::Backtrace;
#[cfg(unix)] #[cfg(unix)]
use crate::bolts::os::unix_signals::{c_void, setup_signal_handler, siginfo_t, Handler, Signal}; use crate::bolts::os::unix_signals::{setup_signal_handler, siginfo_t, Handler, Signal};
use crate::{ use crate::{
bolts::shmem::{ShMem, ShMemDescription, ShMemId, ShMemProvider}, bolts::shmem::{ShMem, ShMemDescription, ShMemId, ShMemProvider},
Error, Error,
}; };
#[cfg(unix)]
use libc::ucontext_t;
/// We'll start off with 256 megabyte maps per fuzzer client /// We'll start off with 256 megabyte maps per fuzzer client
#[cfg(not(feature = "llmp_small_maps"))] #[cfg(not(feature = "llmp_small_maps"))]
@ -215,8 +216,12 @@ fn new_map_size(max_alloc: usize) -> usize {
/// Initialize a new llmp_page. size should be relative to /// Initialize a new llmp_page. size should be relative to
/// llmp_page->messages /// llmp_page->messages
unsafe fn _llmp_page_init<SHM: ShMem>(shmem: &mut SHM, sender: u32, allow_reinit: bool) { unsafe fn _llmp_page_init<SHM: ShMem>(shmem: &mut SHM, sender: u32, allow_reinit: bool) {
let map_size = shmem.map().len(); #[cfg(all(feature = "llmp_debug", feature = "std"))]
dbg!("_llmp_page_init: shmem {}", &shmem);
let map_size = shmem.len();
let page = shmem2page_mut(shmem); let page = shmem2page_mut(shmem);
#[cfg(all(feature = "llmp_debug", feature = "std"))]
dbg!("_llmp_page_init: page {}", *page);
if (*page).magic == PAGE_INITIALIZED_MAGIC && !allow_reinit { if (*page).magic == PAGE_INITIALIZED_MAGIC && !allow_reinit {
panic!( panic!(
"Tried to initialize page {:?} twice (for shmem {:?})", "Tried to initialize page {:?} twice (for shmem {:?})",
@ -234,6 +239,7 @@ unsafe fn _llmp_page_init<SHM: ShMem>(shmem: &mut SHM, sender: u32, allow_reinit
(*(*page).messages.as_mut_ptr()).tag = LLMP_TAG_UNSET; (*(*page).messages.as_mut_ptr()).tag = LLMP_TAG_UNSET;
ptr::write_volatile(&mut (*page).save_to_unmap, 0); ptr::write_volatile(&mut (*page).save_to_unmap, 0);
ptr::write_volatile(&mut (*page).sender_dead, 0); ptr::write_volatile(&mut (*page).sender_dead, 0);
assert!((*page).size_total != 0);
} }
/// Get the next pointer and make sure it's in the current page, and has enough space. /// Get the next pointer and make sure it's in the current page, and has enough space.
@ -365,7 +371,7 @@ where
{ {
#[cfg(feature = "std")] #[cfg(feature = "std")]
/// Creates either a broker, if the tcp port is not bound, or a client, connected to this port. /// Creates either a broker, if the tcp port is not bound, or a client, connected to this port.
pub fn on_port(shmem_provider: &Rc<RefCell<SP>>, port: u16) -> Result<Self, Error> { pub fn on_port(shmem_provider: SP, port: u16) -> Result<Self, Error> {
match TcpListener::bind(format!("127.0.0.1:{}", port)) { match TcpListener::bind(format!("127.0.0.1:{}", port)) {
Ok(listener) => { Ok(listener) => {
// We got the port. We are the broker! :) // We got the port. We are the broker! :)
@ -390,12 +396,6 @@ where
} }
} }
pub fn shmem_provider(&mut self) -> &Rc<RefCell<SP>> {
match self {
LlmpConnection::IsBroker { broker } => &broker.shmem_provider,
LlmpConnection::IsClient { client } => &client.shmem_provider,
}
}
/// Describe this in a reproducable fashion, if it's a client /// Describe this in a reproducable fashion, if it's a client
pub fn describe(&self) -> Result<LlmpClientDescription, Error> { pub fn describe(&self) -> Result<LlmpClientDescription, Error> {
Ok(match self { Ok(match self {
@ -406,7 +406,7 @@ where
/// Recreate an existing client from the stored description /// Recreate an existing client from the stored description
pub fn existing_client_from_description( pub fn existing_client_from_description(
shmem_provider: &Rc<RefCell<SP>>, shmem_provider: SP,
description: &LlmpClientDescription, description: &LlmpClientDescription,
) -> Result<LlmpConnection<SP>, Error> { ) -> Result<LlmpConnection<SP>, Error> {
Ok(LlmpConnection::IsClient { Ok(LlmpConnection::IsClient {
@ -480,7 +480,7 @@ where
/// By keeping the message history around, /// By keeping the message history around,
/// new clients may join at any time in the future. /// new clients may join at any time in the future.
pub keep_pages_forever: bool, pub keep_pages_forever: bool,
shmem_provider: Rc<RefCell<SP>>, shmem_provider: SP,
} }
/// An actor on the sending part of the shared map /// An actor on the sending part of the shared map
@ -488,23 +488,17 @@ impl<SP> LlmpSender<SP>
where where
SP: ShMemProvider, SP: ShMemProvider,
{ {
pub fn new( pub fn new(mut shmem_provider: SP, id: u32, keep_pages_forever: bool) -> Result<Self, Error> {
shmem_provider: &Rc<RefCell<SP>>,
id: u32,
keep_pages_forever: bool,
) -> Result<Self, Error> {
Ok(Self { Ok(Self {
id, id,
last_msg_sent: ptr::null_mut(), last_msg_sent: ptr::null_mut(),
out_maps: vec![LlmpSharedMap::new( out_maps: vec![LlmpSharedMap::new(
0, 0,
shmem_provider shmem_provider.new_map(LLMP_CFG_INITIAL_MAP_SIZE)?,
.borrow_mut()
.new_map(LLMP_CFG_INITIAL_MAP_SIZE)?,
)], )],
// drop pages to the broker if it already read them // drop pages to the broker if it already read them
keep_pages_forever, keep_pages_forever,
shmem_provider: shmem_provider.clone(), shmem_provider,
}) })
} }
@ -520,14 +514,11 @@ where
/// Reattach to a vacant out_map, to with a previous sender stored the information in an env before. /// Reattach to a vacant out_map, to with a previous sender stored the information in an env before.
#[cfg(feature = "std")] #[cfg(feature = "std")]
pub fn on_existing_from_env( pub fn on_existing_from_env(mut shmem_provider: SP, env_name: &str) -> Result<Self, Error> {
shmem_provider: &Rc<RefCell<SP>>,
env_name: &str,
) -> Result<Self, Error> {
let msg_sent_offset = msg_offset_from_env(env_name)?; let msg_sent_offset = msg_offset_from_env(env_name)?;
Self::on_existing_map( Self::on_existing_map(
shmem_provider.clone(), shmem_provider.clone(),
shmem_provider.borrow_mut().existing_from_env(env_name)?, shmem_provider.existing_from_env(env_name)?,
msg_sent_offset, msg_sent_offset,
) )
} }
@ -565,7 +556,7 @@ where
/// It is essential, that the receiver (or someone else) keeps a pointer to this map /// It is essential, that the receiver (or someone else) keeps a pointer to this map
/// else reattach will get a new, empty page, from the OS, or fail. /// else reattach will get a new, empty page, from the OS, or fail.
pub fn on_existing_map( pub fn on_existing_map(
shmem_provider: Rc<RefCell<SP>>, shmem_provider: SP,
current_out_map: SP::Mem, current_out_map: SP::Mem,
last_msg_sent_offset: Option<u64>, last_msg_sent_offset: Option<u64>,
) -> Result<Self, Error> { ) -> Result<Self, Error> {
@ -649,7 +640,7 @@ where
#[cfg(all(feature = "llmp_debug", feature = "std"))] #[cfg(all(feature = "llmp_debug", feature = "std"))]
println!( println!(
"Allocating {} (>={}) bytes on page {:?} / map {:?} (last msg: {:?})", "Allocating {} (>={}) bytes on page {:?} / map {:?} (last msg: {:?})",
complete_msg_size, buf_len, page, map, last_msg complete_msg_size, buf_len, page, &map, last_msg
); );
/* DBG("XXX complete_msg_size %lu (h: %lu)\n", complete_msg_size, sizeof(llmp_message)); */ /* DBG("XXX complete_msg_size %lu (h: %lu)\n", complete_msg_size, sizeof(llmp_message)); */
/* In case we don't have enough space, make sure the next page will be large /* In case we don't have enough space, make sure the next page will be large
@ -795,7 +786,6 @@ where
let mut new_map_shmem = LlmpSharedMap::new( let mut new_map_shmem = LlmpSharedMap::new(
(*old_map).sender, (*old_map).sender,
self.shmem_provider self.shmem_provider
.borrow_mut()
.new_map(new_map_size((*old_map).max_alloc_size))?, .new_map(new_map_size((*old_map).max_alloc_size))?,
); );
let mut new_map = new_map_shmem.page_mut(); let mut new_map = new_map_shmem.page_mut();
@ -909,14 +899,12 @@ where
// Create this client on an existing map from the given description. acquired with `self.describe` // Create this client on an existing map from the given description. acquired with `self.describe`
pub fn on_existing_from_description( pub fn on_existing_from_description(
shmem_provider: &Rc<RefCell<SP>>, mut shmem_provider: SP,
description: &LlmpDescription, description: &LlmpDescription,
) -> Result<Self, Error> { ) -> Result<Self, Error> {
Self::on_existing_map( Self::on_existing_map(
shmem_provider.clone(), shmem_provider.clone(),
shmem_provider shmem_provider.from_description(description.shmem)?,
.borrow_mut()
.from_description(description.shmem)?,
description.last_message_offset, description.last_message_offset,
) )
} }
@ -932,7 +920,7 @@ where
/// Pointer to the last meg this received /// Pointer to the last meg this received
pub last_msg_recvd: *const LlmpMsg, pub last_msg_recvd: *const LlmpMsg,
/// The shmem provider /// The shmem provider
pub shmem_provider: Rc<RefCell<SP>>, pub shmem_provider: SP,
/// current page. After EOP, this gets replaced with the new one /// current page. After EOP, this gets replaced with the new one
pub current_recv_map: LlmpSharedMap<SP::Mem>, pub current_recv_map: LlmpSharedMap<SP::Mem>,
} }
@ -944,13 +932,10 @@ where
{ {
/// Reattach to a vacant recv_map, to with a previous sender stored the information in an env before. /// Reattach to a vacant recv_map, to with a previous sender stored the information in an env before.
#[cfg(feature = "std")] #[cfg(feature = "std")]
pub fn on_existing_from_env( pub fn on_existing_from_env(mut shmem_provider: SP, env_name: &str) -> Result<Self, Error> {
shmem_provider: &Rc<RefCell<SP>>,
env_name: &str,
) -> Result<Self, Error> {
Self::on_existing_map( Self::on_existing_map(
shmem_provider.clone(), shmem_provider.clone(),
shmem_provider.borrow_mut().existing_from_env(env_name)?, shmem_provider.existing_from_env(env_name)?,
msg_offset_from_env(env_name)?, msg_offset_from_env(env_name)?,
) )
} }
@ -968,7 +953,7 @@ where
/// It is essential, that the sender (or someone else) keeps a pointer to the sender_map /// It is essential, that the sender (or someone else) keeps a pointer to the sender_map
/// else reattach will get a new, empty page, from the OS, or fail. /// else reattach will get a new, empty page, from the OS, or fail.
pub fn on_existing_map( pub fn on_existing_map(
shmem_provider: Rc<RefCell<SP>>, shmem_provider: SP,
current_sender_map: SP::Mem, current_sender_map: SP::Mem,
last_msg_recvd_offset: Option<u64>, last_msg_recvd_offset: Option<u64>,
) -> Result<Self, Error> { ) -> Result<Self, Error> {
@ -1053,12 +1038,11 @@ where
ptr::write_volatile(&mut (*page).save_to_unmap, 1); ptr::write_volatile(&mut (*page).save_to_unmap, 1);
// Map the new page. The old one should be unmapped by Drop // Map the new page. The old one should be unmapped by Drop
self.current_recv_map = LlmpSharedMap::existing( self.current_recv_map =
self.shmem_provider.borrow_mut().from_id_and_size( LlmpSharedMap::existing(self.shmem_provider.from_id_and_size(
ShMemId::from_slice(&pageinfo_cpy.shm_str), ShMemId::from_slice(&pageinfo_cpy.shm_str),
pageinfo_cpy.map_size, pageinfo_cpy.map_size,
)?, )?);
);
page = self.current_recv_map.page_mut(); page = self.current_recv_map.page_mut();
// Mark the new page save to unmap also (it's mapped by us, the broker now) // Mark the new page save to unmap also (it's mapped by us, the broker now)
ptr::write_volatile(&mut (*page).save_to_unmap, 1); ptr::write_volatile(&mut (*page).save_to_unmap, 1);
@ -1151,14 +1135,12 @@ where
// Create this client on an existing map from the given description. acquired with `self.describe` // Create this client on an existing map from the given description. acquired with `self.describe`
pub fn on_existing_from_description( pub fn on_existing_from_description(
shmem_provider: &Rc<RefCell<SP>>, mut shmem_provider: SP,
description: &LlmpDescription, description: &LlmpDescription,
) -> Result<Self, Error> { ) -> Result<Self, Error> {
Self::on_existing_map( Self::on_existing_map(
shmem_provider.clone(), shmem_provider.clone(),
shmem_provider shmem_provider.from_description(description.shmem)?,
.borrow_mut()
.from_description(description.shmem)?,
description.last_message_offset, description.last_message_offset,
) )
} }
@ -1220,6 +1202,8 @@ where
if (*ret.page()).magic != PAGE_INITIALIZED_MAGIC { if (*ret.page()).magic != PAGE_INITIALIZED_MAGIC {
panic!("Map was not priviously initialized at {:?}", &ret.shmem); panic!("Map was not priviously initialized at {:?}", &ret.shmem);
} }
#[cfg(all(feature = "llmp_debug", feature = "std"))]
dbg!("PAGE: {}", *ret.page());
} }
ret ret
} }
@ -1328,7 +1312,7 @@ where
/// handlers /// handlers
shutting_down: bool, shutting_down: bool,
/// The ShMemProvider to use /// The ShMemProvider to use
shmem_provider: Rc<RefCell<SP>>, shmem_provider: SP,
} }
#[cfg(unix)] #[cfg(unix)]
@ -1336,9 +1320,9 @@ pub struct LlmpBrokerSignalHandler {
shutting_down: bool, shutting_down: bool,
} }
#[cfg(all(unix))] #[cfg(unix)]
impl Handler for LlmpBrokerSignalHandler { impl Handler for LlmpBrokerSignalHandler {
fn handle(&mut self, _signal: Signal, _info: siginfo_t, _void: *const c_void) { fn handle(&mut self, _signal: Signal, _info: siginfo_t, _context: &mut ucontext_t) {
unsafe { ptr::write_volatile(&mut self.shutting_down, true) }; unsafe { ptr::write_volatile(&mut self.shutting_down, true) };
} }
@ -1354,14 +1338,14 @@ where
SP: ShMemProvider, SP: ShMemProvider,
{ {
/// Create and initialize a new llmp_broker /// Create and initialize a new llmp_broker
pub fn new(shmem_provider: &Rc<RefCell<SP>>) -> Result<Self, Error> { pub fn new(mut shmem_provider: SP) -> Result<Self, Error> {
Ok(LlmpBroker { Ok(LlmpBroker {
llmp_out: LlmpSender { llmp_out: LlmpSender {
id: 0, id: 0,
last_msg_sent: ptr::null_mut(), last_msg_sent: ptr::null_mut(),
out_maps: vec![LlmpSharedMap::new( out_maps: vec![LlmpSharedMap::new(
0, 0,
shmem_provider.borrow_mut().new_map(new_map_size(0))?, shmem_provider.new_map(new_map_size(0))?,
)], )],
// Broker never cleans up the pages so that new // Broker never cleans up the pages so that new
// clients may join at any time // clients may join at any time
@ -1371,7 +1355,7 @@ where
llmp_clients: vec![], llmp_clients: vec![],
socket_name: None, socket_name: None,
shutting_down: false, shutting_down: false,
shmem_provider: shmem_provider.clone(), shmem_provider,
}) })
} }
@ -1515,33 +1499,30 @@ where
// Tcp out map sends messages from background thread tcp server to foreground client // Tcp out map sends messages from background thread tcp server to foreground client
let tcp_out_map = LlmpSharedMap::new( let tcp_out_map = LlmpSharedMap::new(
llmp_tcp_id, llmp_tcp_id,
self.shmem_provider self.shmem_provider.new_map(LLMP_CFG_INITIAL_MAP_SIZE)?,
.borrow_mut()
.new_map(LLMP_CFG_INITIAL_MAP_SIZE)?,
); );
let shmem_id = tcp_out_map.shmem.id(); let shmem_id = tcp_out_map.shmem.id();
let tcp_out_map_str = *shmem_id.as_slice(); let tcp_out_map_str = *shmem_id.as_slice();
let tcp_out_map_size = tcp_out_map.shmem.len(); let tcp_out_map_size = tcp_out_map.shmem.len();
self.register_client(tcp_out_map); self.register_client(tcp_out_map);
let shmem_provider_clone = self.shmem_provider.borrow_mut().clone(); let mut shmem_provider_clone = self.shmem_provider.clone();
Ok(thread::spawn(move || { Ok(thread::spawn(move || {
let shmem_provider = Rc::new(RefCell::new(shmem_provider_clone)); shmem_provider_clone.post_fork();
// Clone so we get a new connection to the AshmemServer if we are using // Clone so we get a new connection to the AshmemServer if we are using
// ServedShMemProvider // ServedShMemProvider
let mut new_client_sender = LlmpSender { let mut new_client_sender = LlmpSender {
id: 0, id: 0,
last_msg_sent: ptr::null_mut(), last_msg_sent: ptr::null_mut(),
out_maps: vec![LlmpSharedMap::existing( out_maps: vec![LlmpSharedMap::existing(
shmem_provider shmem_provider_clone
.borrow_mut()
.from_id_and_size(ShMemId::from_slice(&tcp_out_map_str), tcp_out_map_size) .from_id_and_size(ShMemId::from_slice(&tcp_out_map_str), tcp_out_map_size)
.unwrap(), .unwrap(),
)], )],
// drop pages to the broker if it already read them // drop pages to the broker if it already read them
keep_pages_forever: false, keep_pages_forever: false,
shmem_provider: shmem_provider.clone(), shmem_provider: shmem_provider_clone.clone(),
}; };
loop { loop {
@ -1627,7 +1608,7 @@ where
} else { } else {
let pageinfo = (*msg).buf.as_mut_ptr() as *mut LlmpPayloadSharedMapInfo; let pageinfo = (*msg).buf.as_mut_ptr() as *mut LlmpPayloadSharedMapInfo;
match self.shmem_provider.borrow_mut().from_id_and_size( match self.shmem_provider.from_id_and_size(
ShMemId::from_slice(&(*pageinfo).shm_str), ShMemId::from_slice(&(*pageinfo).shm_str),
(*pageinfo).map_size, (*pageinfo).map_size,
) { ) {
@ -1686,7 +1667,7 @@ pub struct LlmpClient<SP>
where where
SP: ShMemProvider, SP: ShMemProvider,
{ {
shmem_provider: Rc<RefCell<SP>>, shmem_provider: SP,
/// Outgoing channel to the broker /// Outgoing channel to the broker
pub sender: LlmpSender<SP>, pub sender: LlmpSender<SP>,
/// Incoming (broker) broadcast map /// Incoming (broker) broadcast map
@ -1703,7 +1684,7 @@ where
/// It is essential, that the broker (or someone else) kept a pointer to the out_map /// It is essential, that the broker (or someone else) kept a pointer to the out_map
/// else reattach will get a new, empty page, from the OS, or fail /// else reattach will get a new, empty page, from the OS, or fail
pub fn on_existing_map( pub fn on_existing_map(
shmem_provider: Rc<RefCell<SP>>, shmem_provider: SP,
_current_out_map: SP::Mem, _current_out_map: SP::Mem,
_last_msg_sent_offset: Option<u64>, _last_msg_sent_offset: Option<u64>,
current_broker_map: SP::Mem, current_broker_map: SP::Mem,
@ -1726,20 +1707,17 @@ where
/// Recreate this client from a previous client.to_env /// Recreate this client from a previous client.to_env
#[cfg(feature = "std")] #[cfg(feature = "std")]
pub fn on_existing_from_env( pub fn on_existing_from_env(shmem_provider: SP, env_name: &str) -> Result<Self, Error> {
shmem_provider: &Rc<RefCell<SP>>,
env_name: &str,
) -> Result<Self, Error> {
Ok(Self { Ok(Self {
sender: LlmpSender::on_existing_from_env( sender: LlmpSender::on_existing_from_env(
shmem_provider, shmem_provider.clone(),
&format!("{}_SENDER", env_name), &format!("{}_SENDER", env_name),
)?, )?,
receiver: LlmpReceiver::on_existing_from_env( receiver: LlmpReceiver::on_existing_from_env(
shmem_provider, shmem_provider.clone(),
&format!("{}_RECEIVER", env_name), &format!("{}_RECEIVER", env_name),
)?, )?,
shmem_provider: shmem_provider.clone(), shmem_provider,
}) })
} }
@ -1761,16 +1739,19 @@ where
/// Create an existing client from description /// Create an existing client from description
fn existing_client_from_description( fn existing_client_from_description(
shmem_provider: &Rc<RefCell<SP>>, shmem_provider: SP,
description: &LlmpClientDescription, description: &LlmpClientDescription,
) -> Result<Self, Error> { ) -> Result<Self, Error> {
Ok(Self { Ok(Self {
sender: LlmpSender::on_existing_from_description(shmem_provider, &description.sender)?, sender: LlmpSender::on_existing_from_description(
shmem_provider.clone(),
&description.sender,
)?,
receiver: LlmpReceiver::on_existing_from_description( receiver: LlmpReceiver::on_existing_from_description(
shmem_provider, shmem_provider.clone(),
&description.receiver, &description.receiver,
)?, )?,
shmem_provider: shmem_provider.clone(), shmem_provider,
}) })
} }
@ -1787,7 +1768,7 @@ where
/// Creates a new LlmpClient /// Creates a new LlmpClient
pub fn new( pub fn new(
shmem_provider: &Rc<RefCell<SP>>, mut shmem_provider: SP,
initial_broker_map: LlmpSharedMap<SP::Mem>, initial_broker_map: LlmpSharedMap<SP::Mem>,
) -> Result<Self, Error> { ) -> Result<Self, Error> {
Ok(Self { Ok(Self {
@ -1795,9 +1776,7 @@ where
id: 0, id: 0,
last_msg_sent: ptr::null_mut(), last_msg_sent: ptr::null_mut(),
out_maps: vec![LlmpSharedMap::new(0, { out_maps: vec![LlmpSharedMap::new(0, {
shmem_provider shmem_provider.new_map(LLMP_CFG_INITIAL_MAP_SIZE)?
.borrow_mut()
.new_map(LLMP_CFG_INITIAL_MAP_SIZE)?
})], })],
// drop pages to the broker if it already read them // drop pages to the broker if it already read them
keep_pages_forever: false, keep_pages_forever: false,
@ -1810,7 +1789,7 @@ where
last_msg_recvd: ptr::null_mut(), last_msg_recvd: ptr::null_mut(),
shmem_provider: shmem_provider.clone(), shmem_provider: shmem_provider.clone(),
}, },
shmem_provider: shmem_provider.clone(), shmem_provider,
}) })
} }
@ -1887,20 +1866,14 @@ where
#[cfg(feature = "std")] #[cfg(feature = "std")]
/// Creates a new LlmpClient, reading the map id and len from env /// Creates a new LlmpClient, reading the map id and len from env
pub fn create_using_env( pub fn create_using_env(mut shmem_provider: SP, env_var: &str) -> Result<Self, Error> {
shmem_provider: &Rc<RefCell<SP>>, let map = LlmpSharedMap::existing(shmem_provider.existing_from_env(env_var)?);
env_var: &str,
) -> Result<Self, Error> {
let map = LlmpSharedMap::existing(shmem_provider.borrow_mut().existing_from_env(env_var)?);
Self::new(shmem_provider, map) Self::new(shmem_provider, map)
} }
#[cfg(feature = "std")] #[cfg(feature = "std")]
/// Create a LlmpClient, getting the ID from a given port /// Create a LlmpClient, getting the ID from a given port
pub fn create_attach_to_tcp( pub fn create_attach_to_tcp(mut shmem_provider: SP, port: u16) -> Result<Self, Error> {
shmem_provider: &Rc<RefCell<SP>>,
port: u16,
) -> Result<Self, Error> {
let mut stream = TcpStream::connect(format!("127.0.0.1:{}", port))?; let mut stream = TcpStream::connect(format!("127.0.0.1:{}", port))?;
println!("Connected to port {}", port); println!("Connected to port {}", port);
@ -1915,11 +1888,7 @@ where
let broker_map_description: ShMemDescription = postcard::from_bytes(&new_broker_map_str)?; let broker_map_description: ShMemDescription = postcard::from_bytes(&new_broker_map_str)?;
let map = LlmpSharedMap::existing( let map = LlmpSharedMap::existing(shmem_provider.from_description(broker_map_description)?);
shmem_provider
.borrow_mut()
.from_description(broker_map_description)?,
);
let ret = Self::new(shmem_provider, map)?; let ret = Self::new(shmem_provider, map)?;
let own_map_description_bytes = let own_map_description_bytes =
@ -1933,7 +1902,6 @@ where
#[cfg(all(unix, feature = "std"))] #[cfg(all(unix, feature = "std"))]
mod tests { mod tests {
use alloc::rc::Rc;
use std::{thread::sleep, time::Duration}; use std::{thread::sleep, time::Duration};
use super::{ use super::{
@ -1945,18 +1913,16 @@ mod tests {
use crate::bolts::shmem::{ShMemProvider, StdShMemProvider}; use crate::bolts::shmem::{ShMemProvider, StdShMemProvider};
use core::cell::RefCell;
#[test] #[test]
pub fn llmp_connection() { pub fn llmp_connection() {
let shmem_provider = Rc::new(RefCell::new(StdShMemProvider::new())); let shmem_provider = StdShMemProvider::new().unwrap();
let mut broker = match LlmpConnection::on_port(&shmem_provider, 1337).unwrap() { let mut broker = match LlmpConnection::on_port(shmem_provider.clone(), 1337).unwrap() {
IsClient { client: _ } => panic!("Could not bind to port as broker"), IsClient { client: _ } => panic!("Could not bind to port as broker"),
IsBroker { broker } => broker, IsBroker { broker } => broker,
}; };
// Add the first client (2nd, actually, because of the tcp listener client) // Add the first client (2nd, actually, because of the tcp listener client)
let mut client = match LlmpConnection::on_port(&shmem_provider, 1337).unwrap() { let mut client = match LlmpConnection::on_port(shmem_provider.clone(), 1337).unwrap() {
IsBroker { broker: _ } => panic!("Second connect should be a client!"), IsBroker { broker: _ } => panic!("Second connect should be a client!"),
IsClient { client } => client, IsClient { client } => client,
}; };
@ -1973,6 +1939,7 @@ mod tests {
client.send_buf(tag, &arr).unwrap(); client.send_buf(tag, &arr).unwrap();
client.to_env("_ENV_TEST").unwrap(); client.to_env("_ENV_TEST").unwrap();
#[cfg(all(feature = "llmp_debug", feature = "std"))]
dbg!(std::env::vars()); dbg!(std::env::vars());
for (key, value) in std::env::vars_os() { for (key, value) in std::env::vars_os() {
@ -1980,7 +1947,7 @@ mod tests {
} }
/* recreate the client from env, check if it still works */ /* recreate the client from env, check if it still works */
client = LlmpClient::on_existing_from_env(&shmem_provider, "_ENV_TEST").unwrap(); client = LlmpClient::on_existing_from_env(shmem_provider, "_ENV_TEST").unwrap();
client.send_buf(tag, &arr).unwrap(); client.send_buf(tag, &arr).unwrap();

View File

@ -11,10 +11,13 @@ use crate::{
}, },
Error, Error,
}; };
use core::mem::ManuallyDrop;
use hashbrown::HashMap; use hashbrown::HashMap;
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use std::{ use std::{
cell::RefCell,
io::{Read, Write}, io::{Read, Write},
rc::Rc,
sync::{Arc, Condvar, Mutex}, sync::{Arc, Condvar, Mutex},
}; };
@ -39,11 +42,12 @@ const ASHMEM_SERVER_NAME: &str = "@ashmem_server";
pub struct ServedShMemProvider { pub struct ServedShMemProvider {
stream: UnixStream, stream: UnixStream,
inner: AshmemShMemProvider, inner: AshmemShMemProvider,
id: i32,
} }
#[derive(Clone, Debug)] #[derive(Clone, Debug)]
pub struct ServedShMem { pub struct ServedShMem {
inner: AshmemShMem, inner: ManuallyDrop<AshmemShMem>,
server_fd: i32, server_fd: i32,
} }
@ -95,13 +99,13 @@ impl ServedShMemProvider {
impl Default for ServedShMemProvider { impl Default for ServedShMemProvider {
fn default() -> Self { fn default() -> Self {
Self::new() Self::new().unwrap()
} }
} }
impl Clone for ServedShMemProvider { impl Clone for ServedShMemProvider {
fn clone(&self) -> Self { fn clone(&self) -> Self {
Self::new() Self::new().unwrap()
} }
} }
@ -109,22 +113,26 @@ impl ShMemProvider for ServedShMemProvider {
type Mem = ServedShMem; type Mem = ServedShMem;
/// Connect to the server and return a new ServedShMemProvider /// Connect to the server and return a new ServedShMemProvider
fn new() -> Self { fn new() -> Result<Self, Error> {
Self { let mut res = Self {
stream: UnixStream::connect_to_unix_addr( stream: UnixStream::connect_to_unix_addr(
&UnixSocketAddr::new(ASHMEM_SERVER_NAME).unwrap(), &UnixSocketAddr::new(ASHMEM_SERVER_NAME).unwrap(),
) )?,
.expect("Unable to open connection to ashmem service"), inner: AshmemShMemProvider::new()?,
inner: AshmemShMemProvider::new(), id: -1,
} };
let (id, _) = res.send_receive(AshmemRequest::Hello(None));
res.id = id;
Ok(res)
} }
fn new_map(&mut self, map_size: usize) -> Result<Self::Mem, crate::Error> { fn new_map(&mut self, map_size: usize) -> Result<Self::Mem, crate::Error> {
let (server_fd, client_fd) = self.send_receive(AshmemRequest::NewMap(map_size)); let (server_fd, client_fd) = self.send_receive(AshmemRequest::NewMap(map_size));
Ok(ServedShMem { Ok(ServedShMem {
inner: self inner: ManuallyDrop::new(
.inner self.inner
.from_id_and_size(ShMemId::from_string(&format!("{}", client_fd)), map_size)?, .from_id_and_size(ShMemId::from_string(&format!("{}", client_fd)), map_size)?,
),
server_fd, server_fd,
}) })
} }
@ -136,12 +144,30 @@ impl ShMemProvider for ServedShMemProvider {
ShMemDescription::from_string_and_size(server_id_str, size), ShMemDescription::from_string_and_size(server_id_str, size),
)); ));
Ok(ServedShMem { Ok(ServedShMem {
inner: self inner: ManuallyDrop::new(
.inner self.inner
.from_id_and_size(ShMemId::from_string(&format!("{}", client_fd)), size)?, .from_id_and_size(ShMemId::from_string(&format!("{}", client_fd)), size)?,
),
server_fd, server_fd,
}) })
} }
fn post_fork(&mut self) {
self.stream =
UnixStream::connect_to_unix_addr(&UnixSocketAddr::new(ASHMEM_SERVER_NAME).unwrap())
.expect("Unable to reconnect to the ashmem service");
let (id, _) = self.send_receive(AshmemRequest::Hello(Some(self.id)));
self.id = id;
}
fn release_map(&mut self, map: &mut Self::Mem) {
let (refcount, _) = self.send_receive(AshmemRequest::Deregister(map.server_fd));
if refcount == 0 {
unsafe {
ManuallyDrop::drop(&mut map.inner);
}
}
}
} }
/// A request sent to the ShMem server to receive a fd to a shared map /// A request sent to the ShMem server to receive a fd to a shared map
@ -152,38 +178,119 @@ pub enum AshmemRequest {
/// Another client already has a map with this description mapped. /// Another client already has a map with this description mapped.
ExistingMap(ShMemDescription), ExistingMap(ShMemDescription),
/// A client tells us it unregisters the previously allocated map /// A client tells us it unregisters the previously allocated map
Deregister(u32), Deregister(i32),
/// A message that tells us hello, and optionally which other client we were created from, we
/// return a client id.
Hello(Option<i32>),
} }
#[derive(Debug)] #[derive(Debug)]
struct AshmemClient { struct AshmemClient {
stream: UnixStream, stream: UnixStream,
maps: HashMap<i32, Vec<Rc<RefCell<AshmemShMem>>>>,
} }
impl AshmemClient { impl AshmemClient {
fn new(stream: UnixStream) -> Self { fn new(stream: UnixStream) -> Self {
Self { stream } Self {
stream,
maps: HashMap::new(),
}
} }
} }
#[derive(Debug)] #[derive(Debug)]
pub struct AshmemService { pub struct AshmemService {
provider: AshmemShMemProvider, provider: AshmemShMemProvider,
maps: Vec<AshmemShMem>, clients: HashMap<RawFd, AshmemClient>,
all_maps: HashMap<i32, Rc<RefCell<AshmemShMem>>>,
}
#[derive(Debug)]
enum AshmemResponse {
Mapping(Rc<RefCell<AshmemShMem>>),
Id(i32),
RefCount(u32),
} }
impl AshmemService { impl AshmemService {
/// Create a new AshMem service /// Create a new AshMem service
#[must_use] #[must_use]
fn new() -> Self { fn new() -> Result<Self, Error> {
AshmemService { Ok(AshmemService {
provider: AshmemShMemProvider::new(), provider: AshmemShMemProvider::new()?,
maps: Vec::new(), clients: HashMap::new(),
} all_maps: HashMap::new(),
})
} }
/// Read and handle the client request, send the answer over unix fd. /// Read and handle the client request, send the answer over unix fd.
fn handle_client(&mut self, client: &mut AshmemClient) -> Result<(), Error> { fn handle_request(&mut self, client_id: RawFd) -> Result<AshmemResponse, Error> {
let request = self.read_request(client_id)?;
//println!("got ashmem client: {}, request:{:?}", client_id, request);
// Handle the client request
let response = match request {
AshmemRequest::Hello(other_id) => {
if let Some(other_id) = other_id {
if other_id != client_id {
// remove temporarily
let other_client = self.clients.remove(&other_id);
let client = self.clients.get_mut(&client_id).unwrap();
for (id, map) in other_client.as_ref().unwrap().maps.iter() {
client.maps.insert(*id, map.clone());
}
self.clients.insert(other_id, other_client.unwrap());
}
};
Ok(AshmemResponse::Id(client_id))
}
AshmemRequest::NewMap(map_size) => Ok(AshmemResponse::Mapping(Rc::new(RefCell::new(
self.provider.new_map(map_size)?,
)))),
AshmemRequest::ExistingMap(description) => {
let client = self.clients.get_mut(&client_id).unwrap();
if client.maps.contains_key(&description.id.to_int()) {
Ok(AshmemResponse::Mapping(
client
.maps
.get_mut(&description.id.to_int())
.as_mut()
.unwrap()
.first()
.as_mut()
.unwrap()
.clone(),
))
} else if self.all_maps.contains_key(&description.id.to_int()) {
Ok(AshmemResponse::Mapping(
self.all_maps
.get_mut(&description.id.to_int())
.unwrap()
.clone(),
))
} else {
let new_rc =
Rc::new(RefCell::new(self.provider.from_description(description)?));
self.all_maps
.insert(description.id.to_int(), new_rc.clone());
Ok(AshmemResponse::Mapping(new_rc))
}
}
AshmemRequest::Deregister(map_id) => {
let client = self.clients.get_mut(&client_id).unwrap();
let map = client.maps.entry(map_id).or_default().pop().unwrap();
Ok(AshmemResponse::RefCount(Rc::strong_count(&map) as u32))
}
};
//println!("send ashmem client: {}, response: {:?}", client_id, &response);
response
}
fn read_request(&mut self, client_id: RawFd) -> Result<AshmemRequest, Error> {
let client = self.clients.get_mut(&client_id).unwrap();
// Always receive one be u32 of size, then the command. // Always receive one be u32 of size, then the command.
let mut size_bytes = [0u8; 4]; let mut size_bytes = [0u8; 4];
client.stream.read_exact(&mut size_bytes)?; client.stream.read_exact(&mut size_bytes)?;
@ -196,23 +303,36 @@ impl AshmemService {
.expect("Failed to read message body"); .expect("Failed to read message body");
let request: AshmemRequest = postcard::from_bytes(&bytes)?; let request: AshmemRequest = postcard::from_bytes(&bytes)?;
// Handle the client request Ok(request)
let mapping = match request { }
AshmemRequest::NewMap(map_size) => self.provider.new_map(map_size)?, fn handle_client(&mut self, client_id: RawFd) -> Result<(), Error> {
AshmemRequest::ExistingMap(description) => { let response = self.handle_request(client_id)?;
self.provider.from_description(description)?
}
AshmemRequest::Deregister(_) => {
return Ok(());
}
};
let id = mapping.id(); match response {
let server_fd: i32 = id.to_string().parse().unwrap(); AshmemResponse::Mapping(mapping) => {
client let id = mapping.borrow().id();
.stream let server_fd: i32 = id.to_string().parse().unwrap();
.send_fds(&id.to_string().as_bytes(), &[server_fd])?; let client = self.clients.get_mut(&client_id).unwrap();
self.maps.push(mapping); client
.stream
.send_fds(&id.to_string().as_bytes(), &[server_fd])?;
client
.maps
.entry(server_fd)
.or_default()
.push(mapping.clone());
}
AshmemResponse::Id(id) => {
let client = self.clients.get_mut(&client_id).unwrap();
client.stream.send_fds(&id.to_string().as_bytes(), &[])?;
}
AshmemResponse::RefCount(refcount) => {
let client = self.clients.get_mut(&client_id).unwrap();
client
.stream
.send_fds(&refcount.to_string().as_bytes(), &[])?;
}
}
Ok(()) Ok(())
} }
@ -222,7 +342,7 @@ impl AshmemService {
let syncpair = Arc::new((Mutex::new(false), Condvar::new())); let syncpair = Arc::new((Mutex::new(false), Condvar::new()));
let childsyncpair = Arc::clone(&syncpair); let childsyncpair = Arc::clone(&syncpair);
let join_handle = let join_handle =
thread::spawn(move || Self::new().listen(ASHMEM_SERVER_NAME, childsyncpair)); thread::spawn(move || Self::new()?.listen(ASHMEM_SERVER_NAME, childsyncpair));
let (lock, cvar) = &*syncpair; let (lock, cvar) = &*syncpair;
let mut started = lock.lock().unwrap(); let mut started = lock.lock().unwrap();
@ -252,7 +372,6 @@ impl AshmemService {
"The server appears to already be running. We are probably a client".to_string(), "The server appears to already be running. We are probably a client".to_string(),
)); ));
}; };
let mut clients: HashMap<RawFd, AshmemClient> = HashMap::new();
let mut poll_fds: Vec<PollFd> = vec![PollFd::new( let mut poll_fds: Vec<PollFd> = vec![PollFd::new(
listener.as_raw_fd(), listener.as_raw_fd(),
PollFlags::POLLIN | PollFlags::POLLRDNORM | PollFlags::POLLRDBAND, PollFlags::POLLIN | PollFlags::POLLRDNORM | PollFlags::POLLRDBAND,
@ -278,11 +397,10 @@ impl AshmemService {
unsafe { *((&poll_fd as *const PollFd) as *const libc::pollfd) }.fd; unsafe { *((&poll_fd as *const PollFd) as *const libc::pollfd) }.fd;
if revents.contains(PollFlags::POLLHUP) { if revents.contains(PollFlags::POLLHUP) {
poll_fds.remove(poll_fds.iter().position(|item| *item == poll_fd).unwrap()); poll_fds.remove(poll_fds.iter().position(|item| *item == poll_fd).unwrap());
clients.remove(&raw_polled_fd); self.clients.remove(&raw_polled_fd);
} else if revents.contains(PollFlags::POLLIN) { } else if revents.contains(PollFlags::POLLIN) {
if clients.contains_key(&raw_polled_fd) { if self.clients.contains_key(&raw_polled_fd) {
let mut client = clients.get_mut(&raw_polled_fd).unwrap(); match self.handle_client(raw_polled_fd) {
match self.handle_client(&mut client) {
Ok(()) => (), Ok(()) => (),
Err(e) => { Err(e) => {
dbg!("Ignoring failed read from client", e, poll_fd); dbg!("Ignoring failed read from client", e, poll_fd);
@ -304,14 +422,15 @@ impl AshmemService {
PollFlags::POLLIN | PollFlags::POLLRDNORM | PollFlags::POLLRDBAND, PollFlags::POLLIN | PollFlags::POLLRDNORM | PollFlags::POLLRDBAND,
); );
poll_fds.push(pollfd); poll_fds.push(pollfd);
let mut client = AshmemClient::new(stream); let client = AshmemClient::new(stream);
match self.handle_client(&mut client) { let client_id = client.stream.as_raw_fd();
self.clients.insert(client_id, client);
match self.handle_client(client_id) {
Ok(()) => (), Ok(()) => (),
Err(e) => { Err(e) => {
dbg!("Ignoring failed read from client", e); dbg!("Ignoring failed read from client", e);
} }
}; };
clients.insert(client.stream.as_raw_fd(), client);
} }
} else { } else {
//println!("Unknown revents flags: {:?}", revents); //println!("Unknown revents flags: {:?}", revents);

View File

@ -13,9 +13,9 @@ use core::{
use std::ffi::CString; use std::ffi::CString;
use libc::{ use libc::{
c_int, malloc, sigaction, sigaltstack, sigemptyset, stack_t, SA_NODEFER, SA_ONSTACK, c_int, malloc, sigaction, sigaltstack, sigemptyset, stack_t, ucontext_t, SA_NODEFER,
SA_SIGINFO, SIGABRT, SIGALRM, SIGBUS, SIGFPE, SIGHUP, SIGILL, SIGINT, SIGKILL, SIGPIPE, SA_ONSTACK, SA_SIGINFO, SIGABRT, SIGALRM, SIGBUS, SIGFPE, SIGHUP, SIGILL, SIGINT, SIGKILL,
SIGQUIT, SIGSEGV, SIGTERM, SIGUSR2, SIGPIPE, SIGQUIT, SIGSEGV, SIGTERM, SIGTRAP, SIGUSR2,
}; };
use num_enum::{IntoPrimitive, TryFromPrimitive}; use num_enum::{IntoPrimitive, TryFromPrimitive};
@ -40,6 +40,7 @@ pub enum Signal {
SigQuit = SIGQUIT, SigQuit = SIGQUIT,
SigTerm = SIGTERM, SigTerm = SIGTERM,
SigInterrupt = SIGINT, SigInterrupt = SIGINT,
SigTrap = SIGTRAP,
} }
pub static CRASH_SIGNALS: &[Signal] = &[ pub static CRASH_SIGNALS: &[Signal] = &[
@ -77,6 +78,7 @@ impl Display for Signal {
Signal::SigQuit => write!(f, "SIGQUIT")?, Signal::SigQuit => write!(f, "SIGQUIT")?,
Signal::SigTerm => write!(f, "SIGTERM")?, Signal::SigTerm => write!(f, "SIGTERM")?,
Signal::SigInterrupt => write!(f, "SIGINT")?, Signal::SigInterrupt => write!(f, "SIGINT")?,
Signal::SigTrap => write!(f, "SIGTRAP")?,
}; };
Ok(()) Ok(())
@ -85,7 +87,7 @@ impl Display for Signal {
pub trait Handler { pub trait Handler {
/// Handle a signal /// Handle a signal
fn handle(&mut self, signal: Signal, info: siginfo_t, _void: *const c_void); fn handle(&mut self, signal: Signal, info: siginfo_t, _context: &mut ucontext_t);
/// Return a list of signals to handle /// Return a list of signals to handle
fn signals(&self) -> Vec<Signal>; fn signals(&self) -> Vec<Signal>;
} }
@ -113,7 +115,7 @@ static mut SIGNAL_HANDLERS: [Option<HandlerHolder>; 32] = [
/// # Safety /// # Safety
/// This should be somewhat safe to call for signals previously registered, /// This should be somewhat safe to call for signals previously registered,
/// unless the signal handlers registered using [setup_signal_handler] are broken. /// unless the signal handlers registered using [setup_signal_handler] are broken.
unsafe fn handle_signal(sig: c_int, info: siginfo_t, void: *const c_void) { unsafe fn handle_signal(sig: c_int, info: siginfo_t, void: *mut c_void) {
let signal = &Signal::try_from(sig).unwrap(); let signal = &Signal::try_from(sig).unwrap();
let handler = { let handler = {
match &SIGNAL_HANDLERS[*signal as usize] { match &SIGNAL_HANDLERS[*signal as usize] {
@ -121,7 +123,7 @@ unsafe fn handle_signal(sig: c_int, info: siginfo_t, void: *const c_void) {
None => return, None => return,
} }
}; };
handler.handle(*signal, info, void); handler.handle(*signal, info, &mut *(void as *mut ucontext_t));
} }
/// Setup signal handlers in a somewhat rusty way. /// Setup signal handlers in a somewhat rusty way.

View File

@ -16,11 +16,11 @@ pub type OsShMemProvider = Win32ShMemProvider;
pub type OsShMem = Win32ShMem; pub type OsShMem = Win32ShMem;
#[cfg(target_os = "android")] #[cfg(target_os = "android")]
use crate::bolts::os::ashmem_server::{ServedShMem, ServedShMemProvider}; use crate::bolts::os::ashmem_server::ServedShMemProvider;
#[cfg(target_os = "android")] #[cfg(target_os = "android")]
pub type StdShMemProvider = ServedShMemProvider; pub type StdShMemProvider = RcShMemProvider<ServedShMemProvider>;
#[cfg(target_os = "android")] #[cfg(target_os = "android")]
pub type StdShMem = ServedShMem; pub type StdShMem = RcShMem<ServedShMemProvider>;
#[cfg(all(feature = "std", not(target_os = "android")))] #[cfg(all(feature = "std", not(target_os = "android")))]
pub type StdShMemProvider = OsShMemProvider; pub type StdShMemProvider = OsShMemProvider;
@ -32,7 +32,9 @@ use serde::{Deserialize, Serialize};
#[cfg(feature = "std")] #[cfg(feature = "std")]
use std::env; use std::env;
use alloc::string::ToString; use alloc::{rc::Rc, string::ToString};
use core::cell::RefCell;
use core::mem::ManuallyDrop;
use crate::Error; use crate::Error;
@ -137,11 +139,11 @@ pub trait ShMem: Sized + Debug + Clone {
} }
} }
pub trait ShMemProvider: Send + Clone + Default { pub trait ShMemProvider: Send + Clone + Default + Debug {
type Mem: ShMem; type Mem: ShMem;
/// Create a new instance of the provider /// Create a new instance of the provider
fn new() -> Self; fn new() -> Result<Self, Error>;
/// Create a new shared memory mapping /// Create a new shared memory mapping
fn new_map(&mut self, map_size: usize) -> Result<Self::Mem, Error>; fn new_map(&mut self, map_size: usize) -> Result<Self::Mem, Error>;
@ -168,6 +170,108 @@ pub trait ShMemProvider: Send + Clone + Default {
map_size, map_size,
)) ))
} }
/// This method should be called after a fork or thread creation event, allowing the ShMem to
/// reset thread specific info.
fn post_fork(&mut self) {
// do nothing
}
/// Release the resources associated with the given ShMem
fn release_map(&mut self, _map: &mut Self::Mem) {
// do nothing
}
}
#[derive(Debug, Clone)]
pub struct RcShMem<T: ShMemProvider> {
internal: ManuallyDrop<T::Mem>,
provider: Rc<RefCell<T>>,
}
impl<T> ShMem for RcShMem<T>
where
T: ShMemProvider + alloc::fmt::Debug,
{
fn id(&self) -> ShMemId {
self.internal.id()
}
fn len(&self) -> usize {
self.internal.len()
}
fn map(&self) -> &[u8] {
self.internal.map()
}
fn map_mut(&mut self) -> &mut [u8] {
self.internal.map_mut()
}
}
impl<T: ShMemProvider> Drop for RcShMem<T> {
fn drop(&mut self) {
self.provider.borrow_mut().release_map(&mut self.internal)
}
}
#[derive(Debug, Clone)]
pub struct RcShMemProvider<T: ShMemProvider> {
internal: Rc<RefCell<T>>,
}
unsafe impl<T: ShMemProvider> Send for RcShMemProvider<T> {}
impl<T> ShMemProvider for RcShMemProvider<T>
where
T: ShMemProvider + alloc::fmt::Debug,
{
type Mem = RcShMem<T>;
fn new() -> Result<Self, Error> {
return Ok(Self {
internal: Rc::new(RefCell::new(T::new()?)),
});
}
fn new_map(&mut self, map_size: usize) -> Result<Self::Mem, Error> {
Ok(Self::Mem {
internal: ManuallyDrop::new(self.internal.borrow_mut().new_map(map_size)?),
provider: self.internal.clone(),
})
}
fn from_id_and_size(&mut self, id: ShMemId, size: usize) -> Result<Self::Mem, Error> {
Ok(Self::Mem {
internal: ManuallyDrop::new(self.internal.borrow_mut().from_id_and_size(id, size)?),
provider: self.internal.clone(),
})
}
fn release_map(&mut self, map: &mut Self::Mem) {
self.internal.borrow_mut().release_map(&mut map.internal)
}
fn clone_ref(&mut self, mapping: &Self::Mem) -> Result<Self::Mem, Error> {
Ok(Self::Mem {
internal: ManuallyDrop::new(self.internal.borrow_mut().clone_ref(&mapping.internal)?),
provider: self.internal.clone(),
})
}
fn post_fork(&mut self) {
self.internal.borrow_mut().post_fork()
}
}
impl<T> Default for RcShMemProvider<T>
where
T: ShMemProvider + alloc::fmt::Debug,
{
fn default() -> Self {
Self::new().unwrap()
}
} }
#[cfg(all(unix, feature = "std"))] #[cfg(all(unix, feature = "std"))]
@ -320,7 +424,7 @@ pub mod unix_shmem {
#[cfg(unix)] #[cfg(unix)]
impl Default for CommonUnixShMemProvider { impl Default for CommonUnixShMemProvider {
fn default() -> Self { fn default() -> Self {
Self::new() Self::new().unwrap()
} }
} }
@ -329,8 +433,8 @@ pub mod unix_shmem {
impl ShMemProvider for CommonUnixShMemProvider { impl ShMemProvider for CommonUnixShMemProvider {
type Mem = CommonUnixShMem; type Mem = CommonUnixShMem;
fn new() -> Self { fn new() -> Result<Self, Error> {
Self {} Ok(Self {})
} }
fn new_map(&mut self, map_size: usize) -> Result<Self::Mem, Error> { fn new_map(&mut self, map_size: usize) -> Result<Self::Mem, Error> {
CommonUnixShMem::new(map_size) CommonUnixShMem::new(map_size)
@ -507,7 +611,6 @@ pub mod unix_shmem {
impl Drop for AshmemShMem { impl Drop for AshmemShMem {
fn drop(&mut self) { fn drop(&mut self) {
unsafe { unsafe {
//let fd = Self::fd_from_id(self.id).unwrap();
let fd: i32 = self.id.to_string().parse().unwrap(); let fd: i32 = self.id.to_string().parse().unwrap();
let length = ioctl(fd, ASHMEM_GET_SIZE); let length = ioctl(fd, ASHMEM_GET_SIZE);
@ -533,7 +636,7 @@ pub mod unix_shmem {
#[cfg(unix)] #[cfg(unix)]
impl Default for AshmemShMemProvider { impl Default for AshmemShMemProvider {
fn default() -> Self { fn default() -> Self {
Self::new() Self::new().unwrap()
} }
} }
@ -542,8 +645,8 @@ pub mod unix_shmem {
impl ShMemProvider for AshmemShMemProvider { impl ShMemProvider for AshmemShMemProvider {
type Mem = AshmemShMem; type Mem = AshmemShMem;
fn new() -> Self { fn new() -> Result<Self, Error> {
Self {} Ok(Self {})
} }
fn new_map(&mut self, map_size: usize) -> Result<Self::Mem, Error> { fn new_map(&mut self, map_size: usize) -> Result<Self::Mem, Error> {
@ -693,7 +796,7 @@ pub mod win32_shmem {
impl Default for Win32ShMemProvider { impl Default for Win32ShMemProvider {
fn default() -> Self { fn default() -> Self {
Self::new() Self::new().unwrap()
} }
} }
@ -701,8 +804,8 @@ pub mod win32_shmem {
impl ShMemProvider for Win32ShMemProvider { impl ShMemProvider for Win32ShMemProvider {
type Mem = Win32ShMem; type Mem = Win32ShMem;
fn new() -> Self { fn new() -> Result<Self, Error> {
Self {} Ok(Self {})
} }
fn new_map(&mut self, map_size: usize) -> Result<Self::Mem, Error> { fn new_map(&mut self, map_size: usize) -> Result<Self::Mem, Error> {
Win32ShMem::new_map(map_size) Win32ShMem::new_map(map_size)

View File

@ -1,7 +1,7 @@
//! LLMP-backed event manager for scalable multi-processed fuzzing //! LLMP-backed event manager for scalable multi-processed fuzzing
use alloc::{rc::Rc, string::ToString, vec::Vec}; use alloc::{string::ToString, vec::Vec};
use core::{cell::RefCell, marker::PhantomData, time::Duration}; use core::{marker::PhantomData, time::Duration};
use serde::{de::DeserializeOwned, Serialize}; use serde::{de::DeserializeOwned, Serialize};
#[cfg(feature = "std")] #[cfg(feature = "std")]
@ -12,6 +12,7 @@ use crate::bolts::{
llmp::{LlmpClient, LlmpReceiver}, llmp::{LlmpClient, LlmpReceiver},
shmem::StdShMemProvider, shmem::StdShMemProvider,
}; };
use crate::{ use crate::{
bolts::{ bolts::{
llmp::{self, LlmpClientDescription, LlmpSender, Tag}, llmp::{self, LlmpClientDescription, LlmpSender, Tag},
@ -32,10 +33,7 @@ use crate::{
use crate::utils::startable_self; use crate::utils::startable_self;
#[cfg(all(feature = "std", unix))] #[cfg(all(feature = "std", unix))]
use crate::{ use crate::utils::{fork, ForkResult};
bolts::shmem::UnixShMemProvider,
utils::{fork, ForkResult},
};
#[cfg(all(feature = "std", target_os = "android"))] #[cfg(all(feature = "std", target_os = "android"))]
use crate::bolts::os::ashmem_server::AshmemService; use crate::bolts::os::ashmem_server::AshmemService;
@ -65,41 +63,6 @@ where
phantom: PhantomData<(I, S)>, phantom: PhantomData<(I, S)>,
} }
#[cfg(feature = "std")]
#[cfg(unix)]
impl<I, S, ST> LlmpEventManager<I, S, UnixShMemProvider, ST>
where
I: Input,
S: IfInteresting<I>,
ST: Stats,
{
/// Create llmp on a port
/// If the port is not yet bound, it will act as broker
/// Else, it will act as client.
#[cfg(feature = "std")]
pub fn new_on_port_std(
shmem_provider: &Rc<RefCell<UnixShMemProvider>>,
stats: ST,
port: u16,
) -> Result<Self, Error> {
Ok(Self {
stats: Some(stats),
llmp: llmp::LlmpConnection::on_port(shmem_provider, port)?,
phantom: PhantomData,
})
}
/// If a client respawns, it may reuse the existing connection, previously stored by LlmpClient::to_env
/// Std uses UnixShMem.
#[cfg(feature = "std")]
pub fn existing_client_from_env_std(
shmem_provider: &Rc<RefCell<UnixShMemProvider>>,
env_name: &str,
) -> Result<Self, Error> {
Self::existing_client_from_env(shmem_provider, env_name)
}
}
impl<I, S, SP, ST> Drop for LlmpEventManager<I, S, SP, ST> impl<I, S, SP, ST> Drop for LlmpEventManager<I, S, SP, ST>
where where
I: Input, I: Input,
@ -124,11 +87,7 @@ where
/// If the port is not yet bound, it will act as broker /// If the port is not yet bound, it will act as broker
/// Else, it will act as client. /// Else, it will act as client.
#[cfg(feature = "std")] #[cfg(feature = "std")]
pub fn new_on_port( pub fn new_on_port(shmem_provider: SP, stats: ST, port: u16) -> Result<Self, Error> {
shmem_provider: &Rc<RefCell<SP>>,
stats: ST,
port: u16,
) -> Result<Self, Error> {
Ok(Self { Ok(Self {
stats: Some(stats), stats: Some(stats),
llmp: llmp::LlmpConnection::on_port(shmem_provider, port)?, llmp: llmp::LlmpConnection::on_port(shmem_provider, port)?,
@ -138,10 +97,7 @@ where
/// If a client respawns, it may reuse the existing connection, previously stored by LlmpClient::to_env /// If a client respawns, it may reuse the existing connection, previously stored by LlmpClient::to_env
#[cfg(feature = "std")] #[cfg(feature = "std")]
pub fn existing_client_from_env( pub fn existing_client_from_env(shmem_provider: SP, env_name: &str) -> Result<Self, Error> {
shmem_provider: &Rc<RefCell<SP>>,
env_name: &str,
) -> Result<Self, Error> {
Ok(Self { Ok(Self {
stats: None, stats: None,
llmp: llmp::LlmpConnection::IsClient { llmp: llmp::LlmpConnection::IsClient {
@ -160,7 +116,7 @@ where
/// Create an existing client from description /// Create an existing client from description
pub fn existing_client_from_description( pub fn existing_client_from_description(
shmem_provider: &Rc<RefCell<SP>>, shmem_provider: SP,
description: &LlmpClientDescription, description: &LlmpClientDescription,
) -> Result<Self, Error> { ) -> Result<Self, Error> {
Ok(Self { Ok(Self {
@ -304,7 +260,7 @@ where
let observers: OT = postcard::from_bytes(&observers_buf)?; let observers: OT = postcard::from_bytes(&observers_buf)?;
// TODO include ExitKind in NewTestcase // TODO include ExitKind in NewTestcase
let fitness = state.is_interesting(&input, &observers, ExitKind::Ok)?; let fitness = state.is_interesting(&input, &observers, &ExitKind::Ok)?;
if fitness > 0 if fitness > 0
&& state && state
.add_if_interesting(&input, fitness, scheduler)? .add_if_interesting(&input, fitness, scheduler)?
@ -400,7 +356,7 @@ where
/// Deserialize the state and corpus tuple, previously serialized with `serialize_state_corpus(...)` /// Deserialize the state and corpus tuple, previously serialized with `serialize_state_corpus(...)`
#[allow(clippy::type_complexity)] #[allow(clippy::type_complexity)]
pub fn deserialize_state_mgr<I, S, SP, ST>( pub fn deserialize_state_mgr<I, S, SP, ST>(
shmem_provider: &Rc<RefCell<SP>>, shmem_provider: SP,
state_corpus_serialized: &[u8], state_corpus_serialized: &[u8],
) -> Result<(S, LlmpEventManager<I, S, SP, ST>), Error> ) -> Result<(S, LlmpEventManager<I, S, SP, ST>), Error>
where where
@ -525,7 +481,7 @@ where
#[cfg(target_os = "android")] #[cfg(target_os = "android")]
AshmemService::start().expect("Error starting Ashmem Service"); AshmemService::start().expect("Error starting Ashmem Service");
setup_restarting_mgr(StdShMemProvider::new(), stats, broker_port) setup_restarting_mgr(StdShMemProvider::new()?, stats, broker_port)
} }
/// A restarting state is a combination of restarter and runner, that can be used on systems without `fork`. /// A restarting state is a combination of restarter and runner, that can be used on systems without `fork`.
@ -537,7 +493,7 @@ where
clippy::similar_names clippy::similar_names
)] // for { mgr = LlmpEventManager... } )] // for { mgr = LlmpEventManager... }
pub fn setup_restarting_mgr<I, S, SP, ST>( pub fn setup_restarting_mgr<I, S, SP, ST>(
shmem_provider: SP, mut shmem_provider: SP,
//mgr: &mut LlmpEventManager<I, S, SH, ST>, //mgr: &mut LlmpEventManager<I, S, SH, ST>,
stats: ST, stats: ST,
broker_port: u16, broker_port: u16,
@ -548,13 +504,13 @@ where
SP: ShMemProvider, SP: ShMemProvider,
ST: Stats, ST: Stats,
{ {
let shmem_provider = Rc::new(RefCell::new(shmem_provider));
let mut mgr = let mut mgr =
LlmpEventManager::<I, S, SP, ST>::new_on_port(&shmem_provider, stats, broker_port)?; LlmpEventManager::<I, S, SP, ST>::new_on_port(shmem_provider.clone(), stats, broker_port)?;
// We start ourself as child process to actually fuzz // We start ourself as child process to actually fuzz
let (sender, mut receiver, shmem_provider) = if std::env::var(_ENV_FUZZER_SENDER).is_err() { let (sender, mut receiver, mut new_shmem_provider) = if std::env::var(_ENV_FUZZER_SENDER)
.is_err()
{
if mgr.is_broker() { if mgr.is_broker() {
// Yep, broker. Just loop here. // Yep, broker. Just loop here.
println!("Doing broker things. Run this tool again to start fuzzing in a client."); println!("Doing broker things. Run this tool again to start fuzzing in a client.");
@ -566,13 +522,9 @@ where
mgr.to_env(_ENV_FUZZER_BROKER_CLIENT_INITIAL); mgr.to_env(_ENV_FUZZER_BROKER_CLIENT_INITIAL);
// First, create a channel from the fuzzer (sender) to us (receiver) to report its state for restarts. // First, create a channel from the fuzzer (sender) to us (receiver) to report its state for restarts.
let sender = { LlmpSender::new(&shmem_provider, 0, false)? }; let sender = { LlmpSender::new(shmem_provider.clone(), 0, false)? };
let map = { let map = { shmem_provider.clone_ref(&sender.out_maps.last().unwrap().shmem)? };
shmem_provider
.borrow_mut()
.clone_ref(&sender.out_maps.last().unwrap().shmem)?
};
let receiver = LlmpReceiver::on_existing_map(shmem_provider.clone(), map, None)?; let receiver = LlmpReceiver::on_existing_map(shmem_provider.clone(), map, None)?;
// Store the information to a map. // Store the information to a map.
sender.to_env(_ENV_FUZZER_SENDER)?; sender.to_env(_ENV_FUZZER_SENDER)?;
@ -613,14 +565,16 @@ where
// A sender and a receiver for single communication // A sender and a receiver for single communication
// Clone so we get a new connection to the AshmemServer if we are using // Clone so we get a new connection to the AshmemServer if we are using
// ServedShMemProvider // ServedShMemProvider
let shmem_provider = Rc::new(RefCell::new(shmem_provider.borrow_mut().clone())); shmem_provider.post_fork();
( (
LlmpSender::on_existing_from_env(&shmem_provider, _ENV_FUZZER_SENDER)?, LlmpSender::on_existing_from_env(shmem_provider.clone(), _ENV_FUZZER_SENDER)?,
LlmpReceiver::on_existing_from_env(&shmem_provider, _ENV_FUZZER_RECEIVER)?, LlmpReceiver::on_existing_from_env(shmem_provider.clone(), _ENV_FUZZER_RECEIVER)?,
shmem_provider, shmem_provider,
) )
}; };
new_shmem_provider.post_fork();
println!("We're a client, let's fuzz :)"); println!("We're a client, let's fuzz :)");
for (var, val) in std::env::vars() { for (var, val) in std::env::vars() {
@ -633,7 +587,7 @@ where
println!("First run. Let's set it all up"); println!("First run. Let's set it all up");
// Mgr to send and receive msgs from/to all other fuzzer instances // Mgr to send and receive msgs from/to all other fuzzer instances
let client_mgr = LlmpEventManager::<I, S, SP, ST>::existing_client_from_env( let client_mgr = LlmpEventManager::<I, S, SP, ST>::existing_client_from_env(
&shmem_provider, new_shmem_provider,
_ENV_FUZZER_BROKER_CLIENT_INITIAL, _ENV_FUZZER_BROKER_CLIENT_INITIAL,
)?; )?;
@ -643,7 +597,7 @@ where
Some((_sender, _tag, msg)) => { Some((_sender, _tag, msg)) => {
println!("Subsequent run. Let's load all data from shmem (received {} bytes from previous instance)", msg.len()); println!("Subsequent run. Let's load all data from shmem (received {} bytes from previous instance)", msg.len());
let (state, mgr): (S, LlmpEventManager<I, S, SP, ST>) = let (state, mgr): (S, LlmpEventManager<I, S, SP, ST>) =
deserialize_state_mgr(&shmem_provider, &msg)?; deserialize_state_mgr(new_shmem_provider, &msg)?;
(Some(state), LlmpRestartingEventManager::new(mgr, sender)) (Some(state), LlmpRestartingEventManager::new(mgr, sender))
} }

View File

@ -235,7 +235,7 @@ where
mod unix_signal_handler { mod unix_signal_handler {
use alloc::vec::Vec; use alloc::vec::Vec;
use core::ptr; use core::ptr;
use libc::{c_void, siginfo_t}; use libc::{c_void, siginfo_t, ucontext_t};
#[cfg(feature = "std")] #[cfg(feature = "std")]
use std::io::{stdout, Write}; use std::io::{stdout, Write};
@ -273,8 +273,8 @@ mod unix_signal_handler {
pub event_mgr_ptr: *mut c_void, pub event_mgr_ptr: *mut c_void,
pub observers_ptr: *const c_void, pub observers_ptr: *const c_void,
pub current_input_ptr: *const c_void, pub current_input_ptr: *const c_void,
pub crash_handler: unsafe fn(Signal, siginfo_t, *const c_void, data: &mut Self), pub crash_handler: unsafe fn(Signal, siginfo_t, &mut ucontext_t, data: &mut Self),
pub timeout_handler: unsafe fn(Signal, siginfo_t, *const c_void, data: &mut Self), pub timeout_handler: unsafe fn(Signal, siginfo_t, &mut ucontext_t, data: &mut Self),
} }
unsafe impl Send for InProcessExecutorHandlerData {} unsafe impl Send for InProcessExecutorHandlerData {}
@ -283,21 +283,21 @@ mod unix_signal_handler {
unsafe fn nop_handler( unsafe fn nop_handler(
_signal: Signal, _signal: Signal,
_info: siginfo_t, _info: siginfo_t,
_void: *const c_void, _context: &mut ucontext_t,
_data: &mut InProcessExecutorHandlerData, _data: &mut InProcessExecutorHandlerData,
) { ) {
} }
#[cfg(unix)] #[cfg(unix)]
impl Handler for InProcessExecutorHandlerData { impl Handler for InProcessExecutorHandlerData {
fn handle(&mut self, signal: Signal, info: siginfo_t, void: *const c_void) { fn handle(&mut self, signal: Signal, info: siginfo_t, context: &mut ucontext_t) {
unsafe { unsafe {
let data = &mut GLOBAL_STATE; let data = &mut GLOBAL_STATE;
match signal { match signal {
Signal::SigUser2 | Signal::SigAlarm => { Signal::SigUser2 | Signal::SigAlarm => {
(data.timeout_handler)(signal, info, void, data) (data.timeout_handler)(signal, info, context, data)
} }
_ => (data.crash_handler)(signal, info, void, data), _ => (data.crash_handler)(signal, info, context, data),
} }
} }
} }
@ -312,6 +312,7 @@ mod unix_signal_handler {
Signal::SigFloatingPointException, Signal::SigFloatingPointException,
Signal::SigIllegalInstruction, Signal::SigIllegalInstruction,
Signal::SigSegmentationFault, Signal::SigSegmentationFault,
Signal::SigTrap,
] ]
} }
} }
@ -320,7 +321,7 @@ mod unix_signal_handler {
pub unsafe fn inproc_timeout_handler<EM, I, OC, OFT, OT, S>( pub unsafe fn inproc_timeout_handler<EM, I, OC, OFT, OT, S>(
_signal: Signal, _signal: Signal,
_info: siginfo_t, _info: siginfo_t,
_void: *const c_void, _context: &mut ucontext_t,
data: &mut InProcessExecutorHandlerData, data: &mut InProcessExecutorHandlerData,
) where ) where
EM: EventManager<I, S>, EM: EventManager<I, S>,
@ -348,7 +349,7 @@ mod unix_signal_handler {
let obj_fitness = state let obj_fitness = state
.objectives_mut() .objectives_mut()
.is_interesting_all(&input, observers, ExitKind::Timeout) .is_interesting_all(&input, observers, &ExitKind::Timeout)
.expect("In timeout handler objectives failure."); .expect("In timeout handler objectives failure.");
if obj_fitness > 0 { if obj_fitness > 0 {
state state
@ -382,7 +383,7 @@ mod unix_signal_handler {
pub unsafe fn inproc_crash_handler<EM, I, OC, OFT, OT, S>( pub unsafe fn inproc_crash_handler<EM, I, OC, OFT, OT, S>(
_signal: Signal, _signal: Signal,
_info: siginfo_t, _info: siginfo_t,
_void: *const c_void, _context: &mut ucontext_t,
data: &mut InProcessExecutorHandlerData, data: &mut InProcessExecutorHandlerData,
) where ) where
EM: EventManager<I, S>, EM: EventManager<I, S>,
@ -392,6 +393,10 @@ mod unix_signal_handler {
S: HasObjectives<OFT, I> + HasSolutions<OC, I>, S: HasObjectives<OFT, I> + HasSolutions<OC, I>,
I: Input + HasTargetBytes, I: Input + HasTargetBytes,
{ {
#[cfg(all(target_os = "android", target_arch = "aarch64"))]
let _context = *(((_context as *mut _ as *mut c_void as usize) + 128) as *mut c_void
as *mut ucontext_t);
#[cfg(feature = "std")] #[cfg(feature = "std")]
println!("Crashed with {}", _signal); println!("Crashed with {}", _signal);
if !data.current_input_ptr.is_null() { if !data.current_input_ptr.is_null() {
@ -401,6 +406,45 @@ mod unix_signal_handler {
#[cfg(feature = "std")] #[cfg(feature = "std")]
println!("Child crashed!"); println!("Child crashed!");
#[cfg(all(
feature = "std",
any(target_os = "linux", target_os = "android"),
target_arch = "aarch64"
))]
{
use crate::utils::find_mapping_for_address;
println!("{:━^100}", " CRASH ");
println!(
"Received signal {} at 0x{:016x}, fault address: 0x{:016x}",
_signal, _context.uc_mcontext.pc, _context.uc_mcontext.fault_address
);
if let Ok((start, _, _, path)) =
find_mapping_for_address(_context.uc_mcontext.pc as usize)
{
println!(
"pc is at offset 0x{:08x} in {}",
_context.uc_mcontext.pc as usize - start,
path
);
}
println!("{:━^100}", " REGISTERS ");
for reg in 0..31 {
print!(
"x{:02}: 0x{:016x} ",
reg, _context.uc_mcontext.regs[reg as usize]
);
if reg % 4 == 3 {
println!();
}
}
println!("pc : 0x{:016x} ", _context.uc_mcontext.pc);
//println!("{:━^100}", " BACKTRACE ");
//println!("{:?}", backtrace::Backtrace::new())
}
#[cfg(feature = "std")] #[cfg(feature = "std")]
let _ = stdout().flush(); let _ = stdout().flush();
@ -410,7 +454,7 @@ mod unix_signal_handler {
let obj_fitness = state let obj_fitness = state
.objectives_mut() .objectives_mut()
.is_interesting_all(&input, observers, ExitKind::Crash) .is_interesting_all(&input, observers, &ExitKind::Crash)
.expect("In crash handler objectives failure."); .expect("In crash handler objectives failure.");
if obj_fitness > 0 { if obj_fitness > 0 {
let new_input = input.clone(); let new_input = input.clone();
@ -574,7 +618,7 @@ mod windows_exception_handler {
let obj_fitness = state let obj_fitness = state
.objectives_mut() .objectives_mut()
.is_interesting_all(&input, observers, ExitKind::Crash) .is_interesting_all(&input, observers, &ExitKind::Crash)
.expect("In crash handler objectives failure."); .expect("In crash handler objectives failure.");
if obj_fitness > 0 { if obj_fitness > 0 {
let new_input = input.clone(); let new_input = input.clone();

View File

@ -5,24 +5,28 @@ pub use inprocess::InProcessExecutor;
pub mod timeout; pub mod timeout;
pub use timeout::TimeoutExecutor; pub use timeout::TimeoutExecutor;
use core::cmp::PartialEq;
use core::marker::PhantomData; use core::marker::PhantomData;
use crate::{ use crate::{
bolts::tuples::Named, bolts::{serdeany::SerdeAny, tuples::Named},
events::EventManager, events::EventManager,
inputs::{HasTargetBytes, Input}, inputs::{HasTargetBytes, Input},
observers::ObserversTuple, observers::ObserversTuple,
Error, Error,
}; };
use alloc::boxed::Box;
pub trait CustomExitKind: core::fmt::Debug + SerdeAny + 'static {}
/// How an execution finished. /// How an execution finished.
#[derive(Debug, Clone, PartialEq)] #[derive(Debug)]
pub enum ExitKind { pub enum ExitKind {
Ok, Ok,
Crash, Crash,
Oom, Oom,
Timeout, Timeout,
Custom(Box<dyn CustomExitKind>),
} }
pub trait HasObservers<OT> pub trait HasObservers<OT>

View File

@ -91,6 +91,10 @@ where
phantom: PhantomData, phantom: PhantomData,
} }
} }
pub fn inner(&mut self) -> &mut E {
&mut self.executor
}
} }
impl<E, I, OT> Executor<I> for TimeoutExecutor<E, I, OT> impl<E, I, OT> Executor<I> for TimeoutExecutor<E, I, OT>
@ -126,6 +130,11 @@ where
null_mut(), null_mut(),
); );
} }
#[cfg(windows)]
{
// TODO
let _ = self.exec_tmout.as_millis();
}
self.executor.pre_exec(_state, _event_mgr, _input) self.executor.pre_exec(_state, _event_mgr, _input)
} }
@ -155,6 +164,10 @@ where
null_mut(), null_mut(),
); );
} }
#[cfg(windows)]
{
// TODO
}
self.executor.post_exec(_state, _event_mgr, _input) self.executor.post_exec(_state, _event_mgr, _input)
} }

View File

@ -154,7 +154,7 @@ where
&mut self, &mut self,
_input: &I, _input: &I,
observers: &OT, observers: &OT,
_exit_kind: ExitKind, _exit_kind: &ExitKind,
) -> Result<u32, Error> { ) -> Result<u32, Error> {
let mut interesting = 0; let mut interesting = 0;
// TODO optimize // TODO optimize

View File

@ -29,7 +29,7 @@ where
&mut self, &mut self,
input: &I, input: &I,
observers: &OT, observers: &OT,
exit_kind: ExitKind, exit_kind: &ExitKind,
) -> Result<u32, Error>; ) -> Result<u32, Error>;
/// Append to the testcase the generated metadata in case of a new corpus item /// Append to the testcase the generated metadata in case of a new corpus item
@ -54,7 +54,7 @@ where
&mut self, &mut self,
input: &I, input: &I,
observers: &OT, observers: &OT,
exit_kind: ExitKind, exit_kind: &ExitKind,
) -> Result<u32, Error>; ) -> Result<u32, Error>;
/// Write metadata for this testcase /// Write metadata for this testcase
@ -73,7 +73,7 @@ where
&mut self, &mut self,
_: &I, _: &I,
_: &OT, _: &OT,
_: ExitKind, _: &ExitKind,
) -> Result<u32, Error> { ) -> Result<u32, Error> {
Ok(0) Ok(0)
} }
@ -99,9 +99,9 @@ where
&mut self, &mut self,
input: &I, input: &I,
observers: &OT, observers: &OT,
exit_kind: ExitKind, exit_kind: &ExitKind,
) -> Result<u32, Error> { ) -> Result<u32, Error> {
Ok(self.0.is_interesting(input, observers, exit_kind.clone())? Ok(self.0.is_interesting(input, observers, exit_kind)?
+ self.1.is_interesting_all(input, observers, exit_kind)?) + self.1.is_interesting_all(input, observers, exit_kind)?)
} }
@ -128,9 +128,9 @@ where
&mut self, &mut self,
_input: &I, _input: &I,
_observers: &OT, _observers: &OT,
exit_kind: ExitKind, exit_kind: &ExitKind,
) -> Result<u32, Error> { ) -> Result<u32, Error> {
if exit_kind == ExitKind::Crash { if let ExitKind::Crash = exit_kind {
Ok(1) Ok(1)
} else { } else {
Ok(0) Ok(0)
@ -168,9 +168,9 @@ where
&mut self, &mut self,
_input: &I, _input: &I,
_observers: &OT, _observers: &OT,
exit_kind: ExitKind, exit_kind: &ExitKind,
) -> Result<u32, Error> { ) -> Result<u32, Error> {
if exit_kind == ExitKind::Timeout { if let ExitKind::Timeout = exit_kind {
Ok(1) Ok(1)
} else { } else {
Ok(0) Ok(0)
@ -211,7 +211,7 @@ where
&mut self, &mut self,
_input: &I, _input: &I,
observers: &OT, observers: &OT,
_exit_kind: ExitKind, _exit_kind: &ExitKind,
) -> Result<u32, Error> { ) -> Result<u32, Error> {
let observer = observers.match_first_type::<TimeObserver>().unwrap(); let observer = observers.match_first_type::<TimeObserver>().unwrap();
self.exec_time = *observer.last_runtime(); self.exec_time = *observer.last_runtime();

View File

@ -169,7 +169,7 @@ where
&mut self, &mut self,
input: &I, input: &I,
observers: &OT, observers: &OT,
exit_kind: ExitKind, exit_kind: &ExitKind,
) -> Result<u32, Error> ) -> Result<u32, Error>
where where
OT: ObserversTuple; OT: ObserversTuple;
@ -448,7 +448,7 @@ where
&mut self, &mut self,
input: &I, input: &I,
observers: &OT, observers: &OT,
exit_kind: ExitKind, exit_kind: &ExitKind,
) -> Result<u32, Error> ) -> Result<u32, Error>
where where
OT: ObserversTuple, OT: ObserversTuple,
@ -654,13 +654,13 @@ where
executor.post_exec_observers()?; executor.post_exec_observers()?;
let observers = executor.observers(); let observers = executor.observers();
let fitness = let fitness = self
self.feedbacks_mut() .feedbacks_mut()
.is_interesting_all(&input, observers, exit_kind.clone())?; .is_interesting_all(&input, observers, &exit_kind)?;
let is_solution = self let is_solution = self
.objectives_mut() .objectives_mut()
.is_interesting_all(&input, observers, exit_kind)? .is_interesting_all(&input, observers, &exit_kind)?
> 0; > 0;
Ok((fitness, is_solution)) Ok((fitness, is_solution))
} }

View File

@ -438,6 +438,77 @@ pub fn startable_self() -> Result<Command, Error> {
Ok(startable) Ok(startable)
} }
/// Allows one to walk the mappings in /proc/self/maps, caling a callback function for each
/// mapping.
/// If the callback returns true, we stop the walk.
#[cfg(all(feature = "std", any(target_os = "linux", target_os = "android")))]
pub fn walk_self_maps(visitor: &mut dyn FnMut(usize, usize, String, String) -> bool) {
use regex::Regex;
use std::{
fs::File,
io::{BufRead, BufReader},
};
let re = Regex::new(r"^(?P<start>[0-9a-f]{8,16})-(?P<end>[0-9a-f]{8,16}) (?P<perm>[-rwxp]{4}) (?P<offset>[0-9a-f]{8}) [0-9a-f]+:[0-9a-f]+ [0-9]+\s+(?P<path>.*)$")
.unwrap();
let mapsfile = File::open("/proc/self/maps").expect("Unable to open /proc/self/maps");
for line in BufReader::new(mapsfile).lines() {
let line = line.unwrap();
if let Some(caps) = re.captures(&line) {
if visitor(
usize::from_str_radix(caps.name("start").unwrap().as_str(), 16).unwrap(),
usize::from_str_radix(caps.name("end").unwrap().as_str(), 16).unwrap(),
caps.name("perm").unwrap().as_str().to_string(),
caps.name("path").unwrap().as_str().to_string(),
) {
break;
};
}
}
}
/// Get the start and end address, permissions and path of the mapping containing a particular address
#[cfg(all(feature = "std", any(target_os = "linux", target_os = "android")))]
pub fn find_mapping_for_address(address: usize) -> Result<(usize, usize, String, String), Error> {
let mut result = (0, 0, "".to_string(), "".to_string());
walk_self_maps(&mut |start, end, permissions, path| {
if start <= address && address < end {
result = (start, end, permissions, path);
true
} else {
false
}
});
if result.0 != 0 {
Ok(result)
} else {
Err(Error::Unknown(
"Couldn't find a mapping for this address".to_string(),
))
}
}
/// Get the start and end address of the mapping containing with a particular path
#[cfg(all(feature = "std", any(target_os = "linux", target_os = "android")))]
pub fn find_mapping_for_path(libpath: &str) -> (usize, usize) {
let mut libstart = 0;
let mut libend = 0;
walk_self_maps(&mut |start, end, _permissions, path| {
if libpath == path {
if libstart == 0 {
libstart = start;
}
libend = end;
}
false
});
(libstart, libend)
}
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
//use xxhash_rust::xxh3::xxh3_64_with_seed; //use xxhash_rust::xxh3::xxh3_64_with_seed;

33
libafl_frida/Cargo.toml Normal file
View File

@ -0,0 +1,33 @@
[package]
name = "libafl_frida"
version = "0.1.0"
authors = ["s1341 <github@shmarya.net>"]
edition = "2018"
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
[build-dependencies]
cc = { version = "1.0", features = ["parallel"] }
[dependencies]
libafl = { path = "../libafl", version = "0.1.0", features = ["std", "libafl_derive"] }
libafl_targets = { path = "../libafl_targets", version = "0.1.0" }
nix = "0.20.0"
libc = "0.2.92"
hashbrown = "0.11"
libloading = "0.7.0"
rangemap = "0.1.10"
frida-gum = { version = "0.4.0", features = [ "auto-download", "backtrace", "event-sink", "invocation-listener"] }
frida-gum-sys = { version = "0.2.4", features = [ "auto-download", "event-sink", "invocation-listener"] }
regex = "1.4"
dynasmrt = "1.0.1"
capstone = "0.8.0"
color-backtrace ={ version = "0.5", features = [ "resolve-modules" ] }
termcolor = "1.1.2"
serde = "1.0"
backtrace = { version = "0.3.58", default-features = false, features = ["std", "serde"] }
num-traits = "0.2.14"
seahash = "4.1.0"
[target.'cfg(unix)'.dependencies]
gothook = { version = "0.1" }

5
libafl_frida/build.rs Normal file
View File

@ -0,0 +1,5 @@
// build.rs
fn main() {
cc::Build::new().file("src/gettls.c").compile("libgettls.a");
}

1880
libafl_frida/src/asan_rt.rs Normal file

File diff suppressed because it is too large Load Diff

View File

@ -0,0 +1,9 @@
#ifdef _MSC_VER
__declspec( thread ) int i = 0;
#else
__thread int i = 0;
#endif
void * get_tls_ptr() {
return (void*)&i;
}

617
libafl_frida/src/helper.rs Normal file
View File

@ -0,0 +1,617 @@
use libafl::inputs::{HasTargetBytes, Input};
#[cfg(any(target_os = "linux", target_os = "android"))]
use libafl::utils::find_mapping_for_path;
use libafl_targets::drcov::{DrCovBasicBlock, DrCovWriter};
#[cfg(target_arch = "aarch64")]
use capstone::arch::{arm64::Arm64OperandType, ArchOperand::Arm64Operand};
use capstone::{
arch::{self, BuildsCapstone},
Capstone, Insn,
};
use core::cell::RefCell;
#[cfg(target_arch = "x86_64")]
use frida_gum::instruction_writer::X86Register;
#[cfg(target_arch = "aarch64")]
use frida_gum::instruction_writer::{Aarch64Register, IndexMode};
use frida_gum::{
instruction_writer::InstructionWriter,
stalker::{StalkerOutput, Transformer},
CpuContext,
};
use frida_gum::{Gum, Module, PageProtection};
use num_traits::cast::FromPrimitive;
use rangemap::RangeMap;
use std::rc::Rc;
use crate::{asan_rt::AsanRuntime, FridaOptions};
/// An helper that feeds FridaInProcessExecutor with user-supplied instrumentation
pub trait FridaHelper<'a> {
fn transformer(&self) -> &Transformer<'a>;
fn register_thread(&self);
fn pre_exec<I: Input + HasTargetBytes>(&mut self, input: &I);
fn post_exec<I: Input + HasTargetBytes>(&mut self, input: &I);
fn stalker_enabled(&self) -> bool;
fn map_ptr(&mut self) -> *mut u8;
}
pub const MAP_SIZE: usize = 64 * 1024;
/// An helper that feeds FridaInProcessExecutor with edge-coverage instrumentation
pub struct FridaInstrumentationHelper<'a> {
map: [u8; MAP_SIZE],
previous_pc: [u64; 1],
current_log_impl: u64,
/// Transformer that has to be passed to FridaInProcessExecutor
transformer: Option<Transformer<'a>>,
capstone: Capstone,
asan_runtime: Rc<RefCell<AsanRuntime>>,
ranges: RangeMap<usize, (u16, &'a str)>,
options: FridaOptions,
drcov_basic_blocks: Vec<DrCovBasicBlock>,
}
impl<'a> FridaHelper<'a> for FridaInstrumentationHelper<'a> {
fn transformer(&self) -> &Transformer<'a> {
self.transformer.as_ref().unwrap()
}
/// Register the current thread with the FridaInstrumentationHelper
fn register_thread(&self) {
self.asan_runtime.borrow().register_thread();
}
fn pre_exec<I: Input + HasTargetBytes>(&mut self, input: &I) {
let target_bytes = input.target_bytes();
let slice = target_bytes.as_slice();
//println!("target_bytes: {:02x?}", slice);
self.asan_runtime
.borrow()
.unpoison(slice.as_ptr() as usize, slice.len());
}
fn post_exec<I: Input + HasTargetBytes>(&mut self, input: &I) {
if self.options.drcov_enabled() {
let filename = format!(
"./coverage/{:016x}.drcov",
seahash::hash(input.target_bytes().as_slice())
);
DrCovWriter::new(&filename, &self.ranges, &mut self.drcov_basic_blocks).write();
}
if self.options.asan_enabled() {
if self.options.asan_detect_leaks() {
self.asan_runtime.borrow_mut().check_for_leaks();
}
self.asan_runtime.borrow_mut().reset_allocations();
}
}
fn stalker_enabled(&self) -> bool {
self.options.stalker_enabled()
}
fn map_ptr(&mut self) -> *mut u8 {
self.map.as_mut_ptr()
}
}
/// Helper function to get the size of a module's CODE section from frida
pub fn get_module_size(module_name: &str) -> usize {
let mut code_size = 0;
let code_size_ref = &mut code_size;
Module::enumerate_ranges(module_name, PageProtection::ReadExecute, move |details| {
*code_size_ref = details.memory_range().size() as usize;
true
});
code_size
}
/// A minimal maybe_log implementation. We insert this into the transformed instruction stream
/// every time we need a copy that is within a direct branch of the start of the transformed basic
/// block.
#[cfg(target_arch = "x86_64")]
const MAYBE_LOG_CODE: [u8; 47] = [
0x9c, /* pushfq */
0x50, /* push rax */
0x51, /* push rcx */
0x52, /* push rdx */
0x48, 0x8d, 0x05, 0x24, 0x00, 0x00, 0x00, /* lea rax, sym._afl_area_ptr_ptr */
0x48, 0x8b, 0x00, /* mov rax, qword [rax] */
0x48, 0x8d, 0x0d, 0x22, 0x00, 0x00, 0x00, /* lea rcx, sym.previous_pc */
0x48, 0x8b, 0x11, /* mov rdx, qword [rcx] */
0x48, 0x8b, 0x12, /* mov rdx, qword [rdx] */
0x48, 0x31, 0xfa, /* xor rdx, rdi */
0xfe, 0x04, 0x10, /* inc byte [rax + rdx] */
0x48, 0xd1, 0xef, /* shr rdi, 1 */
0x48, 0x8b, 0x01, /* mov rax, qword [rcx] */
0x48, 0x89, 0x38, /* mov qword [rax], rdi */
0x5a, /* pop rdx */
0x59, /* pop rcx */
0x58, /* pop rax */
0x9d, /* popfq */
0xc3, /* ret */
/* Read-only data goes here: */
/* uint8_t* afl_area_ptr */
/* uint64_t* afl_prev_loc_ptr */
];
#[cfg(target_arch = "aarch64")]
const MAYBE_LOG_CODE: [u8; 60] = [
// __afl_area_ptr[current_pc ^ previous_pc]++;
// previous_pc = current_pc >> 1;
0xE1, 0x0B, 0xBF, 0xA9, // stp x1, x2, [sp, -0x10]!
0xE3, 0x13, 0xBF, 0xA9, // stp x3, x4, [sp, -0x10]!
// x0 = current_pc
0xa1, 0x01, 0x00, 0x58, // ldr x1, #0x30, =__afl_area_ptr
0x82, 0x01, 0x00, 0x58, // ldr x2, #0x38, =&previous_pc
0x44, 0x00, 0x40, 0xf9, // ldr x4, [x2] (=previous_pc)
// __afl_area_ptr[current_pc ^ previous_pc]++;
0x84, 0x00, 0x00, 0xca, // eor x4, x4, x0
0x84, 0x3c, 0x40, 0x92, // and x4, x4, 0xffff (=MAP_SIZE - 1)
//0x20, 0x13, 0x20, 0xd4,
0x23, 0x68, 0x64, 0xf8, // ldr x3, [x1, x4]
0x63, 0x04, 0x00, 0x91, // add x3, x3, #1
0x23, 0x68, 0x24, 0xf8, // str x3, [x1, x4]
// previous_pc = current_pc >> 1;
0xe0, 0x07, 0x40, 0x8b, // add x0, xzr, x0, LSR #1
0x40, 0x00, 0x00, 0xf9, // str x0, [x2]
0xE3, 0x13, 0xc1, 0xA8, // ldp x3, x4, [sp], #0x10
0xE1, 0x0B, 0xc1, 0xA8, // ldp x1, x2, [sp], #0x10
0xC0, 0x03, 0x5F, 0xD6, // ret
// &afl_area_ptr
// &afl_prev_loc_ptr
];
#[cfg(target_arch = "aarch64")]
fn get_pc(context: &CpuContext) -> usize {
context.pc() as usize
}
#[cfg(target_arch = "x86_64")]
fn get_pc(context: &CpuContext) -> usize {
context.rip() as usize
}
/// The implementation of the FridaInstrumentationHelper
impl<'a> FridaInstrumentationHelper<'a> {
/// Constructor function to create a new FridaInstrumentationHelper, given a module_name.
pub fn new(
gum: &'a Gum,
options: FridaOptions,
_harness_module_name: &str,
modules_to_instrument: &'a Vec<&str>,
) -> Self {
let mut helper = Self {
map: [0u8; MAP_SIZE],
previous_pc: [0u64; 1],
current_log_impl: 0,
transformer: None,
capstone: Capstone::new()
.arm64()
.mode(arch::arm64::ArchMode::Arm)
.detail(true)
.build()
.expect("Failed to create Capstone object"),
asan_runtime: AsanRuntime::new(options),
ranges: RangeMap::new(),
options,
drcov_basic_blocks: vec![],
};
if options.stalker_enabled() {
for (id, module_name) in modules_to_instrument.iter().enumerate() {
let (lib_start, lib_end) = find_mapping_for_path(module_name);
println!("including range {:x}-{:x}", lib_start, lib_end);
helper
.ranges
.insert(lib_start..lib_end, (id as u16, module_name));
}
if helper.options.drcov_enabled() {
std::fs::create_dir_all("./coverage")
.expect("failed to create directory for coverage files");
}
let transformer = Transformer::from_callback(gum, |basic_block, output| {
let mut first = true;
for instruction in basic_block {
let instr = instruction.instr();
let address = instr.address();
//println!("address: {:x} contains: {:?}", address, helper.ranges.contains(&(address as usize)));
if helper.ranges.contains_key(&(address as usize)) {
if first {
first = false;
//println!("block @ {:x} transformed to {:x}", address, output.writer().pc());
if helper.options.coverage_enabled() {
helper.emit_coverage_mapping(address, &output);
}
if helper.options.drcov_enabled() {
instruction.put_callout(|context| {
let real_address = match helper
.asan_runtime
.borrow()
.real_address_for_stalked(get_pc(&context))
{
Some(address) => *address,
_ => get_pc(&context),
};
//let (range, (id, name)) = helper.ranges.get_key_value(&real_address).unwrap();
//println!("{}:0x{:016x}", name, real_address - range.start);
helper
.drcov_basic_blocks
.push(DrCovBasicBlock::new(real_address, real_address + 4));
})
}
}
if helper.options.asan_enabled() {
#[cfg(not(target_arch = "aarch64"))]
todo!("Implement ASAN for non-aarch64 targets");
#[cfg(target_arch = "aarch64")]
if let Ok((basereg, indexreg, displacement, width)) =
helper.is_interesting_instruction(address, instr)
{
helper.emit_shadow_check(
address,
&output,
basereg,
indexreg,
displacement,
width,
);
}
}
if helper.options.asan_enabled() || helper.options.drcov_enabled() {
helper.asan_runtime.borrow_mut().add_stalked_address(
output.writer().pc() as usize - 4,
address as usize,
);
}
}
instruction.keep()
}
});
helper.transformer = Some(transformer);
if helper.options.asan_enabled() || helper.options.drcov_enabled() {
helper.asan_runtime.borrow_mut().init(modules_to_instrument);
}
}
helper
}
#[cfg(target_arch = "aarch64")]
#[inline]
fn get_writer_register(&self, reg: capstone::RegId) -> Aarch64Register {
let regint: u16 = reg.0;
Aarch64Register::from_u32(regint as u32).unwrap()
}
#[cfg(target_arch = "aarch64")]
#[inline]
fn emit_shadow_check(
&self,
_address: u64,
output: &StalkerOutput,
basereg: capstone::RegId,
indexreg: capstone::RegId,
displacement: i32,
width: u32,
) {
let writer = output.writer();
let basereg = self.get_writer_register(basereg);
let indexreg = if indexreg.0 != 0 {
Some(self.get_writer_register(indexreg))
} else {
None
};
//writer.put_brk_imm(1);
// Preserve x0, x1:
writer.put_stp_reg_reg_reg_offset(
Aarch64Register::X0,
Aarch64Register::X1,
Aarch64Register::Sp,
-(16 + frida_gum_sys::GUM_RED_ZONE_SIZE as i32) as i64,
IndexMode::PreAdjust,
);
// Make sure the base register is copied into x0
match basereg {
Aarch64Register::X0 | Aarch64Register::W0 => {}
Aarch64Register::X1 | Aarch64Register::W1 => {
writer.put_mov_reg_reg(Aarch64Register::X0, Aarch64Register::X1);
}
_ => {
if !writer.put_mov_reg_reg(Aarch64Register::X0, basereg) {
writer.put_mov_reg_reg(Aarch64Register::W0, basereg);
}
}
}
// Make sure the index register is copied into x1
if indexreg.is_some() {
if let Some(indexreg) = indexreg {
match indexreg {
Aarch64Register::X0 | Aarch64Register::W0 => {
writer.put_ldr_reg_reg_offset(
Aarch64Register::X1,
Aarch64Register::Sp,
0u64,
);
}
Aarch64Register::X1 | Aarch64Register::W1 => {}
_ => {
if !writer.put_mov_reg_reg(Aarch64Register::X1, indexreg) {
writer.put_mov_reg_reg(Aarch64Register::W1, indexreg);
}
}
}
}
writer.put_add_reg_reg_reg(
Aarch64Register::X0,
Aarch64Register::X0,
Aarch64Register::X1,
);
}
let displacement = displacement
+ if basereg == Aarch64Register::Sp {
16 + frida_gum_sys::GUM_RED_ZONE_SIZE as i32
} else {
0
};
#[allow(clippy::comparison_chain)]
if displacement < 0 {
if displacement > -4096 {
// Subtract the displacement into x0
writer.put_sub_reg_reg_imm(
Aarch64Register::X0,
Aarch64Register::X0,
displacement.abs() as u64,
);
} else {
let displacement_hi = displacement.abs() / 4096;
let displacement_lo = displacement.abs() % 4096;
writer.put_bytes(&(0xd1400000u32 | ((displacement_hi as u32) << 10)).to_le_bytes());
writer.put_sub_reg_reg_imm(
Aarch64Register::X0,
Aarch64Register::X0,
displacement_lo as u64,
);
}
} else if displacement > 0 {
if displacement < 4096 {
// Add the displacement into x0
writer.put_add_reg_reg_imm(
Aarch64Register::X0,
Aarch64Register::X0,
displacement as u64,
);
} else {
let displacement_hi = displacement / 4096;
let displacement_lo = displacement % 4096;
writer.put_bytes(&(0x91400000u32 | ((displacement_hi as u32) << 10)).to_le_bytes());
writer.put_add_reg_reg_imm(
Aarch64Register::X0,
Aarch64Register::X0,
displacement_lo as u64,
);
}
}
// Insert the check_shadow_mem code blob
match width {
1 => writer.put_bytes(&self.asan_runtime.borrow().blob_check_mem_byte()),
2 => writer.put_bytes(&self.asan_runtime.borrow().blob_check_mem_halfword()),
3 => writer.put_bytes(&self.asan_runtime.borrow().blob_check_mem_3bytes()),
4 => writer.put_bytes(&self.asan_runtime.borrow().blob_check_mem_dword()),
6 => writer.put_bytes(&self.asan_runtime.borrow().blob_check_mem_6bytes()),
8 => writer.put_bytes(&self.asan_runtime.borrow().blob_check_mem_qword()),
12 => writer.put_bytes(&self.asan_runtime.borrow().blob_check_mem_12bytes()),
16 => writer.put_bytes(&self.asan_runtime.borrow().blob_check_mem_16bytes()),
24 => writer.put_bytes(&self.asan_runtime.borrow().blob_check_mem_24bytes()),
32 => writer.put_bytes(&self.asan_runtime.borrow().blob_check_mem_32bytes()),
_ => false,
};
// Restore x0, x1
assert!(writer.put_ldp_reg_reg_reg_offset(
Aarch64Register::X0,
Aarch64Register::X1,
Aarch64Register::Sp,
16 + frida_gum_sys::GUM_RED_ZONE_SIZE as i64,
IndexMode::PostAdjust,
));
}
#[cfg(target_arch = "aarch64")]
#[inline]
fn get_instruction_width(&self, instr: &Insn, operands: &Vec<arch::ArchOperand>) -> u32 {
use capstone::arch::arm64::Arm64Insn as I;
use capstone::arch::arm64::Arm64Reg as R;
use capstone::arch::arm64::Arm64Vas as V;
let num_registers = match instr.id().0.into() {
I::ARM64_INS_STP
| I::ARM64_INS_STXP
| I::ARM64_INS_STNP
| I::ARM64_INS_STLXP
| I::ARM64_INS_LDP
| I::ARM64_INS_LDXP
| I::ARM64_INS_LDNP => 2,
_ => 1,
};
let mnemonic = instr.mnemonic().unwrap();
match mnemonic.as_bytes().last().unwrap() {
b'b' => return 1,
b'h' => return 2,
b'w' => return 4 * num_registers,
_ => (),
}
if let Arm64Operand(operand) = operands.first().unwrap() {
if operand.vas != V::ARM64_VAS_INVALID {
let count_byte: u32 = if mnemonic.starts_with("st") || mnemonic.starts_with("ld") {
mnemonic.chars().nth(2).unwrap().to_digit(10).unwrap()
} else {
1
};
return match operand.vas {
V::ARM64_VAS_1B => 1 * count_byte,
V::ARM64_VAS_1H => 2 * count_byte,
V::ARM64_VAS_4B | V::ARM64_VAS_1S | V::ARM64_VAS_1D | V::ARM64_VAS_2H => {
4 * count_byte
}
V::ARM64_VAS_8B
| V::ARM64_VAS_4H
| V::ARM64_VAS_2S
| V::ARM64_VAS_2D
| V::ARM64_VAS_1Q => 8 * count_byte,
V::ARM64_VAS_8H | V::ARM64_VAS_4S | V::ARM64_VAS_16B => 16 * count_byte,
V::ARM64_VAS_INVALID => {
panic!("should not be reached");
}
};
} else if let Arm64OperandType::Reg(operand) = operand.op_type {
match operand.0 as u32 {
R::ARM64_REG_W0..=R::ARM64_REG_W30
| R::ARM64_REG_WZR
| R::ARM64_REG_WSP
| R::ARM64_REG_S0..=R::ARM64_REG_S31 => return 4 * num_registers,
R::ARM64_REG_D0..=R::ARM64_REG_D31 => return 8 * num_registers,
R::ARM64_REG_Q0..=R::ARM64_REG_Q31 => return 16,
_ => (),
}
};
};
8 * num_registers
}
#[cfg(target_arch = "aarch64")]
#[inline]
fn is_interesting_instruction(
&self,
_address: u64,
instr: &Insn,
) -> Result<(capstone::RegId, capstone::RegId, i32, u32), ()> {
// We have to ignore these instructions. Simulating them with their side effects is
// complex, to say the least.
match instr.mnemonic().unwrap() {
"ldaxr" | "stlxr" | "ldxr" | "stxr" | "ldar" | "stlr" | "ldarb" | "ldarh" | "ldaxp"
| "ldaxrb" | "ldaxrh" | "stlrb" | "stlrh" | "stlxp" | "stlxrb" | "stlxrh" | "ldxrb"
| "ldxrh" | "stxrb" | "stxrh" => return Err(()),
_ => (),
}
let operands = self
.capstone
.insn_detail(instr)
.unwrap()
.arch_detail()
.operands();
if operands.len() < 2 {
return Err(());
}
if let Arm64Operand(arm64operand) = operands.last().unwrap() {
if let Arm64OperandType::Mem(opmem) = arm64operand.op_type {
return Ok((
opmem.base(),
opmem.index(),
opmem.disp(),
self.get_instruction_width(instr, &operands),
));
}
}
Err(())
}
#[inline]
fn emit_coverage_mapping(&mut self, address: u64, output: &StalkerOutput) {
let writer = output.writer();
if self.current_log_impl == 0
|| !writer.can_branch_directly_to(self.current_log_impl)
|| !writer.can_branch_directly_between(writer.pc() + 128, self.current_log_impl)
{
let after_log_impl = writer.code_offset() + 1;
#[cfg(target_arch = "x86_64")]
writer.put_jmp_near_label(after_log_impl);
#[cfg(target_arch = "aarch64")]
writer.put_b_label(after_log_impl);
self.current_log_impl = writer.pc();
writer.put_bytes(&MAYBE_LOG_CODE);
let prev_loc_pointer = self.previous_pc.as_ptr() as usize;
let map_pointer = self.map.as_ptr() as usize;
writer.put_bytes(&map_pointer.to_ne_bytes());
writer.put_bytes(&prev_loc_pointer.to_ne_bytes());
writer.put_label(after_log_impl);
}
#[cfg(target_arch = "x86_64")]
{
println!("here");
writer.put_lea_reg_reg_offset(
X86Register::Rsp,
X86Register::Rsp,
-(frida_gum_sys::GUM_RED_ZONE_SIZE as i32),
);
writer.put_push_reg(X86Register::Rdi);
writer.put_mov_reg_address(
X86Register::Rdi,
((address >> 4) ^ (address << 8)) & (MAP_SIZE - 1) as u64,
);
writer.put_call_address(self.current_log_impl);
writer.put_pop_reg(X86Register::Rdi);
writer.put_lea_reg_reg_offset(
X86Register::Rsp,
X86Register::Rsp,
frida_gum_sys::GUM_RED_ZONE_SIZE as i32,
);
}
#[cfg(target_arch = "aarch64")]
{
writer.put_stp_reg_reg_reg_offset(
Aarch64Register::Lr,
Aarch64Register::X0,
Aarch64Register::Sp,
-(16 + frida_gum_sys::GUM_RED_ZONE_SIZE as i32) as i64,
IndexMode::PreAdjust,
);
writer.put_ldr_reg_u64(
Aarch64Register::X0,
((address >> 4) ^ (address << 8)) & (MAP_SIZE - 1) as u64,
);
writer.put_bl_imm(self.current_log_impl);
writer.put_ldp_reg_reg_reg_offset(
Aarch64Register::Lr,
Aarch64Register::X0,
Aarch64Register::Sp,
16 + frida_gum_sys::GUM_RED_ZONE_SIZE as i64,
IndexMode::PostAdjust,
);
}
}
}

122
libafl_frida/src/lib.rs Normal file
View File

@ -0,0 +1,122 @@
pub mod asan_rt;
pub mod helper;
/// A representation of the various Frida options
#[derive(Clone, Copy, Debug)]
pub struct FridaOptions {
enable_asan: bool,
enable_asan_leak_detection: bool,
enable_asan_continue_after_error: bool,
enable_asan_allocation_backtraces: bool,
enable_coverage: bool,
enable_drcov: bool,
}
impl FridaOptions {
/// Parse the frida options from the LIBAFL_FRIDA_OPTIONS environment variable.
///
/// Options are ':' separated, and each options is a 'name=value' string.
pub fn parse_env_options() -> Self {
let mut options = Self::default();
if let Ok(env_options) = std::env::var("LIBAFL_FRIDA_OPTIONS") {
for option in env_options.trim().to_lowercase().split(':') {
let (name, mut value) =
option.split_at(option.find('=').expect("Expected a '=' in option string"));
value = value.get(1..).unwrap();
match name {
"asan" => {
options.enable_asan = value.parse().unwrap();
#[cfg(not(target_arch = "aarch64"))]
if options.enable_asan {
panic!("ASAN is not currently supported on targets other than aarch64");
}
}
"asan-detect-leaks" => {
options.enable_asan_leak_detection = value.parse().unwrap();
}
"asan-continue-after-error" => {
options.enable_asan_continue_after_error = value.parse().unwrap();
}
"asan-allocation-backtraces" => {
options.enable_asan_allocation_backtraces = value.parse().unwrap();
}
"coverage" => {
options.enable_coverage = value.parse().unwrap();
}
"drcov" => {
options.enable_drcov = value.parse().unwrap();
#[cfg(not(target_arch = "aarch64"))]
if options.enable_drcov {
panic!(
"DrCov is not currently supported on targets other than aarch64"
);
}
}
_ => {
panic!("unknown FRIDA option: '{}'", option);
}
}
}
}
options
}
/// Is ASAN enabled?
#[inline]
pub fn asan_enabled(&self) -> bool {
self.enable_asan
}
/// Is coverage enabled?
#[inline]
pub fn coverage_enabled(&self) -> bool {
self.enable_coverage
}
/// Is DrCov enabled?
#[inline]
pub fn drcov_enabled(&self) -> bool {
self.enable_drcov
}
/// Should ASAN detect leaks
#[inline]
pub fn asan_detect_leaks(&self) -> bool {
self.enable_asan_leak_detection
}
/// Should ASAN continue after a memory error is detected
#[inline]
pub fn asan_continue_after_error(&self) -> bool {
self.enable_asan_continue_after_error
}
/// Should ASAN gather (and report) allocation-/free-site backtraces
#[inline]
pub fn asan_allocation_backtraces(&self) -> bool {
self.enable_asan_allocation_backtraces
}
/// Whether stalker should be enabled. I.e. whether at least one stalker requiring option is
/// enabled.
#[inline]
pub fn stalker_enabled(&self) -> bool {
self.enable_asan || self.enable_coverage || self.enable_drcov
}
}
impl Default for FridaOptions {
fn default() -> Self {
Self {
enable_asan: false,
enable_asan_leak_detection: false,
enable_asan_continue_after_error: false,
enable_asan_allocation_backtraces: true,
enable_coverage: true,
enable_drcov: false,
}
}
}

View File

@ -22,3 +22,4 @@ pcguard = ["pcguard_hitcounts"]
cc = { version = "1.0", features = ["parallel"] } cc = { version = "1.0", features = ["parallel"] }
[dependencies] [dependencies]
rangemap = "0.1.10"

View File

@ -0,0 +1,90 @@
use rangemap::RangeMap;
use std::{
fs::File,
io::{BufWriter, Write},
};
#[derive(Clone, Copy)]
pub struct DrCovBasicBlock {
start: usize,
end: usize,
}
pub struct DrCovWriter<'a> {
writer: BufWriter<File>,
module_mapping: &'a RangeMap<usize, (u16, &'a str)>,
basic_blocks: &'a mut Vec<DrCovBasicBlock>,
}
#[repr(C)]
struct DrCovBasicBlockEntry {
start: u32,
size: u16,
mod_id: u16,
}
impl DrCovBasicBlock {
pub fn new(start: usize, end: usize) -> Self {
Self { start, end }
}
}
impl<'a> DrCovWriter<'a> {
pub fn new(
path: &str,
module_mapping: &'a RangeMap<usize, (u16, &str)>,
basic_blocks: &'a mut Vec<DrCovBasicBlock>,
) -> Self {
Self {
writer: BufWriter::new(
File::create(path).expect("unable to create file for coverage data"),
),
module_mapping,
basic_blocks,
}
}
pub fn write(&mut self) {
self.writer
.write_all(b"DRCOV VERSION: 2\nDRCOV FLAVOR: libafl\n")
.unwrap();
let modules: Vec<(&std::ops::Range<usize>, &(u16, &str))> =
self.module_mapping.iter().collect();
self.writer
.write_all(format!("Module Table: version 2, count {}\n", modules.len()).as_bytes())
.unwrap();
self.writer
.write_all(b"Columns: id, base, end, entry, checksum, timestamp, path\n")
.unwrap();
for module in modules {
let (range, (id, path)) = module;
self.writer
.write_all(
format!(
"{:03}, 0x{:x}, 0x{:x}, 0x00000000, 0x00000000, 0x00000000, {}\n",
id, range.start, range.end, path
)
.as_bytes(),
)
.unwrap();
}
self.writer
.write_all(format!("BB Table: {} bbs\n", self.basic_blocks.len()).as_bytes())
.unwrap();
for block in self.basic_blocks.drain(0..) {
let (range, (id, _)) = self.module_mapping.get_key_value(&block.start).unwrap();
let basic_block = DrCovBasicBlockEntry {
start: (block.start - range.start) as u32,
size: (block.end - block.start) as u16,
mod_id: *id,
};
self.writer
.write_all(unsafe {
std::slice::from_raw_parts(&basic_block as *const _ as *const u8, 8)
})
.unwrap();
}
self.writer.flush().unwrap();
}
}

View File

@ -19,3 +19,5 @@ pub use libfuzzer::*;
pub mod cmplog; pub mod cmplog;
#[cfg(feature = "cmplog")] #[cfg(feature = "cmplog")]
pub use cmplog::*; pub use cmplog::*;
pub mod drcov;