diff --git a/custom_mutators/rust/Cargo.toml b/custom_mutators/rust/Cargo.toml index e36d24b5..00e6675f 100644 --- a/custom_mutators/rust/Cargo.toml +++ b/custom_mutators/rust/Cargo.toml @@ -5,4 +5,5 @@ members = [ "example", # Lain needs a nightly toolchain # "example_lain", -] \ No newline at end of file + # "example_lain_post_process", +] diff --git a/custom_mutators/rust/README.md b/custom_mutators/rust/README.md index e2cc38b4..a23b19df 100644 --- a/custom_mutators/rust/README.md +++ b/custom_mutators/rust/README.md @@ -5,7 +5,15 @@ Bindings to create custom mutators in Rust. These bindings are documented with rustdoc. To view the documentation run ```cargo doc -p custom_mutator --open```. -A minimal example can be found in `example`. Build it using `cargo build --example example_mutator`. +A minimal example can be found in `example`. Build it using `cargo build --example example_mutator`. An example using [lain](https://github.com/microsoft/lain) for structured fuzzing can be found in `example_lain`. Since lain requires a nightly rust toolchain, you need to set one up before you can play with it. + +An example for the use of the post_process function, using [lain](https://github.com/microsoft/lain) with [serde](https://github.com/serde-rs/serde) and [bincode](https://github.com/bincode-org/bincode) can be found in `example_lain_post_process`. +In order for it to work you need to: + +- disable input trimming with `AFL_DISABLE_TRIM=1` +- provide an initial instance serialized with `bincode` or use the `AFL_NO_STARTUP_CALIBRATION=1` environment variable. + +Note that `bincode` can also be used to serialize/deserialize the lain-generated structure and mutate it rather than generating a new one at each iteration, but it requires some structure serialized with `bincode` as input seed. diff --git a/custom_mutators/rust/custom_mutator/src/lib.rs b/custom_mutators/rust/custom_mutator/src/lib.rs index 3b635eb5..6c0e2f38 100644 --- a/custom_mutators/rust/custom_mutator/src/lib.rs +++ b/custom_mutators/rust/custom_mutator/src/lib.rs @@ -73,6 +73,8 @@ pub trait RawCustomMutator { None } + fn post_process<'b, 's: 'b>(&'s mut self, buffer: &'b mut [u8]) -> Option<&'b [u8]>; + /*fn post_process(&self, buffer: &[u8], unsigned char **out_buf)-> usize; int afl_custom_init_trim(&self, buffer: &[u8]); size_t afl_custom_trim(&self, unsigned char **out_buf); @@ -353,6 +355,33 @@ pub mod wrappers { Err(err) => panic_handler("afl_custom_queue_get", &err), } } + + /// Internal function used in the macro + pub unsafe fn afl_custom_post_process( + data: *mut c_void, + buf: *mut u8, + buf_size: usize, + out_buf: *mut *const u8, + ) -> usize { + match catch_unwind(|| { + let mut context = FFIContext::::from(data); + + assert!(!buf.is_null(), "null buf passed to afl_custom_post_process"); + assert!( + !out_buf.is_null(), + "null out_buf passed to afl_custom_post_process" + ); + let buff_slice = slice::from_raw_parts_mut(buf, buf_size); + if let Some(buffer) = context.mutator.post_process(buff_slice) { + *out_buf = buffer.as_ptr(); + return buffer.len(); + } + 0 + }) { + Ok(ret) => ret, + Err(err) => panic_handler("afl_custom_post_process", &err), + } + } } /// An exported macro to defined afl_custom_init meant for insternal usage @@ -480,6 +509,16 @@ macro_rules! export_mutator { pub unsafe extern "C" fn afl_custom_deinit(data: *mut ::std::os::raw::c_void) { $crate::wrappers::afl_custom_deinit_::<$mutator_type>(data) } + + #[no_mangle] + pub unsafe extern "C" fn afl_custom_post_process( + data: *mut ::std::os::raw::c_void, + buf: *mut u8, + buf_size: usize, + out_buf: *mut *const u8, + ) -> usize { + $crate::wrappers::afl_custom_post_process::<$mutator_type>(data, buf, buf_size, out_buf) + } }; } @@ -512,6 +551,10 @@ mod sanity_test { ) -> Option<&'b [u8]> { unimplemented!() } + + fn post_process<'b, 's: 'b>(&'s mut self, buffer: &'b mut [u8]) -> Option<&'b [u8]> { + unimplemented!() + } } export_mutator!(ExampleMutator); @@ -579,6 +622,13 @@ pub trait CustomMutator { fn introspection(&mut self) -> Result, Self::Error> { Ok(None) } + + fn post_process<'b, 's: 'b>( + &'s mut self, + buffer: &'b mut [u8], + ) -> Result, Self::Error> { + Ok(Some(buffer)) + } } impl RawCustomMutator for M @@ -682,6 +732,16 @@ where } } } + + fn post_process<'b, 's: 'b>(&'s mut self, buffer: &'b mut [u8]) -> Option<&'b [u8]> { + match self.post_process(buffer) { + Ok(r) => r, + Err(e) => { + Self::handle_error(e); + None + } + } + } } /// the default value to return from [`CustomMutator::describe`]. diff --git a/custom_mutators/rust/example_lain/Cargo.toml b/custom_mutators/rust/example_lain/Cargo.toml index c52bf86f..ae7f24c1 100644 --- a/custom_mutators/rust/example_lain/Cargo.toml +++ b/custom_mutators/rust/example_lain/Cargo.toml @@ -8,9 +8,9 @@ edition = "2021" [dependencies] custom_mutator = { path = "../custom_mutator" } -lain="0.5" +lain = { git = "https://github.com/AFLplusplus/lain.git" } [[example]] name = "example_lain" path = "./src/lain_mutator.rs" -crate-type = ["cdylib"] \ No newline at end of file +crate-type = ["cdylib"] diff --git a/custom_mutators/rust/example_lain_post_process/Cargo.toml b/custom_mutators/rust/example_lain_post_process/Cargo.toml new file mode 100644 index 00000000..fc7499d4 --- /dev/null +++ b/custom_mutators/rust/example_lain_post_process/Cargo.toml @@ -0,0 +1,21 @@ +[package] +name = "example_lain_post_process" +version = "0.1.0" +authors = [ + "Julius Hohnerlein ", + "jma <94166787+jma-qb@users.noreply.github.com>", +] +edition = "2021" + +# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html + +[dependencies] +custom_mutator = { path = "../custom_mutator" } +lain = { git = "https://github.com/AFLplusplus/lain.git" } +bincode = "1.3.3" +serde = { version = "1.0.214", features = ["derive"] } + +[[example]] +name = "example_lain_post_process" +path = "./src/lain_mutator.rs" +crate-type = ["cdylib"] diff --git a/custom_mutators/rust/example_lain_post_process/rust-toolchain b/custom_mutators/rust/example_lain_post_process/rust-toolchain new file mode 100644 index 00000000..07ade694 --- /dev/null +++ b/custom_mutators/rust/example_lain_post_process/rust-toolchain @@ -0,0 +1 @@ +nightly \ No newline at end of file diff --git a/custom_mutators/rust/example_lain_post_process/src/lain_mutator.rs b/custom_mutators/rust/example_lain_post_process/src/lain_mutator.rs new file mode 100644 index 00000000..8d03c0c0 --- /dev/null +++ b/custom_mutators/rust/example_lain_post_process/src/lain_mutator.rs @@ -0,0 +1,70 @@ +#![cfg(unix)] + +use custom_mutator::{export_mutator, CustomMutator}; +use lain::{ + mutator::Mutator, + prelude::*, + rand::{rngs::StdRng, SeedableRng}, +}; +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Deserialize, Serialize, Mutatable, NewFuzzed, BinarySerialize)] +struct MyStruct { + tag: u8, + #[lain(ignore)] + length: u32, + #[lain(min = 0, max = 10)] + data: Vec, +} + +struct LainMutator { + mutator: Mutator, + buffer: Vec, + post_buffer: Vec, +} + +impl CustomMutator for LainMutator { + type Error = (); + + fn init(seed: u32) -> Result { + Ok(Self { + mutator: Mutator::new(StdRng::seed_from_u64(seed as u64)), + buffer: Vec::new(), + post_buffer: Vec::new(), + }) + } + + fn fuzz<'b, 's: 'b>( + &'s mut self, + _buffer: &'b mut [u8], + _add_buff: Option<&[u8]>, + max_size: usize, + ) -> Result, ()> { + // we just sample an instance of MyStruct, ignoring the current input + let instance = MyStruct::new_fuzzed(&mut self.mutator, None); + let serialized = bincode::serialize(&instance).unwrap(); + let size = serialized.len(); + if size > max_size { + return Err(()); + } + self.buffer.clear(); + self.buffer.reserve(size); + self.buffer.extend_from_slice(&serialized); + Ok(Some(self.buffer.as_slice())) + } + + fn post_process<'b, 's: 'b>( + &'s mut self, + buffer: &'b mut [u8], + ) -> Result, Self::Error> { + let mut instance = bincode::deserialize::(&buffer).unwrap(); + instance.length = instance.data.len() as u32; + let size = instance.serialized_size(); + self.post_buffer.clear(); + self.post_buffer.reserve(size); + instance.binary_serialize::<_, BigEndian>(&mut self.post_buffer); + Ok(Some(&self.post_buffer)) + } +} + +export_mutator!(LainMutator);