From 9918852e57c57809a58b7d4d6cf0a2475934d072 Mon Sep 17 00:00:00 2001 From: boehme Date: Mon, 14 Aug 2017 17:39:26 +0800 Subject: [PATCH] AFLGo initial commit --- Makefile | 2 +- afl-fuzz.c | 197 ++++++++++++- config.h | 4 + llvm_mode/afl-clang-fast.c | 10 +- llvm_mode/afl-llvm-pass.so.cc | 513 +++++++++++++++++++++++++++++++--- llvm_mode/afl-llvm-rt.o.c | 178 +++++++++++- scripts/add_edges.py | 61 ++++ scripts/distance.py | 149 ++++++++++ scripts/genDistance.sh | 150 ++++++++++ 9 files changed, 1214 insertions(+), 50 deletions(-) create mode 100755 scripts/add_edges.py create mode 100755 scripts/distance.py create mode 100755 scripts/genDistance.sh diff --git a/Makefile b/Makefile index 44d1ffa2..15f5c71d 100644 --- a/Makefile +++ b/Makefile @@ -33,7 +33,7 @@ CFLAGS += -Wall -D_FORTIFY_SOURCE=2 -g -Wno-pointer-sign \ -DBIN_PATH=\"$(BIN_PATH)\" ifneq "$(filter Linux GNU%,$(shell uname))" "" - LDFLAGS += -ldl + LDFLAGS += -ldl -lm endif ifeq "$(findstring clang, $(shell $(CC) --version 2>/dev/null))" "" diff --git a/afl-fuzz.c b/afl-fuzz.c index c7c00cd6..aefe19a3 100644 --- a/afl-fuzz.c +++ b/afl-fuzz.c @@ -56,6 +56,8 @@ #include #include +#include + #if defined(__APPLE__) || defined(__FreeBSD__) || defined (__OpenBSD__) # include #endif /* __APPLE__ || __FreeBSD__ || __OpenBSD__ */ @@ -98,6 +100,14 @@ EXP_ST u64 mem_limit = MEM_LIMIT; /* Memory cap for child (MB) */ static u32 stats_update_freq = 1; /* Stats update frequency (execs) */ +static u8 cooling_schedule = 0; /* Cooling schedule for directed fuzzing */ +enum { + /* 00 */ SAN_EXP, /* Exponential schedule */ + /* 01 */ SAN_LOG, /* Logarithmical schedule */ + /* 02 */ SAN_LIN, /* Linear schedule */ + /* 03 */ SAN_QUAD /* Quadratic schedule */ +}; + EXP_ST u8 skip_deterministic, /* Skip deterministic stages? */ force_deterministic, /* Force deterministic stages? */ use_splicing, /* Recombine input files? */ @@ -245,6 +255,8 @@ struct queue_entry { u8* trace_mini; /* Trace bytes, if kept */ u32 tc_ref; /* Trace bytes ref count */ + double distance; /* Distance to targets */ + struct queue_entry *next, /* Next element, if any */ *next_100; /* 100 elements ahead */ @@ -270,6 +282,11 @@ static u32 extras_cnt; /* Total number of tokens read */ static struct extra_data* a_extras; /* Automatically selected extras */ static u32 a_extras_cnt; /* Total number of tokens available */ +static double cur_distance = -1.0; /* Distance of executed input */ +static double max_distance = -1.0; /* Maximal distance for any input */ +static double min_distance = -1.0; /* Minimal distance for any input */ +static u32 t_x = 10; /* Time to exploitation (Default: 10 min) */ + static u8* (*post_handler)(u8* buf, u32* len); /* Interesting values, as per config.h */ @@ -780,6 +797,18 @@ static void add_to_queue(u8* fname, u32 len, u8 passed_det) { q->depth = cur_depth + 1; q->passed_det = passed_det; + q->distance = cur_distance; + if (cur_distance > 0) { + + if (max_distance <= 0) { + max_distance = cur_distance; + min_distance = cur_distance; + } + if (cur_distance > max_distance) max_distance = cur_distance; + if (cur_distance < min_distance) min_distance = cur_distance; + + } + if (q->depth > max_depth) max_depth = q->depth; if (queue_top) { @@ -882,6 +911,15 @@ static inline u8 has_new_bits(u8* virgin_map) { u32 i = (MAP_SIZE >> 3); + /* Calculate distance of current input to targets */ + u64* total_distance = (u64*) (trace_bits + MAP_SIZE); + u64* total_count = (u64*) (trace_bits + MAP_SIZE + 8); + + if (*total_count > 0) + cur_distance = (double) (*total_distance) / (double) (*total_count); + else + cur_distance = -1.0; + #else u32* current = (u32*)trace_bits; @@ -889,6 +927,15 @@ static inline u8 has_new_bits(u8* virgin_map) { u32 i = (MAP_SIZE >> 2); + /* Calculate distance of current input to targets */ + u32* total_distance = (u32*)(trace_bits + MAP_SIZE); + u32* total_count = (u32*)(trace_bits + MAP_SIZE + 4); + + if (*total_count > 0) { + cur_distance = (double) (*total_distance) / (double) (*total_count); + else + cur_distance = -1.0; + #endif /* ^__x86_64__ */ u8 ret = 0; @@ -1346,7 +1393,8 @@ EXP_ST void setup_shm(void) { memset(virgin_tmout, 255, MAP_SIZE); memset(virgin_crash, 255, MAP_SIZE); - shm_id = shmget(IPC_PRIVATE, MAP_SIZE, IPC_CREAT | IPC_EXCL | 0600); + /* Allocate 24 byte more for distance info */ + shm_id = shmget(IPC_PRIVATE, MAP_SIZE + 16, IPC_CREAT | IPC_EXCL | 0600); if (shm_id < 0) PFATAL("shmget() failed"); @@ -2270,7 +2318,7 @@ static u8 run_target(char** argv, u32 timeout) { must prevent any earlier operations from venturing into that territory. */ - memset(trace_bits, 0, MAP_SIZE); + memset(trace_bits, 0, MAP_SIZE + 16); MEM_BARRIER(); /* If we're running in "dumb" mode, we can't rely on the fork server @@ -2582,6 +2630,27 @@ static u8 calibrate_case(char** argv, struct queue_entry* q, u8* use_mem, cksum = hash32(trace_bits, MAP_SIZE, HASH_CONST); + /* This is relevant when test cases are added w/out save_if_interesting */ + + if (q->distance <= 0) { + + /* This calculates cur_distance */ + has_new_bits(virgin_bits); + + q->distance = cur_distance; + if (cur_distance > 0) { + + if (max_distance <= 0) { + max_distance = cur_distance; + min_distance = cur_distance; + } + if (cur_distance > max_distance) max_distance = cur_distance; + if (cur_distance < min_distance) min_distance = cur_distance; + + } + + } + if (q->exec_cksum != cksum) { u8 hnb = has_new_bits(virgin_bits); @@ -4713,10 +4782,68 @@ static u32 calculate_score(struct queue_entry* q) { } + u64 cur_ms = get_cur_time(); + u64 t = (cur_ms - start_time) / 1000; + double progress_to_tx = ((double) t) / ((double) t_x * 60.0); + + double T; + + //TODO Substitute functions of exp and log with faster bitwise operations on integers + switch (cooling_schedule) { + case SAN_EXP: + + T = 1.0 / pow(20.0, progress_to_tx); + + break; + + case SAN_LOG: + + // alpha = 2 and exp(19/2) - 1 = 13358.7268297 + T = 1.0 / (1.0 + 2.0 * log(1.0 + progress_to_tx * 13358.7268297)); + + break; + + case SAN_LIN: + + T = 1.0 / (1.0 + 19.0 * progress_to_tx); + + break; + + case SAN_QUAD: + + T = 1.0 / (1.0 + 19.0 * pow(progress_to_tx, 2)); + + break; + + default: + PFATAL ("Unkown Power Schedule for Directed Fuzzing"); + } + + double power_factor = 1.0; + if (q->distance > 0) { + + double normalized_d = q->distance; + if (max_distance != min_distance) + normalized_d = (q->distance - min_distance) / (max_distance - min_distance); + + if (normalized_d >= 0) { + + double p = (1.0 - normalized_d) * (1.0 - T) + 0.5 * T; + power_factor = pow(2.0, 2.0 * (double) log2(MAX_FACTOR) * (p - 0.5)); + + }// else WARNF ("Normalized distance negative: %f", normalized_d); + + } + + perf_score *= power_factor; + /* Make sure that we don't go over limit. */ if (perf_score > HAVOC_MAX_MULT * 100) perf_score = HAVOC_MAX_MULT * 100; + /* AFLGO-DEBUGGING */ + // fprintf(stderr, "[Time %llu] q->distance: %4lf, max_distance: %4lf min_distance: %4lf, T: %4.3lf, power_factor: %4.3lf, adjusted perf_score: %4d\n", t, q->distance, max_distance, min_distance, T, power_factor, perf_score); + return perf_score; } @@ -7023,6 +7150,13 @@ static void usage(u8* argv0) { " -i dir - input directory with test cases\n" " -o dir - output directory for fuzzer findings\n\n" + "Directed fuzzing specific settings:\n\n" + + " -z schedule - temperature-based power schedules\n" + " {exp, log, lin, quad} (Default: exp)\n" + " -c min - time from start when SA enters exploitation\n" + " in secs (s), mins (m), hrs (h), or days (d)\n\n" + "Execution control settings:\n\n" " -f file - location read by the fuzzed program (stdin)\n" @@ -7665,9 +7799,17 @@ static void save_cmdline(u32 argc, char** argv) { } - #ifndef AFL_LIB +int stricmp(char const *a, char const *b) { + int d; + for (;; a++, b++) { + d = tolower(*a) - tolower(*b); + if (d != 0 || !*a) + return d; + } +} + /* Main entry point */ int main(int argc, char** argv) { @@ -7683,14 +7825,14 @@ int main(int argc, char** argv) { struct timeval tv; struct timezone tz; - SAYF(cCYA "afl-fuzz " cBRI VERSION cRST " by \n"); + SAYF(cCYA "aflgo (yeah!) " cBRI VERSION cRST "\n"); doc_path = access(DOC_PATH, F_OK) ? "docs" : DOC_PATH; gettimeofday(&tv, &tz); srandom(tv.tv_sec ^ tv.tv_usec ^ getpid()); - while ((opt = getopt(argc, argv, "+i:o:f:m:t:T:dnCB:S:M:x:Q")) > 0) + while ((opt = getopt(argc, argv, "+i:o:f:m:t:T:dnCB:S:M:x:Qz:c:")) > 0) switch (opt) { @@ -7858,6 +8000,43 @@ int main(int argc, char** argv) { break; + case 'z': /* Cooling schedule for Directed Fuzzing */ + + if (!stricmp(optarg, "exp")) + cooling_schedule = SAN_EXP; + else if (!stricmp(optarg, "log")) + cooling_schedule = SAN_LOG; + else if (!stricmp(optarg, "lin")) + cooling_schedule = SAN_LIN; + else if (!stricmp(optarg, "quad")) + cooling_schedule = SAN_QUAD; + else + PFATAL ("Unknown value for option -z"); + + break; + + case 'c': { /* cut-off time for cooling schedule */ + + u8 suffix = 'm'; + + if (sscanf(optarg, "%u%c", &t_x, &suffix) < 1 || + optarg[0] == '-') FATAL("Bad syntax used for -c"); + + switch (suffix) { + + case 's': t_x /= 60; break; + case 'm': break; + case 'h': t_x *= 60; break; + case 'd': t_x *= 60 * 24; break; + + default: FATAL("Unsupported suffix or bad syntax for -c"); + + } + + } + + break; + default: usage(argv[0]); @@ -7866,6 +8045,14 @@ int main(int argc, char** argv) { if (optind == argc || !in_dir || !out_dir) usage(argv[0]); + OKF("Running with " cBRI "%s" cRST " schedule and time-to-exploitation set to " cBRI "%d minutes" cRST, + cooling_schedule == SAN_EXP ? "EXP" : + cooling_schedule == SAN_LOG ? "LOG" : + cooling_schedule == SAN_LIN ? "LIN" : + cooling_schedule == SAN_QUAD ? "QUAD" : "???", + t_x + ); + setup_signal_handlers(); check_asan_opts(); diff --git a/config.h b/config.h index 4ba09a23..5ea17001 100644 --- a/config.h +++ b/config.h @@ -347,4 +347,8 @@ // #define IGNORE_FINDS +#define MAX_FACTOR 32 + +//#define AFLGO_TRACING + #endif /* ! _HAVE_CONFIG_H */ diff --git a/llvm_mode/afl-clang-fast.c b/llvm_mode/afl-clang-fast.c index d4202a63..9fd8ea40 100644 --- a/llvm_mode/afl-clang-fast.c +++ b/llvm_mode/afl-clang-fast.c @@ -123,6 +123,7 @@ static void edit_params(u32 argc, char** argv) { cc_params[cc_par_cnt++] = "-fsanitize-coverage=trace-pc-guard"; cc_params[cc_par_cnt++] = "-mllvm"; cc_params[cc_par_cnt++] = "-sanitizer-coverage-block-threshold=0"; + WARNF("Disabling AFLGO features..\n"); #else cc_params[cc_par_cnt++] = "-Xclang"; cc_params[cc_par_cnt++] = "-load"; @@ -139,6 +140,11 @@ static void edit_params(u32 argc, char** argv) { while (--argc) { u8* cur = *(++argv); + if (!strncmp(cur, "-distance", 9) + || !strncmp(cur, "-targets", 8) + || !strncmp(cur, "-outdir", 7)) + cc_params[cc_par_cnt++] = "-mllvm"; + if (!strcmp(cur, "-m32")) bit_mode = 32; if (!strcmp(cur, "-m64")) bit_mode = 64; @@ -316,9 +322,9 @@ int main(int argc, char** argv) { if (isatty(2) && !getenv("AFL_QUIET")) { #ifdef USE_TRACE_PC - SAYF(cCYA "afl-clang-fast [tpcg] " cBRI VERSION cRST " by \n"); + SAYF(cCYA "aflgo-compiler (yeah!) [tpcg] " cBRI VERSION cRST "\n"); #else - SAYF(cCYA "afl-clang-fast " cBRI VERSION cRST " by \n"); + SAYF(cCYA "aflgo-compiler (yeah!) " cBRI VERSION cRST "\n"); #endif /* ^USE_TRACE_PC */ } diff --git a/llvm_mode/afl-llvm-pass.so.cc b/llvm_mode/afl-llvm-pass.so.cc index 04f3ec48..921f949a 100644 --- a/llvm_mode/afl-llvm-pass.so.cc +++ b/llvm_mode/afl-llvm-pass.so.cc @@ -31,15 +31,53 @@ #include #include +#include +#include +#include +#include +#include +#include +#include +#include + #include "llvm/ADT/Statistic.h" #include "llvm/IR/IRBuilder.h" #include "llvm/IR/LegacyPassManager.h" #include "llvm/IR/Module.h" #include "llvm/Support/Debug.h" #include "llvm/Transforms/IPO/PassManagerBuilder.h" +#include "llvm/Support/CommandLine.h" +#include "llvm/Support/raw_ostream.h" +#include "llvm/Analysis/CFGPrinter.h" + +#if defined(LLVM34) +#include "llvm/DebugInfo.h" +#else +#include "llvm/IR/DebugInfo.h" +#endif + +#if defined(LLVM34) || defined(LLVM35) || defined(LLVM36) +#define LLVM_OLD_DEBUG_API +#endif using namespace llvm; +cl::opt DistanceFile( + "distance", + cl::desc("Distance file containing the distance of each basic block to the provided targets."), + cl::value_desc("filename") +); + +cl::opt TargetsFile( + "targets", + cl::desc("Input file containing the target lines of code."), + cl::value_desc("targets")); + +cl::opt OutDirectory( + "outdir", + cl::desc("Output directory where Ftargets.txt, Fnames.txt, and BBnames.txt are generated."), + cl::value_desc("outdir")); + namespace { class AFLCoverage : public ModulePass { @@ -51,24 +89,73 @@ namespace { bool runOnModule(Module &M) override; - // StringRef getPassName() const override { - // return "American Fuzzy Lop Instrumentation"; - // } - }; } - char AFLCoverage::ID = 0; - bool AFLCoverage::runOnModule(Module &M) { + bool is_aflgo = false; + bool is_aflgo_preprocessing = false; + + if (!TargetsFile.empty() && !DistanceFile.empty()) { + FATAL("Cannot specify both '-targets' and '-distance'!"); + return false; + } + + std::list targets; + std::map bb_to_dis; + std::vector basic_blocks; + + if (!TargetsFile.empty()) { + + if (OutDirectory.empty()) { + FATAL("Provide output directory '-outdir '"); + return false; + } + + std::ifstream targetsfile(TargetsFile); + std::string line; + while (std::getline(targetsfile, line)) + targets.push_back(line); + targetsfile.close(); + + is_aflgo_preprocessing = true; + + } else if (!DistanceFile.empty()) { + + std::ifstream cf (DistanceFile.c_str()); + if (cf.is_open()) { + + std::string line; + while (getline(cf,line)) { + + std::size_t pos = line.find(","); + std::string bb_name = line.substr(0, pos); + int bb_dis = (int) (100.0 * atof(line.substr(pos + 1, line.length()).c_str())); + + bb_to_dis.insert(std::pair(bb_name, bb_dis) ); + basic_blocks.push_back(bb_name); + + } + cf.close(); + + is_aflgo = true; + + } else { + FATAL("Unable to find %s.", DistanceFile.c_str()); + return false; + } + + } + LLVMContext &C = M.getContext(); IntegerType *Int8Ty = IntegerType::getInt8Ty(C); IntegerType *Int32Ty = IntegerType::getInt32Ty(C); + IntegerType *Int64Ty = IntegerType::getInt64Ty(C); /* Show a banner */ @@ -76,7 +163,12 @@ bool AFLCoverage::runOnModule(Module &M) { if (isatty(2) && !getenv("AFL_QUIET")) { - SAYF(cCYA "afl-llvm-pass " cBRI VERSION cRST " by \n"); + if (is_aflgo || is_aflgo_preprocessing) + SAYF(cCYA "aflgo-llvm-pass (yeah!) " cBRI VERSION cRST " (%s mode)\n", + (is_aflgo_preprocessing ? "preprocessing" : "distance instrumentation")); + else + SAYF(cCYA "afl-llvm-pass " cBRI VERSION cRST " by \n"); + } else be_quiet = 1; @@ -93,6 +185,24 @@ bool AFLCoverage::runOnModule(Module &M) { } + /* Default: Not selecitive */ + char* is_selective_str = getenv("AFLGO_SELECTIVE"); + unsigned int is_selective = 0; + + if (is_selective_str && sscanf(is_selective_str, "%u", &is_selective) != 1) + FATAL("Bad value of AFLGO_SELECTIVE (must be 0 or 1)"); + + char* dinst_ratio_str = getenv("AFLGO_INST_RATIO"); + unsigned int dinst_ratio = 100; + + if (dinst_ratio_str) { + + if (sscanf(dinst_ratio_str, "%u", &dinst_ratio) != 1 || !dinst_ratio || + dinst_ratio > 100) + FATAL("Bad value of AFLGO_INST_RATIO (must be between 1 and 100)"); + + } + /* Get globals for the SHM region and the previous location. Note that __afl_prev_loc is thread-local. */ @@ -108,60 +218,383 @@ bool AFLCoverage::runOnModule(Module &M) { int inst_blocks = 0; - for (auto &F : M) - for (auto &BB : F) { + if (is_aflgo_preprocessing) { + + std::ofstream bbnames; + std::ofstream bbcalls; + std::ofstream fnames; + std::ofstream ftargets; + struct stat sb; + + bbnames.open(OutDirectory + "/BBnames.txt", std::ofstream::out | std::ofstream::app); + bbcalls.open(OutDirectory + "/BBcalls.txt", std::ofstream::out | std::ofstream::app); + fnames.open(OutDirectory + "/Fnames.txt", std::ofstream::out | std::ofstream::app); + ftargets.open(OutDirectory + "/Ftargets.txt", std::ofstream::out | std::ofstream::app); + + /* Create dot-files directory */ + std::string dotfiles(OutDirectory + "/dot-files"); + if (stat(dotfiles.c_str(), &sb) != 0) { + const int dir_err = mkdir(dotfiles.c_str(), S_IRWXU | S_IRWXG | S_IROTH | S_IXOTH); + if (-1 == dir_err) + FATAL("Could not create directory %s.", dotfiles.c_str()); + } + + for (auto &F : M) { + + bool has_BBs = false; + std::string funcName = F.getName(); - BasicBlock::iterator IP = BB.getFirstInsertionPt(); - IRBuilder<> IRB(&(*IP)); + /* Black list of function names */ + std::vector blacklist = { + "asan.", + "llvm.", + "sancov.", + "free" + "malloc", + "calloc", + "realloc" + }; + for (std::vector::size_type i = 0; i < blacklist.size(); i++) + if (!funcName.compare(0, blacklist[i].size(), blacklist[i])) + continue; - if (AFL_R(100) >= inst_ratio) continue; + bool is_target = false; + for (auto &BB : F) { + + TerminatorInst *TI = BB.getTerminator(); + IRBuilder<> Builder(TI); + + std::string bb_name(""); + std::string filename; + unsigned line; + + for (auto &I : BB) { +#ifdef LLVM_OLD_DEBUG_API + DebugLoc Loc = I.getDebugLoc(); + if (!Loc.isUnknown()) { + + DILocation cDILoc(Loc.getAsMDNode(M.getContext())); + DILocation oDILoc = cDILoc.getOrigLocation(); + + line = oDILoc.getLineNumber(); + filename = oDILoc.getFilename().str(); + + if (filename.empty()) { + line = cDILoc.getLineNumber(); + filename = cDILoc.getFilename().str(); + } +#else + + if (DILocation *Loc = I.getDebugLoc()) { + line = Loc->getLine(); + filename = Loc->getFilename().str(); + + if (filename.empty()) { + DILocation *oDILoc = Loc->getInlinedAt(); + if (oDILoc) { + line = oDILoc->getLine(); + filename = oDILoc->getFilename().str(); + } + } + +#endif /* LLVM_OLD_DEBUG_API */ + + /* Don't worry about external libs */ + std::string Xlibs("/usr/"); + if (filename.empty() || line == 0 || !filename.compare(0, Xlibs.size(), Xlibs)) + continue; + + if (bb_name.empty()) { + + std::size_t found = filename.find_last_of("/\\"); + if (found != std::string::npos) + filename = filename.substr(found + 1); + + bb_name = filename + ":" + std::to_string(line); + + } + + if (!is_target) { + for (std::list::iterator it = targets.begin(); it != targets.end(); ++it) { + + std::string target = *it; + std::size_t found = target.find_last_of("/\\"); + if (found != std::string::npos) + target = target.substr(found + 1); + + std::size_t pos = target.find_last_of(":"); + std::string target_file = target.substr(0, pos); + unsigned int target_line = atoi(target.substr(pos + 1).c_str()); + + if (!target_file.compare(filename) && target_line == line) + is_target = true; + + } + } + + + if (auto *c = dyn_cast(&I)) { + + std::size_t found = filename.find_last_of("/\\"); + if (found != std::string::npos) + filename = filename.substr(found + 1); + + if (c->getCalledFunction()) { + std::string called = c->getCalledFunction()->getName().str(); + + bool blacklisted = false; + for (std::vector::size_type i = 0; i < blacklist.size(); i++) { + if (!called.compare(0, blacklist[i].size(), blacklist[i])) { + blacklisted = true; + break; + } + } + if (!blacklisted) + bbcalls << bb_name << "," << called << "\n"; + } + } + } + } + + if (!bb_name.empty()) { + + BB.setName(bb_name + ":"); + if (!BB.hasName()) { + std::string newname = bb_name + ":"; + Twine t(newname); + SmallString<256> NameData; + StringRef NameRef = t.toStringRef(NameData); + BB.setValueName(ValueName::Create(NameRef)); + } + + bbnames << BB.getName().str() << "\n"; + has_BBs = true; + +#ifdef AFLGO_TRACING + Value *bbnameVal = Builder.CreateGlobalStringPtr(bb_name); + Type *Args[] = { + Type::getInt8PtrTy(M.getContext()) //uint8_t* bb_name + }; + FunctionType *FTy = FunctionType::get(Type::getVoidTy(M.getContext()), Args, false); + Constant *instrumented = M.getOrInsertFunction("llvm_profiling_call", FTy); + Builder.CreateCall(instrumented, {bbnameVal}); +#endif + + } + } + + if (has_BBs) { + /* Print CFG */ + std::string cfgFileName = dotfiles + "/cfg." + funcName + ".dot"; + struct stat buffer; + if (stat (cfgFileName.c_str(), &buffer) != 0) { + FILE *cfgFILE = fopen(cfgFileName.c_str(), "w"); + if (cfgFILE) { + raw_ostream *cfgFile = + new llvm::raw_fd_ostream(fileno(cfgFILE), false, true); + + WriteGraph(*cfgFile, (const Function*)&F, true); + fflush(cfgFILE); + fclose(cfgFILE); + } + } + if (is_target) + ftargets << F.getName().str() << "\n"; + fnames << F.getName().str() << "\n"; + } + } + + bbnames.close(); + bbcalls.close(); + fnames.close(); + ftargets.close(); + + } else { + + for (auto &F : M) { + + int distance = -1; + + for (auto &BB : F) { + + distance = -1; + + if (is_aflgo) { + TerminatorInst *TI = BB.getTerminator(); + IRBuilder<> Builder(TI); + + std::string bb_name; + for (auto &I : BB) { + +#ifdef LLVM_OLD_DEBUG_API + DebugLoc Loc = I.getDebugLoc(); + if (!Loc.isUnknown()) { + + DILocation cDILoc(Loc.getAsMDNode(M.getContext())); + DILocation oDILoc = cDILoc.getOrigLocation(); + + unsigned line = oDILoc.getLineNumber(); + std::string filename = oDILoc.getFilename().str(); + + if (filename.empty()) { + line = cDILoc.getLineNumber(); + filename = cDILoc.getFilename().str(); + } +#else + if (DILocation *Loc = I.getDebugLoc()) { + + unsigned line = Loc->getLine(); + std::string filename = Loc->getFilename().str(); + + if (filename.empty()) { + DILocation *oDILoc = Loc->getInlinedAt(); + if (oDILoc) { + line = oDILoc->getLine(); + filename = oDILoc->getFilename().str(); + } + } +#endif /* LLVM_OLD_DEBUG_API */ + + if (filename.empty() || line == 0) + continue; + std::size_t found = filename.find_last_of("/\\"); + if (found != std::string::npos) + filename = filename.substr(found + 1); + + bb_name = filename + ":" + std::to_string(line); + break; + + } + + } - /* Make up cur_loc */ - unsigned int cur_loc = AFL_R(MAP_SIZE); + if (!bb_name.empty()) { - ConstantInt *CurLoc = ConstantInt::get(Int32Ty, cur_loc); + if (find(basic_blocks.begin(), basic_blocks.end(), bb_name) == basic_blocks.end()) { - /* Load prev_loc */ + if (is_selective) + continue; - LoadInst *PrevLoc = IRB.CreateLoad(AFLPrevLoc); - PrevLoc->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None)); - Value *PrevLocCasted = IRB.CreateZExt(PrevLoc, IRB.getInt32Ty()); + } else { - /* Load SHM pointer */ + /* Find distance for BB */ - LoadInst *MapPtr = IRB.CreateLoad(AFLMapPtr); - MapPtr->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None)); - Value *MapPtrIdx = - IRB.CreateGEP(MapPtr, IRB.CreateXor(PrevLocCasted, CurLoc)); + if (AFL_R(100) < dinst_ratio) { + std::map::iterator it; + for (it = bb_to_dis.begin(); it != bb_to_dis.end(); ++it) + if (it->first.compare(bb_name) == 0) + distance = it->second; - /* Update bitmap */ + ACTF("Distance for %s\t: %d", bb_name.c_str(), distance); - LoadInst *Counter = IRB.CreateLoad(MapPtrIdx); - Counter->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None)); - Value *Incr = IRB.CreateAdd(Counter, ConstantInt::get(Int8Ty, 1)); - IRB.CreateStore(Incr, MapPtrIdx) - ->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None)); + } + } + } + } - /* Set prev_loc to cur_loc >> 1 */ + BasicBlock::iterator IP = BB.getFirstInsertionPt(); + IRBuilder<> IRB(&(*IP)); - StoreInst *Store = - IRB.CreateStore(ConstantInt::get(Int32Ty, cur_loc >> 1), AFLPrevLoc); - Store->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None)); + if (AFL_R(100) >= inst_ratio) continue; - inst_blocks++; + /* Make up cur_loc */ + unsigned int cur_loc = AFL_R(MAP_SIZE); + + ConstantInt *CurLoc = ConstantInt::get(Int32Ty, cur_loc); + + /* Load prev_loc */ + + LoadInst *PrevLoc = IRB.CreateLoad(AFLPrevLoc); + PrevLoc->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None)); + Value *PrevLocCasted = IRB.CreateZExt(PrevLoc, IRB.getInt32Ty()); + + /* Load SHM pointer */ + + LoadInst *MapPtr = IRB.CreateLoad(AFLMapPtr); + MapPtr->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None)); + Value *MapPtrIdx = + IRB.CreateGEP(MapPtr, IRB.CreateXor(PrevLocCasted, CurLoc)); + + /* Update bitmap */ + + LoadInst *Counter = IRB.CreateLoad(MapPtrIdx); + Counter->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None)); + Value *Incr = IRB.CreateAdd(Counter, ConstantInt::get(Int8Ty, 1)); + IRB.CreateStore(Incr, MapPtrIdx) + ->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None)); + + /* Set prev_loc to cur_loc >> 1 */ + + StoreInst *Store = + IRB.CreateStore(ConstantInt::get(Int32Ty, cur_loc >> 1), AFLPrevLoc); + Store->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None)); + + if (distance >= 0) { + + unsigned int udistance = (unsigned) distance; + +#ifdef __x86_64__ + IntegerType *LargestType = Int64Ty; + ConstantInt *MapDistLoc = ConstantInt::get(LargestType, MAP_SIZE); + ConstantInt *MapCntLoc = ConstantInt::get(LargestType, MAP_SIZE + 8); + ConstantInt *Distance = ConstantInt::get(LargestType, udistance); +#else + IntegerType *LargestType = Int32Ty; + ConstantInt *MapDistLoc = ConstantInt::get(LargestType, MAP_SIZE); + ConstantInt *MapCntLoc = ConstantInt::get(LargestType, MAP_SIZE + 4); + ConstantInt *Distance = ConstantInt::get(LargestType, udistance); +#endif + + /* Add distance to shm[MAPSIZE] */ + + Value *MapDistPtr = IRB.CreateGEP(MapPtr, MapDistLoc); +#ifdef LLVM_OLD_DEBUG_API + LoadInst *MapDist = IRB.CreateLoad(MapDistPtr); + MapDist->mutateType(LargestType); +#else + LoadInst *MapDist = IRB.CreateLoad(LargestType, MapDistPtr); +#endif + MapDist->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None)); + Value *IncrDist = IRB.CreateAdd(MapDist, Distance); + IRB.CreateStore(IncrDist, MapDistPtr) + ->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None)); + + /* Increase count at to shm[MAPSIZE + (4 or 8)] */ + + Value *MapCntPtr = IRB.CreateGEP(MapPtr, MapCntLoc); +#ifdef LLVM_OLD_DEBUG_API + LoadInst *MapCnt = IRB.CreateLoad(MapCntPtr); + MapCnt->mutateType(LargestType); +#else + LoadInst *MapCnt = IRB.CreateLoad(LargestType, MapCntPtr); +#endif + MapCnt->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None)); + Value *IncrCnt = IRB.CreateAdd(MapCnt, ConstantInt::get(LargestType, 1)); + IRB.CreateStore(IncrCnt, MapCntPtr) + ->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None)); + + } + + inst_blocks++; + + } } + } /* Say something nice. */ - if (!be_quiet) { + if (!is_aflgo_preprocessing && !be_quiet) { if (!inst_blocks) WARNF("No instrumentation targets found."); - else OKF("Instrumented %u locations (%s mode, ratio %u%%).", - inst_blocks, getenv("AFL_HARDEN") ? "hardened" : - ((getenv("AFL_USE_ASAN") || getenv("AFL_USE_MSAN")) ? - "ASAN/MSAN" : "non-hardened"), inst_ratio); + else OKF("Instrumented %u locations (%s mode, ratio %u%%, dist. ratio %u%%).", + inst_blocks, + getenv("AFL_HARDEN") + ? "hardened" + : ((getenv("AFL_USE_ASAN") || getenv("AFL_USE_MSAN")) + ? "ASAN/MSAN" : "non-hardened"), + inst_ratio, dinst_ratio); } diff --git a/llvm_mode/afl-llvm-rt.o.c b/llvm_mode/afl-llvm-rt.o.c index ed3a664c..45e6e8f9 100644 --- a/llvm_mode/afl-llvm-rt.o.c +++ b/llvm_mode/afl-llvm-rt.o.c @@ -34,6 +34,22 @@ #include #include +#ifdef AFLGO_TRACING +#include "../hash.h" +#include "../hashset.h" +#include + +/* Variables for profiling */ +hashset_t edgeSet; +static FILE* filefd = NULL; +static char edgeStr[1024]; + +static const unsigned int prime_1 = 73; +static const unsigned int prime_2 = 5009; +/* End of profiling variables */ +#endif /* ^AFLGO_TRACING */ + + /* This is a somewhat ugly hack for the experimental 'trace-pc-guard' mode. Basically, we need to make sure that the forkserver is initialized after the LLVM-generated runtime initialization pass, not before. */ @@ -49,7 +65,7 @@ is used for instrumentation output before __afl_map_shm() has a chance to run. It will end up as .comm, so it shouldn't be too wasteful. */ -u8 __afl_area_initial[MAP_SIZE]; +u8 __afl_area_initial[MAP_SIZE + 16]; u8* __afl_area_ptr = __afl_area_initial; __thread u32 __afl_prev_loc; @@ -187,7 +203,7 @@ int __afl_persistent_loop(unsigned int max_cnt) { if (is_persistent) { - memset(__afl_area_ptr, 0, MAP_SIZE); + memset(__afl_area_ptr, 0, MAP_SIZE + 16); __afl_area_ptr[0] = 1; __afl_prev_loc = 0; } @@ -304,3 +320,161 @@ void __sanitizer_cov_trace_pc_guard_init(uint32_t* start, uint32_t* stop) { } } + + +#ifdef AFLGO_TRACING +/* Hashset implementation for C */ +hashset_t hashset_create() +{ + hashset_t set = (hashset_t) calloc(1, sizeof(struct hashset_st)); + + if (set == NULL) { + return NULL; + } + set->nbits = 3; + set->capacity = (size_t)(1 << set->nbits); + set->mask = set->capacity - 1; + set->items = (unsigned long*) calloc(set->capacity, sizeof(size_t)); + if (set->items == NULL) { + hashset_destroy(set); + return NULL; + } + set->nitems = 0; + set->n_deleted_items = 0; + return set; +} + +size_t hashset_num_items(hashset_t set) +{ + return set->nitems; +} + +void hashset_destroy(hashset_t set) +{ + if (set) { + free(set->items); + } + free(set); +} + +static int hashset_add_member(hashset_t set, void *item) +{ + size_t value = (size_t)item; + size_t ii; + + if (value == 0 || value == 1) { + return -1; + } + + ii = set->mask & (prime_1 * value); + + while (set->items[ii] != 0 && set->items[ii] != 1) { + if (set->items[ii] == value) { + return 0; + } else { + /* search free slot */ + ii = set->mask & (ii + prime_2); + } + } + set->nitems++; + if (set->items[ii] == 1) { + set->n_deleted_items--; + } + set->items[ii] = value; + return 1; +} + +static void maybe_rehash(hashset_t set) +{ + size_t *old_items; + size_t old_capacity, ii; + + + if (set->nitems + set->n_deleted_items >= (double)set->capacity * 0.85) { + old_items = set->items; + old_capacity = set->capacity; + set->nbits++; + set->capacity = (size_t)(1 << set->nbits); + set->mask = set->capacity - 1; + set->items = (unsigned long*) calloc(set->capacity, sizeof(size_t)); + set->nitems = 0; + set->n_deleted_items = 0; + assert(set->items); + for (ii = 0; ii < old_capacity; ii++) { + hashset_add_member(set, (void *)old_items[ii]); + } + free(old_items); + } +} + +int hashset_add(hashset_t set, void *item) +{ + int rv = hashset_add_member(set, item); + maybe_rehash(set); + return rv; +} + +int hashset_remove(hashset_t set, void *item) +{ + size_t value = (size_t)item; + size_t ii = set->mask & (prime_1 * value); + + while (set->items[ii] != 0) { + if (set->items[ii] == value) { + set->items[ii] = 1; + set->nitems--; + set->n_deleted_items++; + return 1; + } else { + ii = set->mask & (ii + prime_2); + } + } + return 0; +} + +int hashset_is_member(hashset_t set, void *item) +{ + size_t value = (size_t)item; + size_t ii = set->mask & (prime_1 * value); + + while (set->items[ii] != 0) { + if (set->items[ii] == value) { + return 1; + } else { + ii = set->mask & (ii + prime_2); + } + } + return 0; +} + +/*End of hashset implementation for C */ + +inline __attribute__((always_inline)) +void writeBB(const char* bbname) { + strcat(edgeStr, bbname); + size_t cksum=(size_t)hash32(bbname, strlen(edgeStr), 0xa5b35705); + if(!hashset_is_member(edgeSet,(void*)cksum)) { + fprintf(filefd, "[BB]: %s\n", bbname); + hashset_add(edgeSet, (void*)cksum); + } + strcpy(edgeStr, bbname); + fflush(filefd); +} + +void llvm_profiling_call(const char* bbname) + __attribute__((visibility("default"))); + +void llvm_profiling_call(const char* bbname) { + if (filefd != NULL) { + writeBB(bbname); + } else if (getenv("AFLGO_PROFILER_FILE")) { + filefd = fopen(getenv("AFLGO_PROFILER_FILE"), "a+"); + if (filefd != NULL) { + strcpy(edgeStr, "START"); + edgeSet = hashset_create(); + fprintf(filefd, "--------------------------\n"); + writeBB(bbname); + } + } +} +#endif /* ^AFLGO_TRACING */ diff --git a/scripts/add_edges.py b/scripts/add_edges.py new file mode 100755 index 00000000..77363b94 --- /dev/null +++ b/scripts/add_edges.py @@ -0,0 +1,61 @@ +#!/usr/bin/env python + +import argparse +import networkx as nx + +def node_name (name): + if is_cg: + return "\"{%s}\"" % name + else: + return "\"{%s:" % name + +def parse_edges (line): + edges = line.split ( ) + n1_name = node_name (edges[0]) + n1_list = filter (lambda (_, d): 'label' in d and n1_name in d['label'], G.nodes (data=True)) + if len (n1_list) > 0: + (n1, _) = n1_list[0] + for i in range (2, len(edges)): + n2_name = node_name (edges[i]) + n2_list = filter (lambda (_, d): 'label' in d and n2_name in d['label'], G.nodes (data=True)) + if len (n2_list) > 0: + (n2, _) = n2_list[0] + if G.has_edge (n1, n2): + print "[x] %s -> %s" % (n1_name, n2_name) + else: + print "[v] %s -> %s" % (n1_name, n2_name) + G.add_edge(n1,n2) + was_added = 1 +# else : +# print "Could not find %s" % n1_name + +# Main function +if __name__ == '__main__': + is_cg = 1 + was_added = 0 + parser = argparse.ArgumentParser () + parser.add_argument ('-d', '--dot', type=str, required=True, help="Path to dot-file representing the graph") + parser.add_argument ('-e', '--extra_edges', type=str, required=True, help="Extra edges to add to graph") + args = parser.parse_args () + + print "\nParsing %s .." % args.dot + G = nx.Graph(nx.drawing.nx_pydot.read_dot (args.dot)) + print nx.info (G) + + before = nx.number_connected_components (G) + + is_cg = 1 if "Name: Call graph" in nx.info (G) else 0 + print "\nWorking in %s mode.." % ("CG" if is_cg else "CFG") + + print "Adding edges.." + with open(args.extra_edges, "r") as f: + edges = map(parse_edges, f.readlines ()) + + print "\n############################################" + print "#Connected components reduced from %d to %d." % (before, nx.number_connected_components (G)) + print "############################################" + + +# print "\nWriting %s .." % args.dot + if was_added: + nx.drawing.nx_pydot.write_dot(G, args.dot) diff --git a/scripts/distance.py b/scripts/distance.py new file mode 100755 index 00000000..2224d3ca --- /dev/null +++ b/scripts/distance.py @@ -0,0 +1,149 @@ +#!/usr/bin/env python +import argparse +import networkx as nx +import re + +#Regular expression to find callee +pattern = re.compile('@.*?\(') + +def node_name (name): + if is_cg: + return "\"{%s}\"" % name + else: + return "\"{%s:" % name + +################################# +# Find the graph node for a name +################################# +def find_nodes (name): + n_name = node_name (name) + n_list = list (filter (lambda d: 'label' in d[1] and n_name in d[1]['label'], G.nodes(data=True))) + if len (n_list) > 0: + return n_list + else: + return [] + +################################## +# Calculate Distance +################################## +def distance (name): + + distance = -1 + for (n, _) in find_nodes (name): + d = 0.0 + i = 0 + if is_cg: + for (t, _) in targets: + if nx.has_path (G, n, t): + shortest = nx.dijkstra_path_length (G, n, t) + d += 1.0 / (1.0 + shortest) + i += 1 + else: + for t_name in bb_distance: + di = 0.0 + ii = 0 + for (t, _) in find_nodes(t_name): + #Check if path exists + if nx.has_path (G, n, t) : + shortest = nx.dijkstra_path_length(G, n, t) + di += 1.0 / (1.0 + 10 * bb_distance[t_name] + shortest) + ii += 1 + if ii != 0: + d += di / ii + i += 1 + + if d != 0 and (distance == -1 or distance > i / d) : + distance = i / d + + if distance != -1: + out.write (name) + out.write (",") + out.write (str (distance)) + out.write ("\n") + +# Main function +if __name__ == '__main__': + parser = argparse.ArgumentParser () + parser.add_argument ('-d', '--dot', type=str, required=True, help="Path to dot-file representing the graph.") + parser.add_argument ('-t', '--targets', type=str, required=True, help="Path to file specifying Target nodes.") + parser.add_argument ('-o', '--out', type=str, required=True, help="Path to output file containing distance for each node.") + parser.add_argument ('-n', '--names', type=str, required=True, help="Path to file containing name for each node.") + parser.add_argument ('-c', '--cg_distance', type=str, help="Path to file containing call graph distance.") + parser.add_argument ('-s', '--cg_callsites', type=str, help="Path to file containing mapping between basic blocks and called functions.") + + args = parser.parse_args () + + print ("\nParsing %s .." % args.dot) + G = nx.DiGraph(nx.drawing.nx_pydot.read_dot(args.dot)) + print (nx.info(G)) + + is_cg = 1 if "Name: Call graph" in nx.info(G) else 0 + print ("\nWorking in %s mode.." % ("CG" if is_cg else "CFG")) + + # Process as ControlFlowGraph + caller = "" + cg_distance = {} + bb_distance = {} + if not is_cg : + + if args.cg_distance is None: + print ("Specify file containing CG-level distance (-c).") + exit(1) + + elif args.cg_callsites is None: + print ("Specify file containing mapping between basic blocks and called functions (-s).") + exit(1) + + else: + + caller = args.dot.split(".") + caller = caller[len(caller)-2] + print ("Loading cg_distance for function '%s'.." % caller) + + with open(args.cg_distance, 'r') as f: + for l in f.readlines(): + s = l.strip().split(",") + cg_distance[s[0]] = float(s[1]) + + with open(args.cg_callsites, 'r') as f: + for l in f.readlines(): + s = l.strip().split(",") + if len(find_nodes(s[0])) > 0: + if s[1] in cg_distance: + if s[0] in bb_distance: + if bb_distance[s[0]] > cg_distance[s[1]]: + bb_distance[s[0]] = cg_distance[s[1]] + else: + bb_distance[s[0]] = cg_distance[s[1]] + + print ("Adding target BBs (if any)..") + with open(args.targets, "r") as f: + for l in f.readlines (): + s = l.strip().split("/"); + line = s[len(s) - 1] + nodes = find_nodes(line) + if len(nodes) > 0: + bb_distance[line] = 0 + print ("Added target BB!") + + # Process as CallGraph + else: + + print ("Loading targets..") + with open(args.targets, "r") as f: + targets = [] + for line in f.readlines (): + line = line.strip () + for target in find_nodes(line): + targets.append (target) + + if (len (targets) == 0 and is_cg): + print ("No targets available") + exit(1) + + print ("Calculating distance..") + with open(args.out, "w") as out: + with open(args.names, "r") as f: + for line in f.readlines(): + line = line.strip() + distance (line) diff --git a/scripts/genDistance.sh b/scripts/genDistance.sh new file mode 100755 index 00000000..1a42f4e9 --- /dev/null +++ b/scripts/genDistance.sh @@ -0,0 +1,150 @@ +#!/bin/bash +if [ $# -lt 2 ]; then + echo "Usage: $0 [fuzzer-name]" + echo "" + exit 1 +fi + +BINARIES=$(readlink -e $1) +TMPDIR=$(readlink -e $2) +AFLGO="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )" +fuzzer="" +if [ $# -eq 3 ]; then + fuzzer=$(find $BINARIES -name "$3.0.0.*.bc" | rev | cut -d. -f5- | rev) + if [ $(echo "$fuzzer" | wc -l) -ne 1 ]; then + echo "Couldn't find bytecode for fuzzer $3 in folder $BINARIES." + exit 1 + fi +fi + +SCRIPT=$0 +ARGS=$@ + +#SANITY CHECKS +if [ -z "$BINARIES" ]; then echo "Couldn't find binaries folder ($1)."; exit 1; fi +if ! [ -d "$BINARIES" ]; then echo "No directory: $BINARIES."; exit 1; fi +if [ -z "$TMPDIR" ]; then echo "Couldn't find temporary directory ($3)."; exit 1; fi + +binaries=$(find $BINARIES -name "*.0.0.*.bc" | rev | cut -d. -f5- | rev) +if [ -z "$binaries" ]; then echo "Couldn't find any binaries in folder $BINARIES."; exit; fi + +if [ -z $(which python) ] && [ -z $(which python3) ]; then echo "Please install Python"; exit 1; fi +#if python -c "import pydotplus"; then echo "Install python package: pydotplus (sudo pip install pydotplus)"; exit 1; fi +#if python -c "import pydotplus; import networkx"; then echo "Install python package: networkx (sudo pip install networkx)"; exit 1; fi + +FAIL=0 +STEP=1 + +RESUME=$(if [ -f $TMPDIR/state ]; then cat $TMPDIR/state; else echo 0; fi) + +function next_step { + echo $STEP > $TMPDIR/state + if [ $FAIL -ne 0 ]; then + tail -n30 $TMPDIR/step${STEP}.log + echo "-- Problem in Step $STEP of generating $OUT!" + echo "-- You can resume by executing:" + echo "$ $SCRIPT $ARGS $TMPDIR" + exit 1 + fi + STEP=$((STEP + 1)) +} + + +#------------------------------------------------------------------------------- +# Construct control flow graph and call graph +#------------------------------------------------------------------------------- +if [ $RESUME -le $STEP ]; then + + cd $TMPDIR/dot-files + + if [ -z "$fuzzer" ]; then + for binary in $(echo "$binaries"); do + + echo "($STEP) Constructing CG for $binary.." + while ! opt -dot-callgraph $binary.0.0.*.bc >/dev/null 2> $TMPDIR/step${STEP}.log ; do + echo -e "\e[93;1m[!]\e[0m Could not generate call graph. Repeating.." + done + + #Remove repeated lines and rename + awk '!a[$0]++' callgraph.dot > callgraph.$(basename $binary).dot + rm callgraph.dot + done + + #TODO Integrate several call graphs into one + callgraph=$(ls -1d callgraph.* | head -n1) + cp $callgraph callgraph.dot + echo "($STEP) TODO: Integrate several call graphs into one. Now using $callgraph." + + else + + echo "($STEP) Constructing CG for $fuzzer.." + while ! opt -dot-callgraph $fuzzer.0.0.*.bc >/dev/null 2> $TMPDIR/step${STEP}.log ; do + echo -e "\e[93;1m[!]\e[0m Could not generate call graph. Repeating.." + done + + #Remove repeated lines and rename + awk '!a[$0]++' callgraph.dot > callgraph.1.dot + mv callgraph.1.dot callgraph.dot + + fi +fi +next_step + +#------------------------------------------------------------------------------- +# Generate config file keeping distance information for code instrumentation +#------------------------------------------------------------------------------- +if [ $RESUME -le $STEP ]; then + echo "($STEP) Computing distance for call graph .." + + $AFLGO/distance.py -d $TMPDIR/dot-files/callgraph.dot -t $TMPDIR/Ftargets.txt -n $TMPDIR/Fnames.txt -o $TMPDIR/distance.callgraph.txt > $TMPDIR/step${STEP}.log 2>&1 || FAIL=1 + + if [ $(cat $TMPDIR/distance.callgraph.txt | wc -l) -eq 0 ]; then + FAIL=1 + next_step + fi + + printf "($STEP) Computing distance for control-flow graphs " + for f in $(ls -1d $TMPDIR/dot-files/cfg.*.dot); do + + # Skip CFGs of functions we are not calling + if ! grep "$(basename $f | cut -d. -f2)" $TMPDIR/dot-files/callgraph.dot >/dev/null; then + printf "\nSkipping $f..\n" + continue + fi + + printf "." + + #Clean up duplicate lines and \" in labels (bug in Pydotplus) + awk '!a[$0]++' $f > ${f}.smaller.dot + mv $f $f.bigger.dot + mv $f.smaller.dot $f + sed -i s/\\\\\"//g $f + sed -i 's/\[.\"]//g' $f + sed -i 's/\(^\s*[0-9a-zA-Z_]*\):[a-zA-Z0-9]*\( -> \)/\1\2/g' $f + + #Compute distance + $AFLGO/distance.py -d $f -t $TMPDIR/BBtargets.txt -n $TMPDIR/BBnames.txt -s $TMPDIR/BBcalls.txt -c $TMPDIR/distance.callgraph.txt -o ${f}.distances.txt >> $TMPDIR/step${STEP}.log 2>&1 #|| FAIL=1 + if [ $? -ne 0 ]; then + echo -e "\e[93;1m[!]\e[0m Could not calculate distance for $f." + fi + #if [ $FAIL -eq 1 ]; then + # next_step #Fail asap. + #fi + done + echo "" + + cat $TMPDIR/dot-files/*.distances.txt > $TMPDIR/distance.cfg.txt + +fi +next_step + +echo "" +echo "----------[DONE]----------" +echo "" +echo "Now, you may wish to compile your sources with " +echo "CC=\"$AFLGO/../afl-clang-fast\"" +echo "CXX=\"$AFLGO/../afl-clang-fast++\"" +echo "CFLAGS=\"\$CFLAGS -distance=$(readlink -e $TMPDIR/distance.cfg.txt)\"" +echo "CXXFLAGS=\"\$CXXFLAGS -distance=$(readlink -e $TMPDIR/distance.cfg.txt)\"" +echo "" +echo "--------------------------"