code format

4 files changed, 75 insertions, 109 deletions

diff --git a/src/afl-common.c b/src/afl-common.c
index 134d3180..c0202821 100644
--- a/src/afl-common.c
+++ b/src/afl-common.c
@@ -138,32 +138,19 @@ void argv_cpy_free(char **argv) {
 
 }
 
+u8 *find_binary_own_loc(u8 *fname, u8 *own_loc) {
 
-/* Rewrite argv for QEMU. */
-
-char **get_qemu_argv(u8 *own_loc, u8 **target_path_p, int argc, char **argv) {
-
-  char **new_argv = ck_alloc(sizeof(char *) * (argc + 4));
-  u8 *   tmp, *cp = NULL, *rsl, *own_copy;
-
-  memcpy(&new_argv[3], &argv[1], (int)(sizeof(char *)) * (argc - 1));
-  new_argv[argc - 1] = NULL;
-
-  new_argv[2] = *target_path_p;
-  new_argv[1] = "--";
-
-  /* Now we need to actually find the QEMU binary to put in argv[0]. */
+  u8 *tmp, *rsl, *own_copy, *cp;
 
   tmp = getenv("AFL_PATH");
 
   if (tmp) {
 
-    cp = alloc_printf("%s/afl-qemu-trace", tmp);
+    cp = alloc_printf("%s/%s", tmp, fname);
 
     if (access(cp, X_OK)) { FATAL("Unable to find '%s'", tmp); }
 
-    *target_path_p = new_argv[0] = cp;
-    return new_argv;
+    return cp;
 
   }
 
@@ -174,15 +161,10 @@ char **get_qemu_argv(u8 *own_loc, u8 **target_path_p, int argc, char **argv) {
 
     *rsl = 0;
 
-    cp = alloc_printf("%s/afl-qemu-trace", own_copy);
+    cp = alloc_printf("%s/%s", own_copy, fname);
     ck_free(own_copy);
 
-    if (!access(cp, X_OK)) {
-
-      *target_path_p = new_argv[0] = cp;
-      return new_argv;
-
-    }
+    if (!access(cp, X_OK)) { return cp; }
 
   } else {
 
@@ -190,11 +172,35 @@ char **get_qemu_argv(u8 *own_loc, u8 **target_path_p, int argc, char **argv) {
 
   }
 
-  if (!access(BIN_PATH "/afl-qemu-trace", X_OK)) {
+  cp = alloc_printf("%s/%s", BIN_PATH, fname);
+  if (!access(cp, X_OK)) { return cp; }
+
+  ck_free(cp);
+
+  return NULL;
 
-    if (cp) { ck_free(cp); }
-    *target_path_p = new_argv[0] = ck_strdup(BIN_PATH "/afl-qemu-trace");
+}
+
+/* Rewrite argv for QEMU. */
+
+char **get_qemu_argv(u8 *own_loc, u8 **target_path_p, int argc, char **argv) {
+
+  char **new_argv = ck_alloc(sizeof(char *) * (argc + 4));
+  u8 *   cp = NULL;
+
+  memcpy(&new_argv[3], &argv[1], (int)(sizeof(char *)) * (argc - 1));
+  new_argv[argc - 1] = NULL;
 
+  new_argv[2] = *target_path_p;
+  new_argv[1] = "--";
+
+  /* Now we need to actually find the QEMU binary to put in argv[0]. */
+
+  cp = find_binary_own_loc("afl-qemu-trace", own_loc);
+
+  if (cp) {
+
+    *target_path_p = new_argv[0] = cp;
     return new_argv;
 
   }
@@ -235,66 +241,16 @@ char **get_wine_argv(u8 *own_loc, u8 **target_path_p, int argc, char **argv) {
 
   /* Now we need to actually find the QEMU binary to put in argv[0]. */
 
-  tmp = getenv("AFL_PATH");
-
-  if (tmp) {
+  cp = find_binary_own_loc("afl-qemu-trace", own_loc);
 
-    cp = alloc_printf("%s/afl-qemu-trace", tmp);
-
-    if (access(cp, X_OK)) { FATAL("Unable to find '%s'", tmp); }
+  if (cp) {
 
     ck_free(cp);
+    cp = find_binary_own_loc("afl-wine-trace", own_loc);
 
-    cp = alloc_printf("%s/afl-wine-trace", tmp);
+    if (cp) {
 
-    if (access(cp, X_OK)) { FATAL("Unable to find '%s'", tmp); }
-
-    *target_path_p = new_argv[0] = cp;
-    return new_argv;
-
-  }
-
-  own_copy = ck_strdup(own_loc);
-  rsl = strrchr(own_copy, '/');
-
-  if (rsl) {
-
-    *rsl = 0;
-
-    cp = alloc_printf("%s/afl-qemu-trace", own_copy);
-
-    if (cp && !access(cp, X_OK)) {
-
-      ck_free(cp);
-
-      cp = alloc_printf("%s/afl-wine-trace", own_copy);
-
-      if (!access(cp, X_OK)) {
-
-        *target_path_p = new_argv[0] = cp;
-        return new_argv;
-
-      }
-
-    }
-
-    ck_free(own_copy);
-
-  } else {
-
-    ck_free(own_copy);
-
-  }
-
-  u8 *ncp = BIN_PATH "/afl-qemu-trace";
-
-  if (!access(ncp, X_OK)) {
-
-    ncp = BIN_PATH "/afl-wine-trace";
-
-    if (!access(ncp, X_OK)) {
-
-      *target_path_p = new_argv[0] = ck_strdup(ncp);
+      *target_path_p = new_argv[0] = cp;
       return new_argv;
 
     }
@@ -302,25 +258,21 @@ char **get_wine_argv(u8 *own_loc, u8 **target_path_p, int argc, char **argv) {
   }
 
   SAYF("\n" cLRD "[-] " cRST
-       "Oops, unable to find the '%s' binary. The binary must be "
-       "built\n"
-       "    separately by following the instructions in "
-       "qemu_mode/README.md. "
-       "If you\n"
-       "    already have the binary installed, you may need to specify "
-       "AFL_PATH in the\n"
-       "    environment.\n\n"
-
+       "Oops, unable to find the afl-qemu-trace and afl-wine-trace binaries.\n"
+       "The afl-qemu-trace binary must be built separately by following the "
+       "instructions\n"
+       "in qemu_mode/README.md. If you already have the binary installed, you "
+       "may need\n"
+       "to specify the location via AFL_PATH in the environment.\n\n"
        "    Of course, even without QEMU, afl-fuzz can still work with "
        "binaries that are\n"
        "    instrumented at compile time with afl-gcc. It is also possible to "
        "use it as a\n"
        "    traditional non-instrumented fuzzer by specifying '-n' in the "
        "command "
-       "line.\n",
-       ncp);
+       "line.\n");
 
-  FATAL("Failed to locate '%s'.", ncp);
+  FATAL("Failed to locate 'afl-qemu-trace' and 'afl-wine-trace'.");
 
 }
 
diff --git a/src/afl-forkserver.c b/src/afl-forkserver.c
index b4f92e5b..eeb2f8c3 100644
--- a/src/afl-forkserver.c
+++ b/src/afl-forkserver.c
@@ -481,7 +481,6 @@ void afl_fsrv_start(afl_forkserver_t *fsrv, char **argv,
            "handle_sigill=0",
            0);
 
-fprintf(stderr, "init %p\n", fsrv->init_child_func);
     fsrv->init_child_func(fsrv, argv);
 
     /* Use a distinctive bitmap signature to tell the parent about execv()
@@ -497,19 +496,19 @@ fprintf(stderr, "init %p\n", fsrv->init_child_func);
 
   char pid_buf[16];
   sprintf(pid_buf, "%d", fsrv->fsrv_pid);
-  
+
   if (fsrv->qemu_mode == 2) {
 
     setenv("__AFL_TARGET_PID3", pid_buf, 1);
 
   } else if (fsrv->cmplog_binary) {
-  
+
     setenv("__AFL_TARGET_PID2", pid_buf, 1);
-  
+
   } else {
-  
+
     setenv("__AFL_TARGET_PID1", pid_buf, 1);
-  
+
   }
 
   /* Close the unneeded endpoints. */
diff --git a/src/afl-fuzz-run.c b/src/afl-fuzz-run.c
index 207b3046..badc2239 100644
--- a/src/afl-fuzz-run.c
+++ b/src/afl-fuzz-run.c
@@ -472,16 +472,20 @@ abort_calibration:
   afl->stage_max = old_sm;
 
   /* if taint mode was selected, run the taint */
-  
+
   if (afl->fsrv.taint_mode) {
+
     write_to_testcase(afl, use_mem, q->len);
-    if (afl_fsrv_run_target(&afl->taint_fsrv, use_tmout, &afl->stop_soon) == 0) {
+    if (afl_fsrv_run_target(&afl->taint_fsrv, use_tmout, &afl->stop_soon) ==
+        0) {
+
       u32 len = q->len / 8;
       if (q->len % 8) len++;
       u32 bits = count_bits_len(afl, afl->taint_fsrv.trace_bits, len);
       if (afl->debug) fprintf(stderr, "Debug: tainted bytes: %u\n", bits);
 
     }
+
   }
 
   if (!first_run) { show_stats(afl); }
diff --git a/src/afl-fuzz.c b/src/afl-fuzz.c
index bc780b55..684b123e 100644
--- a/src/afl-fuzz.c
+++ b/src/afl-fuzz.c
@@ -92,7 +92,8 @@ static void usage(u8 *argv0, int more_help) {
       "  -o dir        - output directory for fuzzer findings\n\n"
 
       "Execution control settings:\n"
-      "  -A            - use first level taint analysis (see qemu_taint/README.md)\n"
+      "  -A            - use first level taint analysis (see "
+      "qemu_taint/README.md)\n"
       "  -p schedule   - power schedules compute a seed's performance score. "
       "<explore\n"
       "                  (default), fast, coe, lin, quad, exploit, mmopt, "
@@ -1247,7 +1248,7 @@ int main(int argc, char **argv_orig, char **envp) {
     OKF("Cmplog forkserver successfully started");
 
   }
-  
+
   if (afl->fsrv.taint_mode) {
 
     ACTF("Spawning qemu_taint forkserver");
@@ -1256,11 +1257,21 @@ int main(int argc, char **argv_orig, char **envp) {
     afl->taint_fsrv.trace_bits = afl->fsrv.trace_bits;
     ck_free(afl->taint_fsrv.target_path);
     afl->taint_fsrv.target_path = ck_strdup(afl->fsrv.target_path);
-    afl->argv_taint = get_qemu_argv(argv[0], &afl->taint_fsrv.target_path,
-                                    argc - optind, argv + optind);
-    u32 len = strlen(afl->taint_fsrv.target_path);
-    strcpy(afl->taint_fsrv.target_path + len - 5, "taint");
-    strcpy((afl->argv_taint[0]) + len - 5, "taint");
+    afl->argv_taint = ck_alloc(sizeof(char *) * (argc + 4 - optind));
+    afl->argv_taint[0] = find_binary_own_loc("afl-qemu-taint", argv[0]);
+    if (!afl->argv_taint[0])
+      FATAL(
+          "Cannot find 'afl-qemu-taint', read qemu_taint/README.md on how to "
+          "build it.");
+    u32 idx = optind - 1, offset = 0;
+    do {
+
+      idx++;
+      offset++;
+      afl->argv_taint[offset] = argv[idx];
+
+    } while (argv[idx] != NULL);
+
     if (afl->fsrv.use_stdin)
       unsetenv("AFL_TAINT_INPUT");
     else