initial broken attempt at deriving mapping mutators

Author: riesentoaster

libafl/src/mutators/havoc_mutations.rs

@@ -1,6 +1,6 @@
 //! [`crate::mutators::Mutator`] collection equivalent to AFL++'s havoc mutations
 
-use libafl_bolts::tuples::{Map, Merge};
+use libafl_bolts::tuples::{Map, Merge, NamedTuple};
 use tuple_list::{tuple_list, tuple_list_type};
 
 use crate::mutators::{
@@ -271,3 +271,29 @@ where
             current_input_mapper,
         ))
 }
+
+/// TODO
+pub trait HasHavocMutators {
+    /// TODO
+    fn havoc_mutators<MT: NamedTuple>() -> MT;
+}
+
+#[cfg(test)]
+mod tests {
+    use libafl_derive::HasHavocMutators;
+
+    use super::HasHavocMutators;
+    use crate::mutators::{StdScheduledMutator, Vec};
+
+    #[derive(HasHavocMutators)]
+    struct CustomInput {
+        vec: Vec<u8>,
+    }
+
+    #[test]
+    fn test_derive_has_havoc_mutators() {
+        let input = CustomInput { vec: vec![] };
+        let mutations = CustomInput::havoc_mutators();
+        let scheduler = StdScheduledMutator::new(mutations);
+    }
+}

libafl_derive/Cargo.toml

@@ -24,3 +24,4 @@ proc-macro = true
 syn = { version = "2.0.77", features = ["full", "extra-traits"] }
 quote = "1.0.37"
 proc-macro2 = "1.0.86"
+proc-macro-crate = "3.2"

libafl_derive/src/lib.rs

@@ -58,8 +58,17 @@
 )]
 
 use proc_macro::TokenStream;
-use quote::quote;
-use syn::{parse_macro_input, Data::Struct, DeriveInput, Field, Fields::Named, Type};
+use proc_macro_crate::{crate_name, FoundCrate};
+use quote::{format_ident, quote};
+use syn::{
+    parse_macro_input,
+    punctuated::Punctuated,
+    token::Comma,
+    Data::Struct,
+    DeriveInput, Error, Field,
+    Fields::{Named, Unit, Unnamed},
+    GenericArgument, Ident, PathArguments, PathSegment, Type,
+};
 
 /// Derive macro to implement `SerdeAny`, to use a type in a `SerdeAnyMap`
 #[proc_macro_derive(SerdeAny)]
@@ -157,3 +166,139 @@ fn libafl_display_field_by_type(it: &Field) -> proc_macro2::TokenStream {
         write!(f, #fmt, self.#ident)?;
     }
 }
+
+/// TODO
+#[proc_macro_derive(HasHavocMutators)]
+pub fn derive_has_mutator_bytes(input: TokenStream) -> TokenStream {
+    let input_ast = parse_macro_input!(input as DeriveInput);
+
+    let struct_name = input_ast.ident.clone();
+
+    let fields = match extract_fields(input_ast) {
+        Ok(f) => f,
+        Err(e) => return e.into_compile_error().into(),
+    };
+
+    let (getter_methods, mutator_merge_call) = match create_functions_on_fields(&fields) {
+        Ok(e) => e,
+        Err(e) => return e.into_compile_error().into(),
+    };
+
+    // required to be able to use it from within libafl — used for testing
+    let libafl_source = match crate_name("libafl").expect("Could not figure out current crate") {
+        FoundCrate::Itself => quote! { crate },
+        FoundCrate::Name(_) => quote! { libafl },
+    };
+
+    // Generate the impl block
+    let expanded = quote! {
+        use #libafl_source::{inputs::MutVecInput, mutators::{Mutator, mapped_havoc_mutations}};
+        use libafl_bolts::tuples::{Merge, NamedTuple, tuple_list};
+
+        impl #struct_name {
+            #getter_methods
+        }
+
+        impl HasHavocMutators for #struct_name {
+            fn havoc_mutators<MT: NamedTuple>() -> MT {
+                #mutator_merge_call
+            }
+        }
+    };
+
+    TokenStream::from(expanded)
+}
+
+fn extract_fields(ast: DeriveInput) -> Result<Punctuated<Field, Comma>, Error> {
+    match &ast.data {
+        Struct(data_struct) => match &data_struct.fields {
+            Named(fields_named) => Ok(fields_named.named.clone()),
+            Unnamed(fields_unnamed) => Ok(fields_unnamed.unnamed.clone()),
+            Unit => Err(Error::new_spanned(
+                ast,
+                "HasHavocMutators can not be derived for unit structs",
+            )),
+        },
+        _ => Err(Error::new_spanned(
+            ast,
+            "HasHavocMutators can only be derived for structs",
+        )),
+    }
+}
+
+fn create_functions_on_fields(
+    fields: &Punctuated<Field, Comma>,
+) -> Result<(proc_macro2::TokenStream, proc_macro2::TokenStream), Error> {
+    let functions_res = fields.iter().map(|field| match field.ty.clone() {
+        Type::Path(type_path) => {
+            let segment = type_path.path.segments.last().unwrap();
+            if let Some(tokens) = create_functions_on_type(segment, field.ident.as_ref().unwrap()) {
+                return Ok(tokens);
+            }
+
+            Err(Error::new_spanned(
+                segment.ident.clone(),
+                "HasHavocMutators does not support struct parts of this type",
+            ))
+        }
+        _ => Err(Error::new_spanned(
+            field,
+            "HasHavocMutators can only be derived for structs",
+        )),
+    });
+
+    // check if any fields could not be parsed into functions, combine the errors and return them
+    if let Some(errors) = functions_res
+        .clone()
+        .filter(Result::is_err)
+        .map(Result::unwrap_err)
+        .reduce(|mut acc, e| {
+            acc.combine(e);
+            acc
+        })
+    {
+        return Err(errors);
+    }
+
+    Ok(functions_res.map(Result::unwrap).fold(
+        (quote! {}, quote! { tuple_list!() }),
+        |(acc1, acc2), (e1, e2)| {
+            (
+                quote! {
+                    #acc1
+                    #e1
+                },
+                quote! { #acc2.merge(#e2) },
+            )
+        },
+    ))
+}
+
+fn create_functions_on_type(
+    segment: &PathSegment,
+    field_name: &Ident,
+) -> Option<(proc_macro2::TokenStream, proc_macro2::TokenStream)> {
+    if segment.ident == "Vec" {
+        if let PathArguments::AngleBracketed(args) = &segment.arguments {
+            if let Some(GenericArgument::Type(Type::Path(arg_type))) = args.args.first() {
+                let arg_ident = &arg_type.path.segments.last().unwrap().ident;
+                if arg_ident == "u8" {
+                    let mutable_method_name = format_ident!("{}_mut", field_name);
+                    let immutable_method_name = field_name;
+                    return Some((
+                        quote! {
+                            pub fn #mutable_method_name(&mut self) -> MutVecInput<'_> {
+                                (&mut self.#field_name).into()
+                            }
+                            pub fn #immutable_method_name(&self) -> &[u8] {
+                                &self.#field_name
+                            }
+                        },
+                        quote! { mapped_havoc_mutations(Self::#mutable_method_name, Self::#immutable_method_name) },
+                    ));
+                }
+            }
+        }
+    }
+    None
+}