Forkserver for afl-analyze (#963)

* afl-analyze forkserver

* added missing vars to forkserver

* synchronized a bit more with afl-tmin

* more debugging, runs now, but need to suppress target output

* fix dev/null setting

* afl-analyze info:

Co-authored-by: hexcoder- <heiko@hexco.de>

1 files changed, 81 insertions, 154 deletions

diff --git a/src/afl-analyze.c b/src/afl-analyze.c
index d43278b9..606254d9 100644
--- a/src/afl-analyze.c
+++ b/src/afl-analyze.c
@@ -55,12 +55,7 @@
 #include <sys/types.h>
 #include <sys/resource.h>
 
-static s32 child_pid;                  /* PID of the tested program         */
-
-static u8 *trace_bits;                 /* SHM with instrumentation bitmap   */
-
-static u8 *in_file,                    /* Analyzer input test case          */
-    *prog_in;                          /* Targeted program input file       */
+static u8 *in_file;                    /* Analyzer input test case          */
 
 static u8 *in_data;                    /* Input data for analysis           */
 
@@ -73,20 +68,19 @@ static u64 orig_cksum;                 /* Original checksum                 */
 
 static u64 mem_limit = MEM_LIMIT;      /* Memory limit (MB)                 */
 
-static s32 dev_null_fd = -1;           /* FD to /dev/null                   */
-
 static bool edges_only,                  /* Ignore hit counts?              */
     use_hex_offsets,                   /* Show hex offsets?                 */
     use_stdin = true;                     /* Use stdin for program input?   */
 
-static volatile u8 stop_soon,          /* Ctrl-C pressed?                   */
-    child_timed_out;                   /* Child timed out?                  */
+static volatile u8 stop_soon;          /* Ctrl-C pressed?                   */
 
 static u8 *target_path;
 static u8  frida_mode;
 static u8  qemu_mode;
 static u32 map_size = MAP_SIZE;
 
+static afl_forkserver_t fsrv = {0};   /* The forkserver                     */
+
 /* Constants used for describing byte behavior. */
 
 #define RESP_NONE 0x00                 /* Changing byte is a no-op.         */
@@ -156,7 +150,7 @@ static void classify_counts(u8 *mem) {
 
 static inline u8 anything_set(void) {
 
-  u32 *ptr = (u32 *)trace_bits;
+  u32 *ptr = (u32 *)fsrv.trace_bits;
   u32  i = (map_size >> 2);
 
   while (i--) {
@@ -173,7 +167,7 @@ static inline u8 anything_set(void) {
 
 static void at_exit_handler(void) {
 
-  unlink(prog_in);                                         /* Ignore errors */
+  unlink(fsrv.out_file);                                         /* Ignore errors */
 
 }
 
@@ -205,128 +199,29 @@ static void read_initial_file(void) {
 
 }
 
-/* Write output file. */
-
-static s32 write_to_file(u8 *path, u8 *mem, u32 len) {
-
-  s32 ret;
-
-  unlink(path);                                            /* Ignore errors */
-
-  ret = open(path, O_RDWR | O_CREAT | O_EXCL, DEFAULT_PERMISSION);
-
-  if (ret < 0) { PFATAL("Unable to create '%s'", path); }
-
-  ck_write(ret, mem, len, path);
-
-  lseek(ret, 0, SEEK_SET);
-
-  return ret;
-
-}
-
-/* Handle timeout signal. */
-
-static void handle_timeout(int sig) {
-
-  (void)sig;
-
-  child_timed_out = 1;
-
-  if (child_pid > 0) kill(child_pid, SIGKILL);
-
-}
-
 /* Execute target application. Returns exec checksum, or 0 if program
    times out. */
 
-static u32 analyze_run_target(char **argv, u8 *mem, u32 len, u8 first_run) {
-
-  static struct itimerval it;
-  int                     status = 0;
-
-  s32 prog_in_fd;
-  u64 cksum;
-
-  memset(trace_bits, 0, map_size);
-  MEM_BARRIER();
-
-  prog_in_fd = write_to_file(prog_in, mem, len);
-
-  child_pid = fork();
-
-  if (child_pid < 0) { PFATAL("fork() failed"); }
-
-  if (!child_pid) {
-
-    struct rlimit r;
-
-    if (dup2(use_stdin ? prog_in_fd : dev_null_fd, 0) < 0 ||
-        dup2(dev_null_fd, 1) < 0 || dup2(dev_null_fd, 2) < 0) {
-
-      *(u32 *)trace_bits = EXEC_FAIL_SIG;
-      PFATAL("dup2() failed");
-
-    }
-
-    close(dev_null_fd);
-    close(prog_in_fd);
-
-    if (mem_limit) {
+static u32 analyze_run_target(u8 *mem, u32 len, u8 first_run) {
 
-      r.rlim_max = r.rlim_cur = ((rlim_t)mem_limit) << 20;
+  afl_fsrv_write_to_testcase(&fsrv, mem, len);
+  fsrv_run_result_t ret = afl_fsrv_run_target(&fsrv, exec_tmout, &stop_soon);
 
-#ifdef RLIMIT_AS
+  if (ret == FSRV_RUN_ERROR) {
 
-      setrlimit(RLIMIT_AS, &r);                            /* Ignore errors */
+    FATAL("Error in forkserver");
 
-#else
-
-      setrlimit(RLIMIT_DATA, &r);                          /* Ignore errors */
-
-#endif                                                        /* ^RLIMIT_AS */
-
-    }
+  } else if (ret == FSRV_RUN_NOINST) {
 
-    r.rlim_max = r.rlim_cur = 0;
-    setrlimit(RLIMIT_CORE, &r);                            /* Ignore errors */
+    FATAL("Target not instrumented");
 
-    execv(target_path, argv);
+  } else if (ret == FSRV_RUN_NOBITS) {
 
-    *(u32 *)trace_bits = EXEC_FAIL_SIG;
-    exit(0);
+    FATAL("Failed to run target");
 
   }
 
-  close(prog_in_fd);
-
-  /* Configure timeout, wait for child, cancel timeout. */
-
-  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"); }
-
-  child_pid = 0;
-  it.it_value.tv_sec = 0;
-  it.it_value.tv_usec = 0;
-
-  setitimer(ITIMER_REAL, &it, NULL);
-
-  MEM_BARRIER();
-
-  /* Clean up bitmap, analyze exit condition, etc. */
-
-  if (*(u32 *)trace_bits == EXEC_FAIL_SIG) {
-
-    FATAL("Unable to execute '%s'", argv[0]);
-
-  }
-
-  classify_counts(trace_bits);
+  classify_counts(fsrv.trace_bits);
   total_execs++;
 
   if (stop_soon) {
@@ -338,21 +233,19 @@ static u32 analyze_run_target(char **argv, u8 *mem, u32 len, u8 first_run) {
 
   /* Always discard inputs that time out. */
 
-  if (child_timed_out) {
+  if (fsrv.last_run_timed_out) {
 
     exec_hangs++;
     return 0;
 
   }
 
-  cksum = hash64(trace_bits, map_size, HASH_CONST);
+  u64 cksum = hash64(fsrv.trace_bits, fsrv.map_size, HASH_CONST);
 
-  /* We don't actually care if the target is crashing or not,
-     except that when it does, the checksum should be different. */
+  if (ret == FSRV_RUN_CRASH) {
 
-  if (WIFSIGNALED(status) ||
-      (WIFEXITED(status) && WEXITSTATUS(status) == MSAN_ERROR) ||
-      (WIFEXITED(status) && WEXITSTATUS(status))) {
+    /* We don't actually care if the target is crashing or not,
+       except that when it does, the checksum should be different. */
 
     cksum ^= 0xffffffff;
 
@@ -616,7 +509,7 @@ static void dump_hex(u32 len, u8 *b_data) {
 
 /* Actually analyze! */
 
-static void analyze(char **argv) {
+static void analyze() {
 
   u32 i;
   u32 boring_len = 0, prev_xff = 0, prev_x01 = 0, prev_s10 = 0, prev_a10 = 0;
@@ -642,16 +535,16 @@ static void analyze(char **argv) {
        code. */
 
     in_data[i] ^= 0xff;
-    xor_ff = analyze_run_target(argv, in_data, in_len, 0);
+    xor_ff = analyze_run_target(in_data, in_len, 0);
 
     in_data[i] ^= 0xfe;
-    xor_01 = analyze_run_target(argv, in_data, in_len, 0);
+    xor_01 = analyze_run_target(in_data, in_len, 0);
 
     in_data[i] = (in_data[i] ^ 0x01) - 0x10;
-    sub_10 = analyze_run_target(argv, in_data, in_len, 0);
+    sub_10 = analyze_run_target(in_data, in_len, 0);
 
     in_data[i] += 0x20;
-    add_10 = analyze_run_target(argv, in_data, in_len, 0);
+    add_10 = analyze_run_target(in_data, in_len, 0);
     in_data[i] -= 0x10;
 
     /* Classify current behavior. */
@@ -724,7 +617,7 @@ static void handle_stop_sig(int sig) {
   (void)sig;
   stop_soon = 1;
 
-  if (child_pid > 0) { kill(child_pid, SIGKILL); }
+  afl_fsrv_killall();
 
 }
 
@@ -736,10 +629,10 @@ static void set_up_environment(char **argv) {
   char *afl_preload;
   char *frida_afl_preload = NULL;
 
-  dev_null_fd = open("/dev/null", O_RDWR);
-  if (dev_null_fd < 0) { PFATAL("Unable to open /dev/null"); }
+  fsrv.dev_null_fd = open("/dev/null", O_RDWR);
+  if (fsrv.dev_null_fd < 0) { PFATAL("Unable to open /dev/null"); }
 
-  if (!prog_in) {
+  if (!fsrv.out_file) {
 
     u8 *use_dir = ".";
 
@@ -750,10 +643,15 @@ static void set_up_environment(char **argv) {
 
     }
 
-    prog_in = alloc_printf("%s/.afl-analyze-temp-%u", use_dir, (u32)getpid());
+    fsrv.out_file = alloc_printf("%s/.afl-analyze-temp-%u", use_dir, (u32)getpid());
 
   }
 
+  unlink(fsrv.out_file);
+  fsrv.out_fd = open(fsrv.out_file, O_RDWR | O_CREAT | O_EXCL, DEFAULT_PERMISSION);
+
+  if (fsrv.out_fd < 0) { PFATAL("Unable to create '%s'", fsrv.out_file); }
+
   /* Set sane defaults... */
 
   x = get_afl_env("ASAN_OPTIONS");
@@ -916,11 +814,6 @@ static void setup_signal_handlers(void) {
   sigaction(SIGINT, &sa, NULL);
   sigaction(SIGTERM, &sa, NULL);
 
-  /* Exec timeout notifications. */
-
-  sa.sa_handler = handle_timeout;
-  sigaction(SIGALRM, &sa, NULL);
-
 }
 
 /* Display usage hints. */
@@ -982,6 +875,8 @@ int main(int argc, char **argv_orig, char **envp) {
 
   SAYF(cCYA "afl-analyze" VERSION cRST " by Michal Zalewski\n");
 
+  afl_fsrv_init(&fsrv);
+
   while ((opt = getopt(argc, argv, "+i:f:m:t:eOQUWh")) > 0) {
 
     switch (opt) {
@@ -994,9 +889,9 @@ int main(int argc, char **argv_orig, char **envp) {
 
       case 'f':
 
-        if (prog_in) { FATAL("Multiple -f options not supported"); }
-        use_stdin = 0;
-        prog_in = optarg;
+        if (fsrv.out_file) { FATAL("Multiple -f options not supported"); }
+        fsrv.use_stdin = 0;
+        fsrv.out_file = ck_strdup(optarg);
         break;
 
       case 'e':
@@ -1017,6 +912,7 @@ int main(int argc, char **argv_orig, char **envp) {
         if (!strcmp(optarg, "none")) {
 
           mem_limit = 0;
+          fsrv.mem_limit = 0;
           break;
 
         }
@@ -1055,6 +951,8 @@ int main(int argc, char **argv_orig, char **envp) {
 
         }
 
+        fsrv.mem_limit = mem_limit;
+
       }
 
       break;
@@ -1074,6 +972,8 @@ int main(int argc, char **argv_orig, char **envp) {
 
         }
 
+        fsrv.exec_tmout = exec_tmout;
+
         break;
 
       case 'O':                                               /* FRIDA mode */
@@ -1081,6 +981,7 @@ int main(int argc, char **argv_orig, char **envp) {
         if (frida_mode) { FATAL("Multiple -O options not supported"); }
 
         frida_mode = 1;
+        fsrv.frida_mode = frida_mode;
 
         break;
 
@@ -1090,6 +991,8 @@ int main(int argc, char **argv_orig, char **envp) {
         if (!mem_limit_given) { mem_limit = MEM_LIMIT_QEMU; }
 
         qemu_mode = 1;
+        fsrv.mem_limit = mem_limit;
+        fsrv.qemu_mode = qemu_mode;
         break;
 
       case 'U':
@@ -1098,6 +1001,7 @@ int main(int argc, char **argv_orig, char **envp) {
         if (!mem_limit_given) { mem_limit = MEM_LIMIT_UNICORN; }
 
         unicorn_mode = 1;
+        fsrv.mem_limit = mem_limit;
         break;
 
       case 'W':                                           /* Wine+QEMU mode */
@@ -1107,6 +1011,8 @@ int main(int argc, char **argv_orig, char **envp) {
         use_wine = 1;
 
         if (!mem_limit_given) { mem_limit = 0; }
+        fsrv.qemu_mode = qemu_mode;
+        fsrv.mem_limit = mem_limit;
 
         break;
 
@@ -1125,6 +1031,7 @@ int main(int argc, char **argv_orig, char **envp) {
   if (optind == argc || !in_file) { usage(argv[0]); }
 
   map_size = get_map_size();
+  fsrv.map_size = map_size;
 
   use_hex_offsets = !!get_afl_env("AFL_ANALYZE_HEX");
 
@@ -1134,14 +1041,15 @@ int main(int argc, char **argv_orig, char **envp) {
 
   /* initialize cmplog_mode */
   shm.cmplog_mode = 0;
-  trace_bits = afl_shm_init(&shm, map_size, 0);
+
   atexit(at_exit_handler);
   setup_signal_handlers();
 
   set_up_environment(argv);
 
-  target_path = find_binary(argv[optind]);
-  detect_file_args(argv + optind, prog_in, &use_stdin);
+  fsrv.target_path = find_binary(argv[optind]);
+  fsrv.trace_bits = afl_shm_init(&shm, map_size, 0);
+  detect_file_args(argv + optind, fsrv.out_file, &use_stdin);
 
   if (qemu_mode) {
 
@@ -1165,14 +1073,31 @@ int main(int argc, char **argv_orig, char **envp) {
 
   SAYF("\n");
 
+  if (getenv("AFL_FORKSRV_INIT_TMOUT")) {
+
+    s32 forksrv_init_tmout = atoi(getenv("AFL_FORKSRV_INIT_TMOUT"));
+    if (forksrv_init_tmout < 1) {
+
+      FATAL("Bad value specified for AFL_FORKSRV_INIT_TMOUT");
+
+    }
+
+    fsrv.init_tmout = (u32)forksrv_init_tmout;
+
+  }
+
+  fsrv.kill_signal =
+      parse_afl_kill_signal_env(getenv("AFL_KILL_SIGNAL"), SIGKILL);
+
   read_initial_file();
 
   ACTF("Performing dry run (mem limit = %llu MB, timeout = %u ms%s)...",
        mem_limit, exec_tmout, edges_only ? ", edges only" : "");
 
-  analyze_run_target(use_argv, in_data, in_len, 1);
+  afl_fsrv_start(&fsrv, use_argv, &stop_soon, false);
+  analyze_run_target(in_data, in_len, 1);
 
-  if (child_timed_out) {
+  if (fsrv.last_run_timed_out) {
 
     FATAL("Target binary times out (adjusting -t may help).");
 
@@ -1184,13 +1109,15 @@ int main(int argc, char **argv_orig, char **envp) {
 
   }
 
-  analyze(use_argv);
+  analyze();
 
   OKF("We're done here. Have a nice day!\n");
 
-  if (target_path) { ck_free(target_path); }
-
   afl_shm_deinit(&shm);
+  afl_fsrv_deinit(&fsrv);
+  if (fsrv.target_path) { ck_free(fsrv.target_path); }
+  if (in_data) { ck_free(in_data); }
+
 
   exit(0);