Merge pull request #817 from realmadsci/atat-plusplus

Make @@ handling more flexible

2 files changed, 31 insertions, 46 deletions

diff --git a/src/afl-common.c b/src/afl-common.c
index 37b4788c..1f9839a2 100644
--- a/src/afl-common.c
+++ b/src/afl-common.c
@@ -70,31 +70,26 @@ void detect_file_args(char **argv, u8 *prog_in, bool *use_stdin) {
 
       *use_stdin = false;
 
-      if (prog_in[0] != 0) {  // not afl-showmap special case
+      /* Be sure that we're always using fully-qualified paths. */
 
-        u8 *n_arg;
+      *aa_loc = 0;
 
-        /* Be sure that we're always using fully-qualified paths. */
+      /* Construct a replacement argv value. */
+      u8 *n_arg;
 
-        *aa_loc = 0;
+      if (prog_in[0] == '/') {
 
-        /* Construct a replacement argv value. */
+        n_arg = alloc_printf("%s%s%s", argv[i], prog_in, aa_loc + 2);
 
-        if (prog_in[0] == '/') {
-
-          n_arg = alloc_printf("%s%s%s", argv[i], prog_in, aa_loc + 2);
-
-        } else {
-
-          n_arg = alloc_printf("%s%s/%s%s", argv[i], cwd, prog_in, aa_loc + 2);
-
-        }
+      } else {
 
-        ck_free(argv[i]);
-        argv[i] = n_arg;
+        n_arg = alloc_printf("%s%s/%s%s", argv[i], cwd, prog_in, aa_loc + 2);
 
       }
 
+      ck_free(argv[i]);
+      argv[i] = n_arg;
+
     }
 
     i++;
diff --git a/src/afl-showmap.c b/src/afl-showmap.c
index 077c9248..558665a2 100644
--- a/src/afl-showmap.c
+++ b/src/afl-showmap.c
@@ -72,8 +72,7 @@ static u8 *in_data,                    /* Input data                        */
 static u64 total;                      /* tuple content information         */
 static u32 tcnt, highest;              /* tuple content information         */
 
-static u32 in_len,                     /* Input data length                 */
-    arg_offset;                        /* Total number of execs             */
+static u32 in_len;                     /* Input data length                 */
 
 static u32 map_size = MAP_SIZE;
 
@@ -956,12 +955,30 @@ int main(int argc, char **argv_orig, char **envp) {
 
   }
 
+
   if (in_dir) {
 
-    detect_file_args(argv + optind, "", &fsrv->use_stdin);
+    /* If we don't have a file name chosen yet, use a safe default. */
+    u8 *use_dir = ".";
+
+    if (access(use_dir, R_OK | W_OK | X_OK)) {
+
+      use_dir = get_afl_env("TMPDIR");
+      if (!use_dir) { use_dir = "/tmp"; }
+
+    }
+
+    stdin_file = at_file ? strdup(at_file)
+                         : (char *)alloc_printf("%s/.afl-showmap-temp-%u",
+                                                use_dir, (u32)getpid());
+    unlink(stdin_file);
+
+    // If @@ are in the target args, replace them and also set use_stdin=false.
+    detect_file_args(argv + optind, stdin_file, &fsrv->use_stdin);
 
   } else {
 
+    // If @@ are in the target args, replace them and also set use_stdin=false.
     detect_file_args(argv + optind, at_file, &fsrv->use_stdin);
 
   }
@@ -986,14 +1003,6 @@ int main(int argc, char **argv_orig, char **envp) {
 
   }
 
-  i = 0;
-  while (use_argv[i] != NULL && !arg_offset) {
-
-    if (strcmp(use_argv[i], "@@") == 0) { arg_offset = i; }
-    i++;
-
-  }
-
   shm_fuzz = ck_alloc(sizeof(sharedmem_t));
 
   /* initialize cmplog_mode */
@@ -1104,31 +1113,12 @@ int main(int argc, char **argv_orig, char **envp) {
 
     }
 
-    u8 *use_dir = ".";
-
-    if (access(use_dir, R_OK | W_OK | X_OK)) {
-
-      use_dir = get_afl_env("TMPDIR");
-      if (!use_dir) { use_dir = "/tmp"; }
-
-    }
-
-    stdin_file = at_file ? strdup(at_file)
-                         : (char *)alloc_printf("%s/.afl-showmap-temp-%u",
-                                                use_dir, (u32)getpid());
-    unlink(stdin_file);
     atexit(at_exit_handler);
     fsrv->out_file = stdin_file;
     fsrv->out_fd =
         open(stdin_file, O_RDWR | O_CREAT | O_EXCL, DEFAULT_PERMISSION);
     if (fsrv->out_fd < 0) { PFATAL("Unable to create '%s'", out_file); }
 
-    if (arg_offset && use_argv[arg_offset] != stdin_file) {
-
-      use_argv[arg_offset] = strdup(stdin_file);
-
-    }
-
     if (get_afl_env("AFL_DEBUG")) {
 
       int j = optind;