diff --git a/.gitignore b/.gitignore new file mode 100644 index 00000000..0b386144 --- /dev/null +++ b/.gitignore @@ -0,0 +1,2 @@ +# Ignore the compiled neuzz binary. +neuzz diff --git a/README.md b/README.md index 2ff84d18..ed71eb07 100644 --- a/README.md +++ b/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. diff --git a/neuzz.c b/neuzz.c index 35e33c62..0d621780 100644 --- a/neuzz.c +++ b/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,6 +1790,7 @@ 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); } @@ -1789,6 +1798,7 @@ void copy_file(char* src, char* dst){ 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; } diff --git a/nn.py b/nn.py index fe0dcaec..5761eac9 100644 --- a/nn.py +++ b/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()