zax_driver.cpp

// @HEADER_LANG C++ 
// @HEADER_COMPONENT zafl
// @HEADER_BEGIN
/*
* Copyright (c) 2018-2019 Zephyr Software LLC
*
* This file may be used and modified for non-commercial purposes as long as
* all copyright, permission, and nonwarranty notices are preserved.
* Redistribution is prohibited without prior written consent from Zephyr
* Software.
*
* Please contact the authors for restrictions applying to commercial use.
*
* THIS SOURCE IS PROVIDED "AS IS" AND WITHOUT ANY EXPRESS OR IMPLIED
* WARRANTIES, INCLUDING, WITHOUT LIMITATION, THE IMPLIED WARRANTIES OF
* MERCHANTIBILITY AND FITNESS FOR A PARTICULAR PURPOSE.
*
* Author: Zephyr Software
* e-mail: jwd@zephyr-software.com
* URL   : http://www.zephyr-software.com/
*
* This software was developed with SBIR funding and is subject to SBIR Data Rights, 
* as detailed below.
*
* SBIR DATA RIGHTS
*
* Contract No. __FA8750-17-C-0295___________________________.
* Contractor Name __Zephyr Software LLC_____________________.
* Address __4826 Stony Point Rd, Barboursville, VA 22923____.
*
*/

// @HEADER_END

#include <stdlib.h>
#include <getopt.h>
#include <unistd.h>

#include "zax.hpp"
#include "critical_edge_breaker.hpp"
#include "loop_count.hpp"
#include "zuntracer.hpp"

using namespace std;
using namespace IRDB_SDK;
using namespace Zafl;

static void usage(char* name)
{
	cerr<<"Usage: "<<name<<" <variant_id>\n";
	cerr<<"\t[--verbose | -v]                               Verbose mode                  "<<endl;
	cerr<<"\t[--help,--usage,-?,-h]                         Display this message          "<<endl;
	cerr<<"\t[--entrypoint|-e {<funcName>|<hex_address>}]   Specify where to insert fork server (defaults to main if found)"<<endl;
	cerr<<"\t[--exitpoint|-E {<funcName>|<hex_address>}]    Specify where to insert exit"<<endl;
	cerr<<"\t[--whitelist|-w whitelistFile]                 Specify list of functions/addresses to instrument"<<endl;
	cerr<<"\t[--blacklist|-b blacklistFile]                 Specify list of functions/addresses to omit"<<endl;
	cerr<<"\t[--fixed-map-address <addr>]                   Enable fixed-mapping addressing and Specify the address to use for the map [Default 0x10000]"<<endl;
	cerr<<"\t[--no-fixed-map]                               Disablw fixed-mapping addressing"<<endl;
	cerr<<"\t[--stars|-s]                                   Enable STARS optimizations    "<<endl;
	cerr<<"\t[--enable-bb-graph-optimization|-g]            Elide instrumentation if basic block has 1 successor"<<endl;
	cerr<<"\t[--disable-bb-graph-optimization|-G]           Elide instrumentation if basic block has 1 successor"<<endl;
	cerr<<"\t[--autozafl|-a]                                Auto-initialize fork server (incompatible with --entrypoint)"<<endl;
	cerr<<"\t[--untracer|-u]                                Untracer-style block-coverage instrumentation"<<endl;
	cerr<<"\t[--break-critical-edge-style|-c]               Breakup critical edges"<<endl;
	cerr<<"\t[--enable-loop-count-instr|-j]                 Insert instrumentation to do afl-style loop counting for zuntracer."<<endl;
	cerr<<"\t[--disable-loop-count-instr|-J]                Do not do -j (default)"<<endl;
	cerr<<"\t[--enable-floating-instrumentation|-i]         Select best instrumentation within basic block"<<endl;
	cerr<<"\t[--disable-floating-instrumentation|-I]        Instrument first instruction in basic blocks"<<endl;
	cerr<<"\t[--enable-context-sensitivity <style>]         Use calling context sensitivity, style={callsite,function}"<<endl;
	cerr<<"\t[--disable-context-sensitivity]                Disable calling context sensitivity"<<endl;
	cerr<<"\t[--random-seed|r <value>]                      Specify random seed"<<endl;
}

