[libFuzzer] Let user specify extra stats file.

Summary: If AFL_DRIVER_EXTRA_STATS_FILENAME is set and valid, write to it peak_rss_mb and slowest_unit_time_sec. These are both stats that libFuzzer can print but afl cannot.

Reviewers: kcc, aizatsky, metzman

Subscribers: llvm-commits

Differential Revision: http://reviews.llvm.org/D21742

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@274273 91177308-0d34-0410-b5e6-96231b3b80d8

lib/Fuzzer/afl/afl_driver.cpp

@@ -32,6 +32,22 @@ clang++ afl_driver.cpp test_fuzzer.o afl-llvm-rt.o.o
 rm -rf IN OUT; mkdir IN OUT; echo z > IN/z;
 $AFL_HOME/afl-fuzz -i IN -o OUT ./a.out
 ################################################################################
+Environment Variables:
+There are a few environment variables that can be set to use features that
+afl-fuzz doesn't have.
+
+AFL_DRIVER_STDERR_DUPLICATE_FILENAME: Setting this *appends* stderr to the file
+specified. If the file does not exist, it is created. This is useful for getting
+stack traces (when using ASAN for example) or original error messages on hard to
+reproduce bugs.
+
+AFL_DRIVER_EXTRA_STATS_FILENAME: Setting this causes afl_driver to write extra
+statistics to the file specified. Currently these are peak_rss_mb
+(the peak amount of virtual memory used in MB) and slowest_unit_time_secs. If
+the file does not exist it is created. If the file does exist then
+afl_driver assumes it was restarted by afl-fuzz and will try to read old
+statistics from the file. If that fails then the process will quit.
+
 */
 #include <assert.h>
 #include <stdio.h>
@@ -39,6 +55,31 @@ $AFL_HOME/afl-fuzz -i IN -o OUT ./a.out
 #include <stdlib.h>
 #include <string.h>
 #include <unistd.h>
