afl-tmin forkserver patch

3 files changed, 177 insertions, 20 deletions

diff --git a/afl-tmin.c b/afl-tmin.c
index 2d839041..a42be6e9 100644
--- a/afl-tmin.c
+++ b/afl-tmin.c
@@ -44,7 +44,11 @@
 #include <sys/types.h>
 #include <sys/resource.h>
 
-static s32 child_pid;                 /* PID of the tested program         */
+static s32 forksrv_pid,               /* PID of the fork server           */
+           child_pid;                 /* PID of the tested program        */
+
+static s32 fsrv_ctl_fd,               /* Fork server control pipe (write) */
+           fsrv_st_fd;                /* Fork server status pipe (read)   */
 
 static u8 *trace_bits,                /* SHM with instrumentation bitmap   */
           *mask_bitmap;               /* Mask for trace bits (-B)          */
@@ -55,6 +59,8 @@ static u8 *in_file,                   /* Minimizer input test case         */
           *target_path,               /* Path to target binary             */
           *doc_path;                  /* Path to docs                      */
 
+static s32 prog_in_fd;                /* Persistent fd for prog_in         */
+
 static u8* in_data;                   /* Input data for trimming           */
 
 static u32 in_len,                    /* Input data length                 */
@@ -236,38 +242,70 @@ static s32 write_to_file(u8* path, u8* mem, u32 len) {
 
 }
 
+/* Write modified data to file for testing. If use_stdin is clear, the old file
+   is unlinked and a new one is created. Otherwise, prog_in_fd is rewound and
+   truncated. */
+
+static void write_to_testcase(void* mem, u32 len) {
+
+  s32 fd = prog_in_fd;
+
+  if (!use_stdin) {
+
+    unlink(prog_in); /* Ignore errors. */
+
+    fd = open(prog_in, O_WRONLY | O_CREAT | O_EXCL, 0600);
+
+    if (fd < 0) PFATAL("Unable to create '%s'", prog_in);
+
+  } else lseek(fd, 0, SEEK_SET);
+
+  ck_write(fd, mem, len, prog_in);
+
+  if (use_stdin) {
+
+    if (ftruncate(fd, len)) PFATAL("ftruncate() failed");
+    lseek(fd, 0, SEEK_SET);
+
+  } else close(fd);
+
+}
+
+
 
 /* Handle timeout signal. */
 
 static void handle_timeout(int sig) {
 
+  if (child_pid > 0) {
+
   child_timed_out = 1;
-  if (child_pid > 0) kill(child_pid, SIGKILL);
+    kill(child_pid, SIGKILL);
 
-}
+  } else if (child_pid == -1 && forksrv_pid > 0) {
 
+    child_timed_out = 1;
+    kill(forksrv_pid, SIGKILL);
 
-/* Execute target application. Returns 0 if the changes are a dud, or
-   1 if they should be kept. */
+  }
 
-static u8 run_target(char** argv, u8* mem, u32 len, u8 first_run) {
+}
 
