Add utilities for handling panics and reading from stdin.

ci/run.sh

@@ -9,7 +9,6 @@ cargo build --verbose
 cargo build --example hello
 cargo build --example deferred-init
 cargo build --example integer-overflow
-cargo build --example panic
 cd afl-sys
 cargo build
 cd ..

examples/deferred-init.rs

@@ -12,7 +12,6 @@
 
 extern crate afl;
 
-use std::io::{self, Read};
 use std::thread;
 use std::time::Duration;
 
@@ -24,22 +23,16 @@ fn main() {
         afl::init();
     }
 
-    println!("the blink of an eye.");
-
-    let mut input = String::new();
-    io::stdin().read_to_string(&mut input).unwrap();
-
-    if input.starts_with("x") {
-        println!("going...");
-        if input.starts_with("xy") {
+    afl::handle_string(|input| {
+        println!("the blink of an eye.");
+        if input.starts_with("x") {
             println!("going...");
-            if input.starts_with("xyz") {
-                println!("gone!");
-                unsafe {
-                    let x: *mut usize = 0 as *mut usize;
-                    *x = 0xBEEF;
+            if input.starts_with("xy") {
+                println!("going...");
+                if input.starts_with("xyz") {
+                    panic!("gone!");
                 }
             }
         }
-    }
+    });
 }

examples/hello.rs

@@ -10,23 +10,16 @@
 
 extern crate afl;
 
-use std::io::{self, Read};
-
 fn main() {
-    let mut input = String::new();
-    io::stdin().read_to_string(&mut input).unwrap();
-
-    if input.starts_with("x") {
-        println!("going...");
-        if input.starts_with("xy") {
+    afl::handle_string(|input| {
+        if input.starts_with("x") {
             println!("going...");
-            if input.starts_with("xyz") {
-                println!("gone!");
-                unsafe {
-                    let x: *mut usize = 0 as *mut usize;
-                    *x = 0xBEEF;
+            if input.starts_with("xy") {
+                println!("going...");
+                if input.starts_with("xyz") {
+                    panic!("gone!");
                 }
             }
         }
-    }
+    })
 }

src/lib.rs

@@ -1,9 +1,16 @@
 // Copyright 2015 Keegan McAllister.
+// Copyright 2016 Corey Farwell.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
 // See `LICENSE` in this repository.
 
+#![feature(core_intrinsics)]
+
+use std::intrinsics::abort;
+use std::io::{self, Read};
+use std::panic::{self, UnwindSafe};
+
 #[cfg(not(test))]
 #[link(name="afl-llvm-rt", kind="static")]
 extern "C" {
@@ -26,6 +33,119 @@ pub unsafe fn init() {
     __afl_manual_init();
 }
 
+/// Utility that reads a `Vec` of bytes from standard input (stdin)
+/// and passes it to `closure`. All panics that occur within
+/// `closure` will be treated as aborts. This is done so that
+/// AFL considers a panic to be a crash.
+///
+/// # Examples
+///
+/// ```no_run
+/// extern crate afl;
+/// # struct Image;
+/// # impl Image {
+/// #     fn parse<R: ::std::io::Read>(_: R) {}
+/// # }
+///
+/// fn main() {
+///     afl::handle_read(|read| {
+///         Image::parse(read)
+///     })
+/// }
+/// ```
+pub fn handle_bytes<F>(closure: F)
+    where F: FnOnce(Vec<u8>) + UnwindSafe
+{
+    let mut input = vec![];
+    let result = io::stdin().read_to_end(&mut input);
+    if result.is_err() {
+        return;
+    }
+
+    let result = panic::catch_unwind(|| {
+        closure(input);
+    });
+    if result.is_err() {
+        unsafe {
+            abort();
+        }
+    }
+}
+
+/// Utility that reads a `String` from standard input (stdin) and
+/// passes it to `closure`. If a `String` cannot be constructed from
+/// the data provided by standard input, `closure` will _not_ be
+/// called. All panics that occur within `closure` will be treated as
+/// aborts. This is done so that AFL considers a panic to be a crash.
+///
+/// # Examples
+///
+/// ```no_run
+/// extern crate afl;
+/// # struct Url;
+/// # impl Url {
+/// #     fn parse(_: &str) {}
+/// # }
+///
+/// fn main() {
+///     afl::handle_string(|string| {
+///         Url::parse(&string)
+///     })
+/// }
+/// ```
+pub fn handle_string<F>(closure: F)
+    where F: FnOnce(String) + UnwindSafe
+{
+    let mut input = String::new();
+    let result = io::stdin().read_to_string(&mut input);
+    if result.is_err() {
+        return;
+    }
+    let result = panic::catch_unwind(|| {
+        closure(input);
+    });
+    if result.is_err() {
+        unsafe {
+            abort();
+        }
+    }
+}
+
+/// Utility that passes `Stdin` to `closure` for use with functions
+/// that expect a structure that implements `Read`. All panics that
+/// occur within `closure` will be treated as aborts. This is done
+/// so that
+/// AFL considers a panic to be a crash.
+///
+/// # Examples
+///
+/// ```no_run
+/// extern crate afl;
+/// # struct Image;
+/// # impl Image {
+/// #     fn parse(_: &[u8]) {}
+/// # }
+///
+/// fn main() {
+///     afl::handle_bytes(|bytes| {
+///         Image::parse(&bytes)
+///     })
+/// }
+/// ```
+pub fn handle_read<F>(closure: F)
+    where F: FnOnce(io::Stdin) + UnwindSafe
+{
+    let result = panic::catch_unwind(|| {
+        closure(io::stdin());
+    });
+    if result.is_err() {
+        // TODO: add option to prevent this abort?
+        unsafe {
+            abort();
+        }
+    }
+}
+
 #[cfg(test)]
 mod test {
     use std::process::{Command, Stdio};