+#include <errno.h>
+#include <signal.h>
+#include <sys/resource.h>
+#include <sys/time.h>
+// Platform detection. Copied from FuzzerInternal.h
+#ifdef __linux__
+#define LIBFUZZER_LINUX 1
+#define LIBFUZZER_APPLE 0
+#elif __APPLE__
+#define LIBFUZZER_LINUX 0
+#define LIBFUZZER_APPLE 1
+#else
+#error "Support for your platform has not been implemented"
+#endif
+
+// Used to avoid repeating error checking boilerplate. If cond is false, a
+// fatal error has occured in the program. In this event print error_message
+// to stderr and abort(). Otherwise do nothing. Note that setting
+// AFL_DRIVER_STDERR_DUPLICATE_FILENAME may cause error_message to be appended
+// to the file as well, if the error occurs after the duplication is performed.
+#define CHECK_ERROR(cond, error_message)                                       \
+  if (!(cond)) {                                                               \
+    fprintf(stderr, (error_message));                                          \
+    abort();                                                                   \
+  }
 
 // libFuzzer interface is thin, so we don't include any libFuzzer headers.
 extern "C" {
@@ -60,6 +101,124 @@ static volatile char suppress_warning1 = AFL_DEFER_FORKSVR[0];
 static const size_t kMaxAflInputSize = 1 << 20;
 static uint8_t AflInputBuf[kMaxAflInputSize];
 
+// Variables we need for writing to the extra stats file.
+static FILE *extra_stats_file = NULL;
+static uint32_t previous_peak_rss = 0;
+static time_t slowest_unit_time_secs = 0;
+static const int kNumExtraStats = 2;
+static const char *kExtraStatsFormatString = "peak_rss_mb            : %u\n"
+                                             "slowest_unit_time_sec  : %u\n";
+
+// Copied from FuzzerUtil.cpp.
+size_t GetPeakRSSMb() {
+  struct rusage usage;
+  if (getrusage(RUSAGE_SELF, &usage))
+    return 0;
+  if (LIBFUZZER_LINUX) {
+    // ru_maxrss is in KiB
+    return usage.ru_maxrss >> 10;
+  } else if (LIBFUZZER_APPLE) {
+    // ru_maxrss is in bytes
+    return usage.ru_maxrss >> 20;
+  }
+  assert(0 && "GetPeakRSSMb() is not implemented for your platform");
+  return 0;
+}
+
+// Based on SetSigaction in FuzzerUtil.cpp
+static void SetSigaction(int signum,
+                         void (*callback)(int, siginfo_t *, void *)) {
+  struct sigaction sigact;
+  memset(&sigact, 0, sizeof(sigact));
+  sigact.sa_sigaction = callback;
+  if (sigaction(signum, &sigact, 0)) {
+    fprintf(stderr, "libFuzzer: sigaction failed with %d\n", errno);
+    exit(1);
+  }
+}
+
+// Write extra stats to the file specified by the user. If none is specified
+// this function will never be called.
+static void write_extra_stats() {
+  uint32_t peak_rss = GetPeakRSSMb();
+
+  if (peak_rss < previous_peak_rss)
+    peak_rss = previous_peak_rss;
+
+  int chars_printed = fprintf(extra_stats_file, kExtraStatsFormatString,
+                              peak_rss, slowest_unit_time_secs);
+
+  CHECK_ERROR(chars_printed != 0, "Failed to write extra_stats_file");
+
+  CHECK_ERROR(fclose(extra_stats_file) == 0,
+              "Failed to close extra_stats_file");
+}
+
+// Call write_extra_stats before we exit.
+static void crash_handler(int, siginfo_t *, void *) {
+  // Make sure we don't try calling write_extra_stats again if we crashed while
+  // trying to call it.
+  static bool first_crash = true;
+  CHECK_ERROR(first_crash,
+              "Crashed in crash signal handler. This is a bug in the fuzzer.");
+
+  first_crash = false;
+  write_extra_stats();
+}
+
+// If the user has specified an extra_stats_file through the environment
+// variable AFL_DRIVER_EXTRA_STATS_FILENAME, then perform necessary set up
+// to write stats to it on exit. If no file is specified, do nothing. Otherwise
+// install signal and exit handlers to write to the file when the process exits.
+// Then if the file doesn't exist create it and set extra stats to 0. But if it
+// does exist then read the initial values of the extra stats from the file
+// and check that the file is writable.
+static void maybe_initialize_extra_stats() {
+  // If AFL_DRIVER_EXTRA_STATS_FILENAME isn't set then we have nothing to do.
+  char *extra_stats_filename = getenv("AFL_DRIVER_EXTRA_STATS_FILENAME");
+  if (!extra_stats_filename)
+    return;
+
+  // Open the file and find the previous peak_rss_mb value.
+  // This is necessary because the fuzzing process is restarted after N
+  // iterations are completed. So we may need to get this value from a previous
+  // process to be accurate.
+  extra_stats_file = fopen(extra_stats_filename, "r");
+
+  // If extra_stats_file already exists: read old stats from it.
+  if (extra_stats_file) {
+    int matches = fscanf(extra_stats_file, kExtraStatsFormatString,
+                         &previous_peak_rss, &slowest_unit_time_secs);
+
+    // Make sure we have read a real extra stats file and that we have used it
+    // to set slowest_unit_time_secs and previous_peak_rss.
+    CHECK_ERROR(matches == kNumExtraStats, "Extra stats file is corrupt");
+
+    CHECK_ERROR(fclose(extra_stats_file) == 0, "Failed to close file");
+
+    // Now open the file for writing.
+    extra_stats_file = fopen(extra_stats_filename, "w");
+    CHECK_ERROR(extra_stats_file,
+                "Failed to open extra stats file for writing");
+  } else {
+    // Looks like this is the first time in a fuzzing job this is being called.
+    extra_stats_file = fopen(extra_stats_filename, "w+");
+    CHECK_ERROR(extra_stats_file, "failed to create extra stats file");
+  }
+
+  // Make sure that crash_handler gets called on any kind of fatal error.
+  int crash_signals[] = {SIGSEGV, SIGBUS, SIGABRT, SIGILL, SIGFPE,  SIGINT,
+                         SIGTERM};
+
+  const size_t num_signals = sizeof(crash_signals) / sizeof(crash_signals[0]);
+
+  for (size_t idx = 0; idx < num_signals; idx++)
+    SetSigaction(crash_signals[idx], crash_handler);
+
+  // Make sure it gets called on other kinds of exits.
+  atexit(write_extra_stats);
+}
+
 // If the user asks us to duplicate stderr, then do it.
 static void maybe_duplicate_stderr() {
   char* stderr_duplicate_filename =
@@ -72,9 +231,9 @@ static void maybe_duplicate_stderr() {
       freopen(stderr_duplicate_filename, "a+", stderr);
 
   if (!stderr_duplicate_stream) {
-    fprintf(stderr,
-            "Failed to duplicate stderr to AFL_DRIVER_STDERR_DUPLICATE_FILENAME"
-            );
+    fprintf(
+        stderr,
+        "Failed to duplicate stderr to AFL_DRIVER_STDERR_DUPLICATE_FILENAME");
     abort();
   }
 }
@@ -90,6 +249,7 @@ int main(int argc, char **argv) {
   // Do any other expensive one-time initialization here.
 
   maybe_duplicate_stderr();
+  maybe_initialize_extra_stats();
 
   __afl_manual_init();
 
@@ -97,6 +257,7 @@ int main(int argc, char **argv) {
   if (argc >= 2)
     N = atoi(argv[1]);
   assert(N > 0);
+  time_t unit_time_secs;
   while (__afl_persistent_loop(N)) {
     ssize_t n_read = read(0, AflInputBuf, kMaxAflInputSize);
     if (n_read > 0) {
@@ -104,7 +265,22 @@ int main(int argc, char **argv) {
       // overflows. Don't use unique_ptr/etc to avoid extra dependencies.
       uint8_t *copy = new uint8_t[n_read];
       memcpy(copy, AflInputBuf, n_read);
+
+      struct timeval unit_start_time;
+      CHECK_ERROR(gettimeofday(&unit_start_time, NULL) == 0,
+                  "Calling gettimeofday failed");
+
       LLVMFuzzerTestOneInput(copy, n_read);
+
+      struct timeval unit_stop_time;
+      CHECK_ERROR(gettimeofday(&unit_stop_time, NULL) == 0,
+                  "Calling gettimeofday failed");
+
+      // Update slowest_unit_time_secs if we see a new max.
+      unit_time_secs = unit_stop_time.tv_sec - unit_start_time.tv_sec;
+      if (slowest_unit_time_secs < unit_time_secs)
+        slowest_unit_time_secs = unit_time_secs;
+
       delete[] copy;
     }
   }

lib/Fuzzer/test/afl-driver-extra-stats.test

@@ -0,0 +1,28 @@
+; Test that not specifying an extra stats file isn't broken.
+RUN: unset AFL_DRIVER_EXTRA_STATS_FILENAME
+RUN: AFLDriverTest
+
+; Test that specifying an invalid extra stats file causes a crash.
+RUN: ASAN_OPTIONS= AFL_DRIVER_EXTRA_STATS_FILENAME=%T not --crash AFLDriverTest
+
+; Test that specifying a corrupted stats file causes a crash.
+echo "peak_rss_mb :0" > %t
+ASAN_OPTIONS= AFL_DRIVER_EXTRA_STATS_FILENAME=%t not --crash AFLDriverTest
+
+; Test that specifying a valid nonexistent stats file works.
+RUN: rm -f %t
+RUN: AFL_DRIVER_EXTRA_STATS_FILENAME=%t AFLDriverTest
+RUN: [[ $(grep "peak_rss_mb\|slowest_unit_time_sec" %t | wc -l) -eq 2 ]]
+
+; Test that specifying a valid preexisting stats file works.
+RUN: printf "peak_rss_mb : 0\nslowest_unit_time_sec: 0\n" > %t
+RUN: AFL_DRIVER_EXTRA_STATS_FILENAME=%t AFLDriverTest
+; Check that both lines were printed.
+RUN: [[ $(grep "peak_rss_mb\|slowest_unit_time_sec" %t | wc -l) -eq 2 ]]
+
+; Test that peak_rss_mb and slowest_unit_time_in_secs are only updated when necessary.
+; Check that both lines have 9999 since there's no way we have exceeded that
+; amount of time or virtual memory.
+RUN: printf "peak_rss_mb : 9999\nslowest_unit_time_sec: 9999\n" > %t
+RUN: AFL_DRIVER_EXTRA_STATS_FILENAME=%t AFLDriverTest
+RUN: [[ $(grep "9999" %t | wc -l) -eq 2 ]]

lib/Fuzzer/test/afl-driver-stderr.test

@@ -1,4 +1,4 @@
-; Test that not specifying a file isn't broken.
+; Test that not specifying a stderr file isn't broken.
 RUN: unset AFL_DRIVER_STDERR_DUPLICATE_FILENAME
 RUN: AFLDriverTest