Allow ASAN binaries and add some quality-of-life improvements

.gitignore

@@ -0,0 +1,2 @@
+# Ignore the compiled neuzz binary.
+neuzz

README.md

@@ -26,6 +26,23 @@ If you want to try NEUZZ on a new program,
 1. Compile the new program from source code using afl-gcc.
 2. Collect the training data by running AFL on the binary for a while(about an hour), then copy the queue folder to neuzz_in.
 3. Follow the above two steps to start NN module and NEUZZ module.
+
+### Running with ASAN
+
+If your binary is compiled with ASAN instrumentation, do the following to run
+it properly.
+
+Pass `--enable-asan` to `nn.py`:
+```bash
+    python nn.py --enable-asan ./readelf -a
+```
+
+And pass `-m none` to `./nuezz` as you would to afl:
+```bash
+    ./neuzz -m none -i neuzz_in -o seeds -l 7506 ./readelf -a @@
+```
+
+
 ## Sample programs
 Try 10 real-world programs on NEUZZ. Check setup details at programs/[program names]/README.
 

neuzz.c

@@ -342,8 +342,7 @@ void init_forkserver(char** argv) {
   int st_pipe[2], ctl_pipe[2];
   int status;
   int rlen;
-  char* cwd = getcwd(NULL, 0);
-  out_file = alloc_printf("%s/%s/.cur_input",cwd, out_dir);
+  out_file = alloc_printf("%s/.cur_input", out_dir);
   printf("Spinning up the fork server...\n");
 
   if (pipe(st_pipe) || pipe(ctl_pipe)) perror("pipe() failed");
@@ -1690,9 +1689,18 @@ void dry_run(char* dir, int stage){
     struct dirent *entry;
     struct stat statbuf;
     if((dp = opendir(dir)) == NULL) {
+        perror("opendir");
         fprintf(stderr,"cannot open directory: %s\n", dir);
         return;
     }
+    // Store the original cwd so we can chdir back to it.
+    char original_cwd[PATH_MAX];
+    if (getcwd(original_cwd, sizeof(original_cwd)) == NULL) {
+      perror("getcwd");
+      return;
+    }
+
+
     if(chdir(dir)== -1)
         perror("chdir failed\n");
     int cnt = 0;
@@ -1755,7 +1763,7 @@ void dry_run(char* dir, int stage){
             }
         }
     }
-    if(chdir("..") == -1)
+    if(chdir(original_cwd) == -1)
         perror("chdir failed\n");
     closedir(dp);
 
