Enable keep/reject inputs from the corpus

src/lib.rs

@@ -14,10 +14,37 @@
 pub use arbitrary;
 use once_cell::sync::OnceCell;
 
+/// Indicates whether the input should be kept in the corpus or rejected. This
+/// should be returned by your fuzz target. If your fuzz target does not return
+/// a value (i.e., returns `()`), then the input will be kept in the corpus.
+#[derive(Debug)]
+pub enum Corpus {
+    /// Keep the input in the corpus.
+    Keep,
+
+    /// Reject the input and do not keep it in the corpus.
+    Reject,
+}
+
+impl From<()> for Corpus {
+    fn from(_: ()) -> Self {
+        Self::Keep
+    }
+}
+
+impl From<Corpus> for i32 {
+    fn from(value: Corpus) -> i32 {
+        match value {
+            Corpus::Keep => 0,
+            Corpus::Reject => -1,
+        }
+    }
+}
+
 extern "C" {
     // We do not actually cross the FFI bound here.
     #[allow(improper_ctypes)]
-    fn rust_fuzzer_test_input(input: &[u8]);
+    fn rust_fuzzer_test_input(input: &[u8]) -> i32;
 
     fn LLVMFuzzerMutate(data: *mut u8, size: usize, max_size: usize) -> usize;
 }
@@ -27,14 +54,17 @@ extern "C" {
 pub fn test_input_wrap(data: *const u8, size: usize) -> i32 {
     let test_input = ::std::panic::catch_unwind(|| unsafe {
         let data_slice = ::std::slice::from_raw_parts(data, size);
-        rust_fuzzer_test_input(data_slice);
+        rust_fuzzer_test_input(data_slice)
     });
-    if test_input.err().is_some() {
-        // hopefully the custom panic hook will be called before and abort the
-        // process before the stack frames are unwinded.
-        ::std::process::abort();
+
+    match test_input {
+        Ok(i) => i,
+        Err(_) => {
+            // hopefully the custom panic hook will be called before and abort the
+            // process before the stack frames are unwinded.
+            ::std::process::abort();
+        }
     }
-    0
 }
 
 #[doc(hidden)]
@@ -86,6 +116,30 @@ pub fn initialize(_argc: *const isize, _argv: *const *const *const u8) -> isize
 /// # mod my_crate { pub fn parse(_: &[u8]) -> Result<(), ()> { unimplemented!() } }
 /// ```
 ///
+/// ## Rejecting Inputs
+///
+/// To indicate whether an input should be kept in or rejected from the corpus,
+/// return a [Corpus] value from your fuzz target. For example:
+///
+/// ```no_run
+/// #![no_main]
+///
+/// use libfuzzer_sys::{Corpus, fuzz_target};
+///
+/// fuzz_target!(|input: String| -> Corpus {
+///     let parts: Vec<&str> = input.splitn(2, '=').collect();
+///     if parts.len() != 2 {
+///         return Corpus::Reject;
+///     }
+///
+///     let key = parts[0];
+///     let value = parts[1];
+///     my_crate::parse(key, value);
+///     Corpus::Keep
+/// );
+/// # mod my_crate { pub fn parse(_key: &str, _value: &str) -> Result<(), ()> { unimplemented!() } }
+/// ```
+///
 /// ## Arbitrary Input Types
 ///
 /// The input is a `&[u8]` slice by default, but you can take arbitrary input
@@ -139,7 +193,7 @@ macro_rules! fuzz_target {
         const _: () = {
             /// Auto-generated function
             #[no_mangle]
-            pub extern "C" fn rust_fuzzer_test_input(bytes: &[u8]) {
+            pub extern "C" fn rust_fuzzer_test_input(bytes: &[u8]) -> i32 {
                 // When `RUST_LIBFUZZER_DEBUG_PATH` is set, write the debug
                 // formatting of the input to that file. This is only intended for
                 // `cargo fuzz`'s use!
@@ -154,7 +208,8 @@ macro_rules! fuzz_target {
                     return;
                 }
 
-                run(bytes)
+                run(bytes);
+                0
             }
 
             // Split out the actual fuzzer into a separate function which is
@@ -181,10 +236,14 @@ macro_rules! fuzz_target {
     };
 
     (|$data:ident: $dty: ty| $body:block) => {
+        $crate::fuzz_target!(|$data: $dty| -> () $body);
+    };
+
+    (|$data:ident: $dty: ty| -> $rty: ty $body:block) => {
         const _: () = {
             /// Auto-generated function
             #[no_mangle]
-            pub extern "C" fn rust_fuzzer_test_input(bytes: &[u8]) {
+            pub extern "C" fn rust_fuzzer_test_input(bytes: &[u8]) -> i32 {
                 use $crate::arbitrary::{Arbitrary, Unstructured};
 
                 // Early exit if we don't have enough bytes for the `Arbitrary`
@@ -194,7 +253,7 @@ macro_rules! fuzz_target {
                 // get to longer inputs that actually lead to interesting executions
                 // quicker.
                 if bytes.len() < <$dty as Arbitrary>::size_hint(0).0 {
-                    return;
+                    return -1;
                 }
 
                 let mut u = Unstructured::new(bytes);
@@ -214,20 +273,21 @@ macro_rules! fuzz_target {
                         Err(err) => writeln!(&mut file, "Arbitrary Error: {}", err),
                     })
                     .expect("failed to write to `RUST_LIBFUZZER_DEBUG_PATH` file");
-                    return;
+                    return -1;
                 }
 
                 let data = match data {
                     Ok(d) => d,
-                    Err(_) => return,
+                    Err(_) => return -1,
                 };
 
-                run(data)
+                let result: i32 = ::libfuzzer_sys::Corpus::from(run(data)).into();
+                result
             }
 
             // See above for why this is split to a separate function.
             #[inline(never)]
-            fn run($data: $dty) {
+            fn run($data: $dty) -> $rty {
                 $body
             }
         };