int main(int argc, char **argv)
{
	if(argc < 2)
	{
		usage(argv[0]);
		exit(1);
	}

	string programName(argv[0]);
	auto entry_fork_server = string();
	auto variantID = atoi(argv[1]);
	auto verbose=false;
	auto use_stars=false;
	auto autozafl=false;
	auto whitelistFile=string();
	auto blacklistFile=string();
	auto bb_graph_optimize=false;
	auto domgraph_optimize=false;
	auto forkserver_enabled=true;
	auto untracer_mode=false;
	auto break_critical_edge_style=bceNone;
	auto do_loop_count_instr=false;
	auto do_keep_loop_headers_only=false;
	auto loop_count_buckets=string("1,2,3,4,8,16,32,128");
	auto floating_instrumentation=false;
	auto context_sensitivity=ContextSensitivity_None;
	auto random_seed = 0U;
	auto fixed_map_address=VirtualOffset_t(0x10000);
	set<string> exitpoints;

	srand(getpid()+time(NULL));

	// Parse some options for the transform
	static struct option long_options[] = {
		{"verbose",                          no_argument,       0, 'v'},
		{"help",                             no_argument,       0, 'h'},
		{"usage",                            no_argument,       0, '?'},
		{"stars",                            no_argument,       0, 's'},
		{"entrypoint",                       required_argument, 0, 'e'},
		{"exitpoint",                        required_argument, 0, 'E'},
		{"whitelist",                        required_argument, 0, 'w'},
		{"blacklist",                        required_argument, 0, 'b'},
		{"fixed-map-address",                required_argument, 0, 'm'},
		{"no-fixed-map",                     no_argument,       0, 'M'},
		{"autozafl",                         no_argument,       0, 'a'},
		{"enable-bb-graph-optimization",     no_argument,       0, 'g'},
		{"disable-bb-graph-optimization",    no_argument,       0, 'G'},
		{"enable-domgraph-optimization",     no_argument,       0, 'd'},
		{"disable-domgraph-optimization",    no_argument,       0, 'D'},
		{"enable-forkserver",                no_argument,       0, 'f'},
		{"disable-forkserver",               no_argument,       0, 'F'},
		{"untracer",                         no_argument,       0, 'u'},
		{"break-critical-edge",     	     required_argument, 0, 'c'},
		{"enable-loop-count-instr",          no_argument,       0, 'j'},
		{"disable-loop-count-instr",         no_argument,       0, 'J'},
		{"keep-loop-headers-only",           no_argument,       0, 'k'},
		{"loop-count-buckets",               required_argument, 0, ':'},
		{"enable-floating-instrumentation",  no_argument,       0, 'i'},
		{"disable-floating-instrumentation", no_argument,       0, 'I'},
		{"enable-context-sensitivity",       required_argument, 0, 'z'},
		{"random-seed",                      required_argument, 0, 'r'},
		{0,0,0,0}
	};
	const char* short_opts="r:z:e:E:w:sv?hagGdDfFuc:jJkiIm:M";

	while(true)
	{
		int index = 0;
		int c = getopt_long(argc, argv, short_opts, long_options, &index);
		if(c == -1)
			break;
		switch(c)
		{
			case 'v':
				verbose=true;
				break;
			case '?':
			case 'h':
				usage(argv[0]);
				exit(1);
				break;
			case 's':
				use_stars=true;
				cout << "STARS optimization enabled" << endl;
				break;
			case 'e':
				entry_fork_server = optarg;
				break;
			case 'E':
				exitpoints.insert(optarg);
				break;
			case 'w':
				whitelistFile=optarg;
				break;
			case 'b':
				blacklistFile=optarg;
				break;
			case 'm':
				fixed_map_address = strtoul(optarg, nullptr, 0);
				break;
			case 'M':
				fixed_map_address = 0;
				break;
			case 'a':
				autozafl=true;
				break;
			case 'g':
				bb_graph_optimize=true;
				break;
			case 'G':
				bb_graph_optimize=false;
				break;
			case 'd':
				domgraph_optimize=true;
				break;
			case 'D':
				domgraph_optimize=false;
				break;
			case 'f':
				forkserver_enabled=true;
				break;
			case 'F':
				forkserver_enabled=false;
				break;
			case 'u':
				untracer_mode=true;
				break;
			case 'c':
				if(optarg)
				{
					if(string(optarg)=="all")
						break_critical_edge_style=bceAll;
					else if(string(optarg)=="none")
						break_critical_edge_style=bceNone;
					else if(string(optarg)=="targets")
						break_critical_edge_style=bceTargets;
					else if(string(optarg)=="target")
						break_critical_edge_style=bceTargets;
					else if(string(optarg)=="fallthroughs")
						break_critical_edge_style=bceFallthroughs;
					else if(string(optarg)=="fallthrough")
						break_critical_edge_style=bceFallthroughs;
					else
					{
						cout << "Cannot map " << optarg << "to {all,none,targets,fallthrough}" << endl;
						exit(2);
					}
				}
				else
					break_critical_edge_style=bceAll;
				break;
			case 'j':
				do_loop_count_instr=true;
				break;
			case 'J':
				do_loop_count_instr=false;
				break;
			case 'k':
				do_keep_loop_headers_only=true;
				break;
			case ':':
				loop_count_buckets = string(optarg);
				break;
			case 'i':
				floating_instrumentation=true;
				break;
			case 'I':
				floating_instrumentation=false;
				break;
			case 'z':
				if (optarg == string("callsite"))
					context_sensitivity=ContextSensitivity_Callsite; // Angora fuzzer style
				else if (optarg == string("function"))
					context_sensitivity=ContextSensitivity_Function;
				else
					context_sensitivity=ContextSensitivity_None;
				break;
			case 'r':
				random_seed = strtoul(optarg, NULL, 0);
				srand(random_seed);
				cout << "Setting random seed to: " << random_seed << endl;
				break;
			default:
				break;
		}
	}

	if (entry_fork_server.size() > 0 && autozafl)
	{
		cerr << "--entrypoint and --autozafl are incompatible options" << endl;
		usage(argv[0]);
		exit(1);
	}

	if (floating_instrumentation && !use_stars)
	{
		cerr << "STARS must be turned on when using floating instrumentation" << endl;
		usage(argv[0]);
		exit(1);
	}


	/* setup the interface to the sql server */
	auto pqxx_interface=pqxxDB_t::factory();
	BaseObj_t::setInterface(pqxx_interface.get());

	auto pidp=VariantID_t::factory(variantID);
	assert(pidp->isRegistered()==true);

	bool one_success = false;
	for(set<File_t*>::iterator it=pidp->getFiles().begin();
	        it!=pidp->getFiles().end();
	        ++it)
	{
		File_t* this_file = *it;
		auto firp = FileIR_t::factory(pidp.get(), this_file);

		cout<<"Transforming "<<this_file->getURL()<<endl;

		assert(firp && pidp);

		try
		{

			if (do_loop_count_instr)
			{
				ZedgeNS::Zedge_t(firp.get(), loop_count_buckets).execute();
				firp->setBaseIDS();
				firp->assembleRegistry();
			}
			if (break_critical_edge_style != bceNone)
			{
				CriticalEdgeBreaker_t ceb(firp.get(), {}, break_critical_edge_style, verbose);
				cout << "#ATTRIBUTE num_bb_extra_blocks=" << ceb.getNumberExtraNodes() << endl;
				firp->setBaseIDS();
				firp->assembleRegistry();
			}

			const auto zax = unique_ptr<ZaxBase_t>( 
					untracer_mode ? static_cast<ZaxBase_t*>(new ZUntracer_t(*pqxx_interface, firp.get(), entry_fork_server, exitpoints, use_stars, autozafl)) : 
					                static_cast<ZaxBase_t*>(new Zax_t      (*pqxx_interface, firp.get(), entry_fork_server, exitpoints, use_stars, autozafl))); 

			if (whitelistFile.size()>0)
				zax->setWhitelist(whitelistFile);
			if (blacklistFile.size()>0)
				zax->setBlacklist(blacklistFile);

			zax->setVerbose(verbose);
			zax->setBasicBlockOptimization(bb_graph_optimize);
			zax->setDomgraphOptimization(domgraph_optimize);
			zax->setBasicBlockFloatingInstrumentation(floating_instrumentation);
			zax->setEnableForkServer(forkserver_enabled);
//		 	zax->setBreakCriticalEdgeStyle(break_critical_edge_style);
			zax->setKeepLoopHeadersOnly(do_keep_loop_headers_only);
//			zax->setDoLoopCountInstrumentation(do_loop_count_instr);
//			zax->setLoopCountBuckets(loop_count_buckets);
			zax->setContextSensitivity(context_sensitivity); 
			zax->setFixedMapAddress(fixed_map_address);

			int success=zax->execute();

			if (success)
			{
				cout<<"Writing changes for "<<this_file->getURL()<<endl;
				one_success = true;
				firp->writeToDB();
			}
			else
			{
				cout<<"Skipping (no changes) "<<this_file->getURL()<<endl;
			}
		}
		catch (const DatabaseError_t &pnide)
		{
			cerr << programName << ": Unexpected database error: " << pnide << "file url: " << this_file->getURL() << endl;
			exit(1);
		}
		catch (const std::exception &e)
		{
			cerr << programName << ": Error: " << e.what() << endl;
			exit(1);
		}
		catch (...)
		{
			cerr << programName << ": Unexpected error file url: " << this_file->getURL() << endl;
			exit(1);
		}
	} // end file iterator

	// if any transforms for any files succeeded, we commit
	if (one_success)
	{
		cout<<"Commiting changes...\n";
		pqxx_interface->commit();
	}

	return 0;
}