+/* start the app and it's forkserver */
+static void init_forkserver(char **argv) {
   static struct itimerval it;
+  int st_pipe[2], ctl_pipe[2];
   int status = 0;
+  s32 rlen;
 
-  s32 prog_in_fd;
-  u32 cksum;
-
-  memset(trace_bits, 0, MAP_SIZE);
-  MEM_BARRIER();
-
-  prog_in_fd = write_to_file(prog_in, mem, len);
+  ACTF("Spinning up the fork server...");
+  if (pipe(st_pipe) || pipe(ctl_pipe)) PFATAL("pipe() failed");
 
-  child_pid = fork();
+  forksrv_pid = fork();
 
-  if (child_pid < 0) PFATAL("fork() failed");
+  if (forksrv_pid < 0) PFATAL("fork() failed");
 
-  if (!child_pid) {
+  if (!forksrv_pid) {
 
     struct rlimit r;
 
@@ -304,6 +342,16 @@ static u8 run_target(char** argv, u8* mem, u32 len, u8 first_run) {
     r.rlim_max = r.rlim_cur = 0;
     setrlimit(RLIMIT_CORE, &r); /* Ignore errors */
 
+    /* Set up control and status pipes, close the unneeded original fds. */
+
+    if (dup2(ctl_pipe[0], FORKSRV_FD) < 0) PFATAL("dup2() failed");
+    if (dup2(st_pipe[1], FORKSRV_FD + 1) < 0) PFATAL("dup2() failed");
+
+    close(ctl_pipe[0]);
+    close(ctl_pipe[1]);
+    close(st_pipe[0]);
+    close(st_pipe[1]);
+
     execv(target_path, argv);
 
     *(u32*)trace_bits = EXEC_FAIL_SIG;
@@ -311,17 +359,113 @@ static u8 run_target(char** argv, u8* mem, u32 len, u8 first_run) {
 
   }
 
-  close(prog_in_fd);
+  /* Close the unneeded endpoints. */
+
+  close(ctl_pipe[0]);
+  close(st_pipe[1]);
+
+  fsrv_ctl_fd = ctl_pipe[1];
+  fsrv_st_fd  = st_pipe[0];
+
+  /* Configure timeout, wait for child, cancel timeout. */
+
+  if (exec_tmout) {
+
+    child_timed_out = 0;
+    it.it_value.tv_sec = (exec_tmout * FORK_WAIT_MULT / 1000);
+    it.it_value.tv_usec = ((exec_tmout * FORK_WAIT_MULT) % 1000) * 1000;
+
+  }
+
+  setitimer(ITIMER_REAL, &it, NULL);
+
+  rlen = read(fsrv_st_fd, &status, 4);
+
+  it.it_value.tv_sec = 0;
+  it.it_value.tv_usec = 0;
+  setitimer(ITIMER_REAL, &it, NULL);
+
+  /* If we have a four-byte "hello" message from the server, we're all set.
+     Otherwise, try to figure out what went wrong. */
+
+  if (rlen == 4) {
+    ACTF("All right - fork server is up.");
+    return;
+  }
+
+  if (waitpid(forksrv_pid, &status, 0) <= 0)
+    PFATAL("waitpid() failed");
+
+  u8 child_crashed;
+
+  if (WIFSIGNALED(status))
+    child_crashed = 1;
+
+  if (child_timed_out)
+    SAYF(cLRD "\n+++ Program timed off +++\n" cRST);
+  else if (stop_soon)
+    SAYF(cLRD "\n+++ Program aborted by user +++\n" cRST);
+  else if (child_crashed)
+    SAYF(cLRD "\n+++ Program killed by signal %u +++\n" cRST, WTERMSIG(status));
+
+}
+
+
+/* Execute target application. Returns 0 if the changes are a dud, or
+   1 if they should be kept. */
+
+static u8 run_target(char** argv, u8* mem, u32 len, u8 first_run) {
+
+  static struct itimerval it;
+  static u32 prev_timed_out = 0;
+  int status = 0;
+
+  u32 cksum;
+
+  memset(trace_bits, 0, MAP_SIZE);
+  MEM_BARRIER();
+
+  write_to_testcase(mem, len);
+
+  s32 res;
+
+  /* we have the fork server up and running, so simply
+     tell it to have at it, and then read back PID. */
+
+  if ((res = write(fsrv_ctl_fd, &prev_timed_out, 4)) != 4) {
+
+    if (stop_soon) return 0;
+    RPFATAL(res, "Unable to request new process from fork server (OOM?)");
+
+  }
+
+  if ((res = read(fsrv_st_fd, &child_pid, 4)) != 4) {
+
+    if (stop_soon) return 0;
+    RPFATAL(res, "Unable to request new process from fork server (OOM?)");
+
+  }
+
+  if (child_pid <= 0) FATAL("Fork server is misbehaving (OOM?)");
 
   /* Configure timeout, wait for child, cancel timeout. */
 
+  if (exec_tmout) {
+
   child_timed_out = 0;
   it.it_value.tv_sec = (exec_tmout / 1000);
   it.it_value.tv_usec = (exec_tmout % 1000) * 1000;
 
+  }
+
   setitimer(ITIMER_REAL, &it, NULL);
 
-  if (waitpid(child_pid, &status, 0) <= 0) FATAL("waitpid() failed");
+  if ((res = read(fsrv_st_fd, &status, 4)) != 4) {
+
+    if (stop_soon) return 0;
+    RPFATAL(res, "Unable to communicate with fork server (OOM?)");
+
+  }
 
   child_pid = 0;
   it.it_value.tv_sec = 0;
@@ -687,6 +831,13 @@ static void set_up_environment(void) {
 
   }
 
+  unlink(prog_in);
+
+  prog_in_fd = open(prog_in, O_RDWR | O_CREAT | O_EXCL, 0600);
+
+  if (prog_in_fd < 0) PFATAL("Unable to create '%s'", prog_in);
+
+
   /* Set sane defaults... */
 
   x = getenv("ASAN_OPTIONS");
@@ -1113,6 +1264,8 @@ int main(int argc, char** argv) {
 
   read_initial_file();
 
+  init_forkserver(use_argv);
+
   ACTF("Performing dry run (mem limit = %llu MB, timeout = %u ms%s)...",
        mem_limit, exec_tmout, edges_only ? ", edges only" : "");
 
diff --git a/docs/ChangeLog b/docs/ChangeLog
index abb5fd46..0d730118 100644
--- a/docs/ChangeLog
+++ b/docs/ChangeLog
@@ -17,13 +17,16 @@ sending a mail to <afl-users+subscribe@googlegroups.com>.
 Version ++2.52d (tbd):
 -----------------------------
 
-  - more power to afl-system-config :)
+  - more cpu power for afl-system-config
+  - added forkserver patch to afl-tmin, makes it much faster (originally from
+    github.com/nccgroup/TriforceAFL)
   - added whitelist support for llvm_mode via AFL_LLVM_WHITELIST to allow
     only to instrument what is actually interesting. Gives more speed and less
     map pollution (originally by choller@mozilla)
   - added Python Module mutator support, python2.7-dev is autodetected.
     see docs/python_mutators.txt (originally by choller@mozilla)
-  - added AFL_CAL_FAST for slow applications and AFL_DEBUG_CHILD_OUTPUT for debugging
+  - added AFL_CAL_FAST for slow applications and AFL_DEBUG_CHILD_OUTPUT for
+    debugging
   - added a  -s seed  switch to allow afl run with a fixed initial
     seed that is not updated. this is good for performance and path discovery
     tests as the random numbers are deterministic then
diff --git a/docs/PATCHES b/docs/PATCHES
index c933e031..f61f8d24 100644
--- a/docs/PATCHES
+++ b/docs/PATCHES
@@ -21,6 +21,7 @@ afl-qemu-optimize-map.diff		by mh(at)mh-sec(dot)de
 + Qemu 3.1 upgrade with enhancement patches (github.com/andreafioraldi/afl)
 + Python mutator modules support (github.com/choeller/afl)
 + Whitelisting in LLVM mode (github.com/choeller/afl)
++ forkserver patch for afl-tmin (github.com/nccgroup/TriforceAFL)
 
 
 NOT INSTALLED