@@ -1782,13 +1790,15 @@ void copy_file(char* src, char* dst){
     fptr1 = fopen(src, "r");
     if (fptr1 == NULL)
     {
+        perror("fopen");
         printf("Cannot open file %s \n", src);
         exit(0);
     }
 
     fptr2 = fopen(dst, "w");
     if (fptr2 == NULL)
     {
+        perror("fopen");
         printf("Cannot open file %s \n", dst);
         exit(0);
     }
@@ -1984,9 +1994,9 @@ void start_fuzz_test(int f_len){
 }
 
 
-void main(int argc, char*argv[]){
+int main(int argc, char*argv[]){
     int opt;
-    while ((opt = getopt(argc, argv, "+i:o:l:")) > 0)
+    while ((opt = getopt(argc, argv, "+i:o:l:m:")) > 0)
 
     switch (opt) {
 
@@ -2022,9 +2032,42 @@ void main(int argc, char*argv[]){
          printf("num_index %d %d small %d medium %d large %d\n", num_index[12], num_index[13], havoc_blk_small, havoc_blk_medium, havoc_blk_large);
          printf("mutation len: %ld\n", len);
          break;
+
+      case 'm': /* memory limit */
+          if (!strcmp(optarg, "none")) {
+            mem_limit = 0;
+            break;
+          }
+
+          char suffix = 'M';
+          if (sscanf(optarg, "%llu%c", &mem_limit, &suffix) < 1 || optarg[0] == '-') {
+            fprintf(stderr, "Bad syntax used for -m\n");
+            return -1;
+          }
+
+          switch (suffix) {
+            case 'T': mem_limit *= 1024 * 1024; break;
+            case 'G': mem_limit *= 1024; break;
+            case 'k': mem_limit /= 1024; break;
+            case 'M': break;
+            default:
+              fprintf(stderr, "Unsupported suffix or bad syntax for -m\n");
+              return -1;
+          }
+
+          if (mem_limit < 5) {
+            fprintf(stderr, "Dangerously low value of -m\n");
+            return -1;
+          }
+          if (sizeof(rlim_t) == 4 && mem_limit > 2000) {
+            fprintf(stderr, "Value of -m out of range on 32-bit systems\n");
+            return -1;
+          }
+          break;
 
     default:
         printf("no manual...");
+        return 0;
     }
 
     setup_signal_handlers();
@@ -2043,6 +2086,6 @@ void main(int argc, char*argv[]){
 
     start_fuzz(len);   
     printf("total execs %ld edge coverage %d.\n", total_execs, count_non_255_bytes(virgin_bits));
-    return;
+    return 0;
 }
 

nn.py

@@ -1,3 +1,4 @@
+import argparse
 import os
 import sys
 import glob
@@ -33,6 +34,7 @@
 SPLIT_RATIO = len(seed_list)
 # get binary argv
 argvv = sys.argv[1:]
+args = None
 
 # process training data from afl raw data
 
@@ -44,20 +46,21 @@ def process_data():
     global seed_list
     global new_seeds
 
+    output_folder = os.path.abspath(args.output_folder)
     # shuffle training samples
-    seed_list = glob.glob('./seeds/*')
+    seed_list = glob.glob('{}/*'.format(output_folder))
     seed_list.sort()
     SPLIT_RATIO = len(seed_list)
     rand_index = np.arange(SPLIT_RATIO)
     np.random.shuffle(seed_list)
-    new_seeds = glob.glob('./seeds/id_*')
+    new_seeds = glob.glob('{}/id_*'.format(output_folder))
 
     call = subprocess.check_output
 
     # get MAX_FILE_SIZE
     cwd = os.getcwd()
-    max_file_name = call(['ls', '-S', cwd + '/seeds/']).decode('utf8').split('\n')[0].rstrip('\n')
-    MAX_FILE_SIZE = os.path.getsize(cwd + '/seeds/' + max_file_name)
+    max_file_name = call(['ls', '-S', output_folder]).decode('utf8').split('\n')[0].rstrip('\n')
+    MAX_FILE_SIZE = os.path.getsize(output_folder + '/' + max_file_name)
 
     # create directories to save label, spliced seeds, variant length seeds, crashes and mutated seeds.
     if os.path.isdir("./bitmaps/") == False:
@@ -76,13 +79,14 @@ def process_data():
     for f in seed_list:
         tmp_list = []
         try:
+            mem_limit = '512' if not args.enable_asan else 'none'
             # append "-o tmp_file" to strip's arguments to avoid tampering tested binary.
             if argvv[0] == './strip':
-                out = call(['./afl-showmap', '-q', '-e', '-o', '/dev/stdout', '-m', '512', '-t', '500'] + argvv + [f] + ['-o', 'tmp_file'])
+                out = call(['./afl-showmap', '-q', '-e', '-o', '/dev/stdout', '-m', mem_limit, '-t', '500'] + args.target + [f] + ['-o', 'tmp_file'])
             else:
-                out = call(['./afl-showmap', '-q', '-e', '-o', '/dev/stdout', '-m', '512', '-t', '500'] + argvv + [f])
-        except subprocess.CalledProcessError:
-            print("find a crash")
+                out = call(['./afl-showmap', '-q', '-e', '-o', '/dev/stdout', '-m', mem_limit, '-t', '500'] + args.target + [f])
+        except subprocess.CalledProcessError as e:
+            print("find a crash", e)
         for line in out.splitlines():
             edge = line.split(b':')[0]
             tmp_cnt.append(edge)
@@ -399,6 +403,8 @@ def gen_grad(data):
 
 def setup_server():
     sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    # Allow re-use so we don't have to wait if the process crashes.
+    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
     sock.bind((HOST, PORT))
     sock.listen(1)
     conn, addr = sock.accept()
@@ -415,4 +421,20 @@ def setup_server():
     conn.close()
 
 
-setup_server()
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="""Runs the background machine
+                        learning process for Neuzz.""")
+    parser.add_argument('-e',
+                        '--enable-asan',
+                        help='Enable ASAN (runs afl-showmap with -m none)',
+                        default=False,
+                        action='store_true')
+    parser.add_argument('-o',
+                        '--output-folder',
+                        help="""The -o folder provided to afl, this is where the
+                        Neuzz process reads seeds from.""",
+                        default='seeds')
+    parser.add_argument('target', nargs=argparse.REMAINDER)
+    args = parser.parse_args()
+
+    setup_server()