Skip to content

Commit

Permalink
Baby PT with raw traces diff working
Browse files Browse the repository at this point in the history
  • Loading branch information
@Marcondiro
Marcondiro committed Sep 25, 2024
1 parent 05ba4bb commit f325dcf
Show file tree
Hide file tree
Showing 3 changed files with 112 additions and 118 deletions.
11 changes: 3 additions & 8 deletions fuzzers/baby/baby_fuzzer/src/main.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,13 +15,11 @@ use libafl::{
corpus::{InMemoryCorpus, OnDiskCorpus},
events::SimpleEventManager,
executors::{hooks::IntelPTHook, inprocess::GenericInProcessExecutor, ExitKind},
feedback_and,
feedbacks::{CrashFeedback, IntelPTFeedback, MaxMapFeedback},
feedbacks::{CrashFeedback, IntelPTFeedback},
fuzzer::{Fuzzer, StdFuzzer},
generators::RandPrintablesGenerator,
inputs::{BytesInput, HasTargetBytes},
mutators::{havoc_mutations::havoc_mutations, scheduled::StdScheduledMutator},
observers::StdMapObserver,
schedulers::QueueScheduler,
stages::mutational::StdMutationalStage,
state::StdState,
Expand Down Expand Up @@ -68,11 +66,8 @@ pub fn main() {
ExitKind::Ok
};

// Create an observation channel using the signals map
let observer = unsafe { StdMapObserver::from_mut_ptr("PT trace", SIGNALS_PTR, SIGNALS.len()) };

// Feedback to rate the interestingness of an input
let mut feedback = feedback_and!(IntelPTFeedback::new(pt_trace.clone())); //MaxMapFeedback::new(&observer),
let mut feedback = IntelPTFeedback::new(pt_trace.clone());

// A feedback to choose if an input is a solution or not
let mut objective = CrashFeedback::new();
Expand Down Expand Up @@ -121,7 +116,7 @@ pub fn main() {
let mut executor = PTInProcessExecutor::with_timeout_generic(
tuple_list!(pt_hook),
&mut harness,
tuple_list!(observer),
tuple_list!(),
&mut fuzzer,
&mut state,
&mut mgr,
Expand Down
201 changes: 99 additions & 102 deletions libafl/src/executors/hooks/intel_pt.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -114,46 +114,43 @@ where
assert!(self.pt.is_none(), "Intel PT was already set up");
let pid = process::id();
self.pt = Some(IntelPT::try_new(pid as i32).unwrap());
println!("Init hook");
}

#[allow(clippy::cast_possible_wrap)]
fn pre_exec(&mut self, _state: &mut S, _input: &S::Input) {
self.trace.lock().unwrap().clear();
self.pt.as_mut().unwrap().enable_tracing().unwrap();
println!("Pre exec hook");
}

#[allow(clippy::cast_possible_wrap)]
fn post_exec(&mut self, state: &mut S, _input: &S::Input) {
println!("Post exec hook");
self.pt.as_mut().unwrap().disable_tracing().unwrap();
let mut image = Image::new(Some("test_trace_pid")).unwrap();
// TODO optimize, move this stuff
let pid = process::id();
let maps = get_process_maps(pid as i32).unwrap();
for map in maps {
if map.is_exec() && map.filename().is_some() {
let _ = image.add_file(
map.filename().unwrap().to_str().unwrap(),
map.offset as u64,
map.size() as u64,
None,
map.start() as u64,
);
}
}
// let pid = process::id();
// let maps = get_process_maps(pid as i32).unwrap();
// for map in maps {
// if map.is_exec() && map.filename().is_some() {
// let _ = image.add_file(
// map.filename().unwrap().to_str().unwrap(),
// map.offset as u64,
// map.size() as u64,
// None,
// map.start() as u64,
// );
// }
// }
let mut buff = self.trace.lock().unwrap();
// let ips = self
// .pt
// .as_mut()
// .unwrap()
// .decode_with_image(&mut image, Some(&mut buff));
let s = serde_json::to_vec(&state).unwrap();
dump_corpus(&s).unwrap();
dump_trace_to_file(&buff).unwrap();
let ips = self
.pt
.as_mut()
.unwrap()
.decode_with_image(&mut image, Some(&mut buff));
// let s = serde_json::to_vec(&state).unwrap();
// dump_corpus(&s).unwrap();
// dump_trace_to_file(&buff).unwrap();
// println!("IPs: {ips:x?}");
println!("Post exec hook");
// println!("Post exec hook");
}
}

Expand Down Expand Up @@ -321,89 +318,89 @@ impl IntelPT {
}
};

println!("Intel PT: decoding {len} bytes");
if let Some(copy_buffer) = copy_buffer {
copy_buffer.extend_from_slice(&data);
}

smp_rmb(); // TODO double check impl

// println!("Intel PT: decoding {len} bytes");
if let Some(b) = copy_buffer {
b.extend_from_slice(&data);
}

