Add support for post_process in Rust custom mutator + associated example with lain (#2241)

7 files changed, 165 insertions, 4 deletions

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<M: RawCustomMutator>(
+        data: *mut c_void,
+        buf: *mut u8,
+        buf_size: usize,
+        out_buf: *mut *const u8,
+    ) -> usize {
+        match catch_unwind(|| {
+            let mut context = FFIContext::<M>::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<Option<&str>, Self::Error> {
         Ok(None)
     }
+
+    fn post_process<'b, 's: 'b>(
+        &'s mut self,
+        buffer: &'b mut [u8],
+    ) -> Result<Option<&'b [u8]>, Self::Error> {
+        Ok(Some(buffer))
+    }
 }
 
 impl<M> 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 <julihoh@users.noreply.github.com>",
+    "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<u8>,
+}
+
+struct LainMutator {
+    mutator: Mutator<StdRng>,
+    buffer: Vec<u8>,
+    post_buffer: Vec<u8>,
+}
+
+impl CustomMutator for LainMutator {
+    type Error = ();
+
+    fn init(seed: u32) -> Result<Self, ()> {
+        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<Option<&'b [u8]>, ()> {
+        // 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<Option<&'b [u8]>, Self::Error> {
+        let mut instance = bincode::deserialize::<MyStruct>(&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);