Enable keep/reject inputs from the corpus

src/lib.rs

@@ -14,10 +14,41 @@
 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 Corpus {
+    #[doc(hidden)]
+    /// Convert this Corpus result into the [integer codes used by
+    /// `libFuzzer`](https://llvm.org/docs/LibFuzzer.html#rejecting-unwanted-inputs).
+    /// This is -1 for reject, 0 for keep.
+    pub fn to_libfuzzer_code(self) -> i32 {
+        match self {
+            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 +58,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 +120,39 @@ pub fn initialize(_argc: *const isize, _argv: *const *const *const u8) -> isize
 /// # mod my_crate { pub fn parse(_: &[u8]) -> Result<(), ()> { unimplemented!() } }
 /// ```
 ///
+/// ## Rejecting Inputs
+///
+/// It may be desirable to reject some inputs, i.e. to not add them to the
+/// corpus.
+///
+/// For example, when fuzzing an API consisting of parsing and other logic,
+/// one may want to allow only those inputs into the corpus that parse
+/// successfully. To indicate whether an input should be kept in or rejected
+/// from the corpus, return either [Corpus::Keep] or [Corpus::Reject] from your
+/// fuzz target. The default behavior (e.g. if `()` is returned) is to keep the
+/// input in the corpus.
+/// 
+/// 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];
+///     let _result: Result<_, _> = 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 +206,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!
@@ -151,10 +218,11 @@ macro_rules! fuzz_target {
                         .expect("failed to create `RUST_LIBFUZZER_DEBUG_PATH` file");
                     writeln!(&mut file, "{:?}", bytes)
                         .expect("failed to write to `RUST_LIBFUZZER_DEBUG_PATH` file");
-                    return;
+                    return 0;
                 }
 
-                run(bytes)
+                run(bytes);
+                0
             }
 
             // Split out the actual fuzzer into a separate function which is
@@ -181,10 +249,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 +266,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 +286,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 = ::libfuzzer_sys::Corpus::from(run(data));
+                result.to_libfuzzer_code()
             }
 
             // See above for why this is split to a separate function.
             #[inline(never)]
-            fn run($data: $dty) {
+            fn run($data: $dty) -> $rty {
                 $body
             }
         };