// TODO handle decoding failures with config.decode.callback = <decode function>; config.decode.context = <decode context>;??
// apparently the rust library doesn't have the context parameter for the image.set_callback
// also, under the hood looks like it is passing the callback itself as context to the C fn 🤔
// TODO remove unwrap()
let mut config = ConfigBuilder::new(data.to_mut()).unwrap();
if let Some(cpu) = current_cpu() {
config.cpu(cpu);
}
let mut decoder = BlockDecoder::new(&config.finish()).unwrap();
if let Some(i) = image {
decoder.set_image(Some(i)).expect("Failed to set image");
}
if let Some(rm) = read_memory {
decoder
.image()
.unwrap()
.set_callback(Some(rm))
.expect("Failed to set get memory callback");
}
// TODO rewrite decently
// TODO consider dropping libipt-rs and using sys, or bindgen ourselves
let mut status;
loop {
status = match decoder.sync_forward() {
Ok(s) => s,
Err(e) => {
println!("pterror in sync {e:?}");
break;
}
};

loop {
if loop {
if !status.event_pending() {
break Ok(());
}
match decoder.event() {
Ok((_, s)) => {
// TODO maybe we care about some events?
status = s;
}
Err(e) => {
println!("pterror in event {e:?}");
break Err(e);
}
};
}
.is_err()
{
break;
}

let block = decoder.next();
match block {
Err((b, e)) => {
// libipt-rs library ignores the fact that
// Even in case of errors, we may have succeeded in decoding some instructions.
// https://github.com/intel/libipt/blob/4a06fdffae39dadef91ae18247add91029ff43c0/ptxed/src/ptxed.c#L1954
// Using my fork that fixes this atm
println!("pterror in packet next {e:?}");
println!("err block ip: 0x{:x?}", b.ip());
ips.push(b.ip());
// status = Status::from_bits(e.code() as u32).unwrap();
break;
}
Ok((b, s)) => {
status = s;
ips.push(b.ip());

if status.eos() {
break;
}
}
}
}
}
// let mut config = ConfigBuilder::new(data.to_mut()).unwrap();
// if let Some(cpu) = current_cpu() {
// config.cpu(cpu);
// }
// let mut decoder = BlockDecoder::new(&config.finish()).unwrap();
// if let Some(i) = image {
// decoder.set_image(Some(i)).expect("Failed to set image");
// }
// if let Some(rm) = read_memory {
// decoder
// .image()
// .unwrap()
// .set_callback(Some(rm))
// .expect("Failed to set get memory callback");
// }
// // TODO rewrite decently
// // TODO consider dropping libipt-rs and using sys, or bindgen ourselves
// let mut status;
// loop {
// status = match decoder.sync_forward() {
// Ok(s) => s,
// Err(e) => {
// // println!("pterror in sync {e:?}");
// break;
// }
// };
//
// loop {
// if loop {
// if !status.event_pending() {
// break Ok(());
// }
// match decoder.event() {
// Ok((_, s)) => {
// // TODO maybe we care about some events?
// status = s;
// }
// Err(e) => {
// // println!("pterror in event {e:?}");
// break Err(e);
// }
// };
// }
// .is_err()
// {
// break;
// }
//
// let block = decoder.next();
// match block {
// Err((b, e)) => {
// // libipt-rs library ignores the fact that
// // Even in case of errors, we may have succeeded in decoding some instructions.
// // https://github.com/intel/libipt/blob/4a06fdffae39dadef91ae18247add91029ff43c0/ptxed/src/ptxed.c#L1954
// // Using my fork that fixes this atm
// // println!("pterror in packet next {e:?}");
// // println!("err block ip: 0x{:x?}", b.ip());
// ips.push(b.ip());
// // status = Status::from_bits(e.code() as u32).unwrap();
// break;
// }
// Ok((b, s)) => {
// status = s;
// ips.push(b.ip());
//
// if status.eos() {
// break;
// }
// }
// }
// }
// }

unsafe { aux_tail.write_volatile(head) };

Expand Down
18 changes: 10 additions & 8 deletions libafl/src/feedbacks/intel_pt.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@

use alloc::borrow::Cow;
use std::{
cmp::min,
sync::{Arc, Mutex},
vec::Vec,
};
Expand All @@ -20,6 +21,7 @@ pub struct IntelPTFeedback {
trace: Arc<Mutex<Vec<u8>>>,
past_traces: Vec<Vec<u8>>,
avg_score: f64,
execution_number: usize,
}

impl IntelPTFeedback {
Expand All @@ -28,6 +30,7 @@ impl IntelPTFeedback {
trace,
past_traces: vec![],
avg_score: 0.0,
execution_number: 0,
}
}
}
Expand All @@ -54,13 +57,14 @@ where
EM: EventFirer<State = S>,
OT: ObserversTuple<S>,
{
self.execution_number += 1;
let trace = self.trace.lock().unwrap();
if self.past_traces.is_empty() {
self.past_traces.push(trace.clone());
return Ok(true);
}

let mut tot_score = 0;
let mut min_score = usize::MAX;
for pt in &self.past_traces {
let diff = capture_diff_slices(Algorithm::Myers, &trace, &pt);
let score = diff
Expand All @@ -72,16 +76,14 @@ where
DiffOp::Replace { new_len, .. } => *new_len,
})
.sum::<usize>();
tot_score += score;
min_score = min(min_score, score);
}

let weighted_score = tot_score as f64 / self.past_traces.len() as f64;
let n = self.execution_number as f64;
self.avg_score = (self.avg_score * (n - 1.0) + min_score as f64) / n;

self.past_traces.push(trace.clone());
let n = self.past_traces.len() as f64;
self.avg_score = (self.avg_score * (n - 1.0) + weighted_score) / n;

if n > 200.0 && weighted_score > self.avg_score * 2.0 {
if n > 50.0 && min_score as f64 > self.avg_score * 1.05 {
self.past_traces.push(trace.clone());
Ok(true)
} else {
Ok(false)
Expand Down

0 comments on commit f325dcf

Please sign in to comment.