From 99d24d13f1ad972ab0645cf81155b47a7e693a42 Mon Sep 17 00:00:00 2001 From: Adrian Herrera Date: Thu, 22 Jul 2021 11:19:57 +0000 Subject: optimin: add missing cmin options and auto create output dir * Includes frida, unicorn modes * Automatically create the output directory if it does not already exist --- utils/optimin/src/OptiMin.cpp | 32 +++++++++++++++++++++----------- 1 file changed, 21 insertions(+), 11 deletions(-) (limited to 'utils/optimin/src/OptiMin.cpp') diff --git a/utils/optimin/src/OptiMin.cpp b/utils/optimin/src/OptiMin.cpp index b82acbb6..55e391df 100644 --- a/utils/optimin/src/OptiMin.cpp +++ b/utils/optimin/src/OptiMin.cpp @@ -116,29 +116,38 @@ static const auto StatMsg = [] { }; -static cl::opt CorpusDir("i", cl::desc("Input directory"), +static cl::opt CorpusDir("i", cl::desc("Input directory"), cl::value_desc("dir"), cl::Required); -static cl::opt OutputDir("o", cl::desc("Output directory"), +static cl::opt OutputDir("o", cl::desc("Output directory"), cl::value_desc("dir"), cl::Required); -static cl::opt EdgesOnly("f", cl::desc("Include edge hit counts"), + +static cl::opt ShowProgBar("p", cl::desc("Display progress bar")); +static cl::opt EdgesOnly("f", cl::desc("Include edge hit counts"), cl::init(true)); -static cl::opt ShowProgBar("p", cl::desc("Display progress bar")); -static cl::opt WeightsFile("w", cl::desc("Weights file"), +static cl::opt WeightsFile("w", cl::desc("Weights file"), cl::value_desc("csv")); + static cl::opt TargetProg(cl::Positional, cl::desc(""), cl::Required); static cl::list TargetArgs(cl::ConsumeAfter, cl::desc("[target args...]")); -static cl::opt MemLimit( + +static cl::opt MemLimit( "m", cl::desc("Memory limit for child process (default=none)"), cl::value_desc("megs"), cl::init("none")); static cl::opt Timeout( "t", cl::desc("Run time limit for child process (default=5000)"), cl::value_desc("msec"), cl::init("4000")); + static cl::opt CrashMode( "C", cl::desc("Keep crashing inputs, reject everything else")); -static cl::opt QemuMode("Q", cl::desc("Use binary-only instrumentation")); +static cl::opt FridaMode( + "O", cl::desc("Use binary-only instrumentation (FRIDA mode)")); +static cl::opt QemuMode( + "Q", cl::desc("Use binary-only instrumentation (QEMU mode)")); +static cl::opt UnicornMode( + "U", cl::desc("Use unicorn-based instrumentation (unicorn mode)")); } // anonymous namespace @@ -184,7 +193,6 @@ static void GetWeights(const MemoryBuffer &MB, WeightsMap &Weights) { // Prepare afl-showmap arguments SmallVector AFLShowmapArgs{ - AFLShowmapPath, "-m", MemLimit, "-t", Timeout, "-q", "-o", OutputPath}; if (TargetArgsHasAtAt) @@ -192,8 +200,9 @@ static void GetWeights(const MemoryBuffer &MB, WeightsMap &Weights) { else Redirects[/* stdin */ 0] = Seed; + if (FridaMode) AFLShowmapArgs.push_back("-O"); if (QemuMode) AFLShowmapArgs.push_back("-Q"); - if (CrashMode) AFLShowmapArgs.push_back("-C"); + if (UnicornMode) AFLShowmapArgs.push_back("-U"); AFLShowmapArgs.append({"--", TargetProg}); AFLShowmapArgs.append(TargetArgs.begin(), TargetArgs.end()); @@ -269,9 +278,10 @@ int main(int argc, char *argv[]) { cl::ParseCommandLineOptions(argc, argv, "Optimal corpus minimizer"); - if (!sys::fs::is_directory(OutputDir)) { + if (EC = sys::fs::create_directory(OutputDir)) { - ErrMsg() << "Invalid output directory `" << OutputDir << "`\n"; + ErrMsg() << "Invalid output directory `" << OutputDir + << "`: " << EC.message() << '\n'; return 1; } -- cgit 1.4.1 From 7123dd9d299206e7e21cbc8152abb0666acb170f Mon Sep 17 00:00:00 2001 From: Adrian Herrera Date: Fri, 23 Jul 2021 08:23:29 +0000 Subject: optimin: improve afl-showmap error messages --- utils/optimin/src/OptiMin.cpp | 40 +++++++++++++++++++++++----------------- 1 file changed, 23 insertions(+), 17 deletions(-) (limited to 'utils/optimin/src/OptiMin.cpp') diff --git a/utils/optimin/src/OptiMin.cpp b/utils/optimin/src/OptiMin.cpp index 55e391df..dca433cc 100644 --- a/utils/optimin/src/OptiMin.cpp +++ b/utils/optimin/src/OptiMin.cpp @@ -14,6 +14,7 @@ #include #include #include +#include #include #include #include @@ -180,16 +181,18 @@ static void GetWeights(const MemoryBuffer &MB, WeightsMap &Weights) { } -[[nodiscard]] static std::error_code getAFLCoverage(const StringRef Seed, - AFLCoverageVector &Cov) { +static Error getAFLCoverage(const StringRef Seed, AFLCoverageVector &Cov) { Optional Redirects[] = {None, None, None}; - std::error_code EC; // Create temporary output file - SmallString<64> OutputPath; - if (EC = sys::fs::createTemporaryFile("showmap", "txt", OutputPath)) - return EC; + SmallString<32> OutputPath; + if (const auto EC = sys::fs::createTemporaryFile("showmap", "txt", OutputPath)) { + + return createStringError( + EC, "Failed to create temporary file for afl-showmap output"); + + } // Prepare afl-showmap arguments SmallVector AFLShowmapArgs{ @@ -210,14 +213,16 @@ static void GetWeights(const MemoryBuffer &MB, WeightsMap &Weights) { // Run afl-showmap const int RC = sys::ExecuteAndWait(AFLShowmapPath, AFLShowmapArgs, /*env=*/None, Redirects); - if (RC) return std::make_error_code(std::errc::executable_format_error); + if (RC) + return createStringError(inconvertibleErrorCode(), "afl-showmap failed"); // Parse afl-showmap output const auto CovOrErr = MemoryBuffer::getFile(OutputPath); - if (EC = CovOrErr.getError()) { + if (const auto EC = CovOrErr.getError()) { sys::fs::remove(OutputPath); - return EC; + return createStringError(EC, "Failed to read afl-showmap output file `%s`", + OutputPath.c_str()); } @@ -237,7 +242,8 @@ static void GetWeights(const MemoryBuffer &MB, WeightsMap &Weights) { } - return sys::fs::remove(OutputPath); + sys::fs::remove(OutputPath); + return Error::success(); } @@ -278,7 +284,7 @@ int main(int argc, char *argv[]) { cl::ParseCommandLineOptions(argc, argv, "Optimal corpus minimizer"); - if (EC = sys::fs::create_directory(OutputDir)) { + if ((EC = sys::fs::create_directory(OutputDir))) { ErrMsg() << "Invalid output directory `" << OutputDir << "`: " << EC.message() << '\n'; @@ -316,7 +322,7 @@ int main(int argc, char *argv[]) { StartTimer(/*ShowProgBar=*/false); const auto WeightsOrErr = MemoryBuffer::getFile(WeightsFile); - if (EC = WeightsOrErr.getError()) { + if ((EC = WeightsOrErr.getError())) { ErrMsg() << "Failed to read weights from `" << WeightsFile << "`: " << EC.message() << '\n'; @@ -354,7 +360,7 @@ int main(int argc, char *argv[]) { } const auto &Path = Dir->path(); - if (EC = sys::fs::status(Path, Status)) { + if ((EC = sys::fs::status(Path, Status))) { WarnMsg() << "Failed to access seed file `" << Path << "`: " << EC.message() << ". Skipping...\n"; @@ -402,10 +408,10 @@ int main(int argc, char *argv[]) { // Execute seed Cov.clear(); - if (EC = getAFLCoverage(SeedFile, Cov)) { + if (auto Err = getAFLCoverage(SeedFile, Cov)) { - ErrMsg() << "Failed to get coverage for seed " << SeedFile << ": " - << EC.message() << '\n'; + ErrMsg() << "Failed to get coverage for seed `" << SeedFile << "`: " + << Err << '\n'; return 1; } @@ -519,7 +525,7 @@ int main(int argc, char *argv[]) { OutputSeed = OutputDir; sys::path::append(OutputSeed, sys::path::filename(Seed)); - if (EC = sys::fs::copy_file(Seed, OutputSeed)) { + if ((EC = sys::fs::copy_file(Seed, OutputSeed))) { WarnMsg() << "Failed to copy `" << Seed << "` to `" << OutputDir << "`: " << EC.message() << '\n'; -- cgit 1.4.1 From 64daf1058407985697874df9a4860bddd4fa3f1b Mon Sep 17 00:00:00 2001 From: Adrian Herrera Date: Fri, 23 Jul 2021 08:25:49 +0000 Subject: optimin: recursively iterate over input directory --- utils/optimin/src/OptiMin.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) (limited to 'utils/optimin/src/OptiMin.cpp') diff --git a/utils/optimin/src/OptiMin.cpp b/utils/optimin/src/OptiMin.cpp index dca433cc..b6997939 100644 --- a/utils/optimin/src/OptiMin.cpp +++ b/utils/optimin/src/OptiMin.cpp @@ -348,7 +348,7 @@ int main(int argc, char *argv[]) { std::vector SeedFiles; sys::fs::file_status Status; - for (sys::fs::directory_iterator Dir(CorpusDir, EC), DirEnd; + for (sys::fs::recursive_directory_iterator Dir(CorpusDir, EC), DirEnd; Dir != DirEnd && !EC; Dir.increment(EC)) { if (EC) { -- cgit 1.4.1 From aebb14ceeb5af441c49919802386bcfc18e3574d Mon Sep 17 00:00:00 2001 From: Adrian Herrera Date: Sat, 24 Jul 2021 01:14:53 +0000 Subject: optimin: add support for afl-cmin env vars --- utils/optimin/src/OptiMin.cpp | 64 ++++++++++++++++++++++++++++--------------- 1 file changed, 42 insertions(+), 22 deletions(-) (limited to 'utils/optimin/src/OptiMin.cpp') diff --git a/utils/optimin/src/OptiMin.cpp b/utils/optimin/src/OptiMin.cpp index b6997939..8d839911 100644 --- a/utils/optimin/src/OptiMin.cpp +++ b/utils/optimin/src/OptiMin.cpp @@ -5,6 +5,7 @@ */ #include +#include #include #include @@ -92,6 +93,7 @@ static std::chrono::seconds Duration; static std::string AFLShowmapPath; static bool TargetArgsHasAtAt = false; +static bool KeepTraces = false; static const auto ErrMsg = [] { @@ -185,18 +187,12 @@ static Error getAFLCoverage(const StringRef Seed, AFLCoverageVector &Cov) { Optional Redirects[] = {None, None, None}; - // Create temporary output file - SmallString<32> OutputPath; - if (const auto EC = sys::fs::createTemporaryFile("showmap", "txt", OutputPath)) { - - return createStringError( - EC, "Failed to create temporary file for afl-showmap output"); - - } + SmallString<32> TracePath{OutputDir}; + sys::path::append(TracePath, ".traces", sys::path::filename(Seed)); // Prepare afl-showmap arguments SmallVector AFLShowmapArgs{ - AFLShowmapPath, "-m", MemLimit, "-t", Timeout, "-q", "-o", OutputPath}; + AFLShowmapPath, "-m", MemLimit, "-t", Timeout, "-q", "-o", TracePath}; if (TargetArgsHasAtAt) AFLShowmapArgs.append({"-A", Seed}); @@ -217,12 +213,12 @@ static Error getAFLCoverage(const StringRef Seed, AFLCoverageVector &Cov) { return createStringError(inconvertibleErrorCode(), "afl-showmap failed"); // Parse afl-showmap output - const auto CovOrErr = MemoryBuffer::getFile(OutputPath); + const auto CovOrErr = MemoryBuffer::getFile(TracePath); if (const auto EC = CovOrErr.getError()) { - sys::fs::remove(OutputPath); + sys::fs::remove(TracePath); return createStringError(EC, "Failed to read afl-showmap output file `%s`", - OutputPath.c_str()); + TracePath.c_str()); } @@ -242,7 +238,7 @@ static Error getAFLCoverage(const StringRef Seed, AFLCoverageVector &Cov) { } - sys::fs::remove(OutputPath); + if (!KeepTraces) sys::fs::remove(TracePath); return Error::success(); } @@ -277,20 +273,15 @@ int main(int argc, char *argv[]) { std::error_code EC; // ------------------------------------------------------------------------ // - // Parse command-line options + // Parse command-line options and environment variables // // Also check the target arguments, as this determines how we run afl-showmap. // ------------------------------------------------------------------------ // cl::ParseCommandLineOptions(argc, argv, "Optimal corpus minimizer"); - if ((EC = sys::fs::create_directory(OutputDir))) { - - ErrMsg() << "Invalid output directory `" << OutputDir - << "`: " << EC.message() << '\n'; - return 1; - - } + KeepTraces = !!std::getenv("AFL_KEEP_TRACES"); + const auto AFLPath = std::getenv("AFL_PATH"); for (const auto &Arg : TargetArgs) if (Arg == "@@") TargetArgsHasAtAt = true; @@ -299,7 +290,13 @@ int main(int argc, char *argv[]) { // Find afl-showmap // ------------------------------------------------------------------------ // - const auto AFLShowmapOrErr = sys::findProgramByName("afl-showmap"); + SmallVector EnvPaths; + + if (const char *PathEnv = std::getenv("PATH")) + SplitString(PathEnv, EnvPaths, ":"); + if (AFLPath) EnvPaths.push_back(AFLPath); + + const auto AFLShowmapOrErr = sys::findProgramByName("afl-showmap", EnvPaths); if (AFLShowmapOrErr.getError()) { ErrMsg() << "Failed to find afl-showmap. Check your PATH\n"; @@ -336,6 +333,29 @@ int main(int argc, char *argv[]) { } + // ------------------------------------------------------------------------ // + // Setup output directory + // ------------------------------------------------------------------------ // + + SmallString<32> TraceDir{OutputDir}; + sys::path::append(TraceDir, ".traces"); + + if ((EC = sys::fs::remove_directories(TraceDir))) { + + ErrMsg() << "Failed to remove existing trace directory in `" << OutputDir + << "`: " << EC.message() << '\n'; + return 1; + + } + + if ((EC = sys::fs::create_directories(TraceDir))) { + + ErrMsg() << "Failed to create output directory `" << OutputDir << "`: " + << EC.message() << '\n'; + return 1; + + } + // ------------------------------------------------------------------------ // // Traverse corpus directory // -- cgit 1.4.1 From 1afa0d78619e57fcdfbc19d9cfa37256962dea59 Mon Sep 17 00:00:00 2001 From: Adrian Herrera Date: Sat, 24 Jul 2021 23:31:06 +0000 Subject: optimin: add test of binary target --- utils/optimin/src/OptiMin.cpp | 34 ++++++++++++++++++++++++++++++---- 1 file changed, 30 insertions(+), 4 deletions(-) (limited to 'utils/optimin/src/OptiMin.cpp') diff --git a/utils/optimin/src/OptiMin.cpp b/utils/optimin/src/OptiMin.cpp index 8d839911..84179022 100644 --- a/utils/optimin/src/OptiMin.cpp +++ b/utils/optimin/src/OptiMin.cpp @@ -94,6 +94,7 @@ static std::chrono::seconds Duration; static std::string AFLShowmapPath; static bool TargetArgsHasAtAt = false; static bool KeepTraces = false; +static bool SkipBinCheck = false; static const auto ErrMsg = [] { @@ -183,12 +184,14 @@ static void GetWeights(const MemoryBuffer &MB, WeightsMap &Weights) { } -static Error getAFLCoverage(const StringRef Seed, AFLCoverageVector &Cov) { +static Error getAFLCoverage(const StringRef Seed, AFLCoverageVector &Cov, + bool BinCheck = false) { Optional Redirects[] = {None, None, None}; SmallString<32> TracePath{OutputDir}; - sys::path::append(TracePath, ".traces", sys::path::filename(Seed)); + StringRef TraceName = BinCheck ? ".run_test" : sys::path::filename(Seed); + sys::path::append(TracePath, ".traces", TraceName); // Prepare afl-showmap arguments SmallVector AFLShowmapArgs{ @@ -238,7 +241,7 @@ static Error getAFLCoverage(const StringRef Seed, AFLCoverageVector &Cov) { } - if (!KeepTraces) sys::fs::remove(TracePath); + if (!KeepTraces || BinCheck) sys::fs::remove(TracePath); return Error::success(); } @@ -281,6 +284,7 @@ int main(int argc, char *argv[]) { cl::ParseCommandLineOptions(argc, argv, "Optimal corpus minimizer"); KeepTraces = !!std::getenv("AFL_KEEP_TRACES"); + SkipBinCheck = !!std::getenv("AFL_SKIP_BIN_CHECK"); const auto AFLPath = std::getenv("AFL_PATH"); for (const auto &Arg : TargetArgs) @@ -404,6 +408,29 @@ int main(int argc, char *argv[]) { EndTimer(/*ShowProgBar=*/false); + // ------------------------------------------------------------------------ // + // Test the target binary + // ------------------------------------------------------------------------ // + + AFLCoverageVector Cov; + + if (!SkipBinCheck && SeedFiles.size() > 0) { + + StatMsg() << "Testing the target binary... "; + StartTimer(/*ShowProgBar=*/false); + + if (auto Err = getAFLCoverage(SeedFiles.front(), Cov, /*BinCheck=*/true)) { + + ErrMsg() << "No instrumentation output detected (perhaps crash or timeout)"; + return 1; + + } + + EndTimer(/*ShowProgBar=*/false); + SuccMsg() << "OK, " << Cov.size() << " tuples recorded\n"; + + } + // ------------------------------------------------------------------------ // // Generate seed coverage // @@ -422,7 +449,6 @@ int main(int argc, char *argv[]) { EvalMaxSAT Solver(/*nbMinimizeThread=*/0); MaxSATSeeds SeedVars; MaxSATCoverageMap SeedCoverage; - AFLCoverageVector Cov; for (const auto &SeedFile : SeedFiles) { -- cgit 1.4.1 From 33f9f911c13cfd513366b7b1c9f9c66c0aa944f7 Mon Sep 17 00:00:00 2001 From: Adrian Herrera Date: Sat, 24 Jul 2021 23:37:28 +0000 Subject: optimin: fix default timeout So that it matches the help text --- utils/optimin/README.md | 2 +- utils/optimin/src/OptiMin.cpp | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) (limited to 'utils/optimin/src/OptiMin.cpp') diff --git a/utils/optimin/README.md b/utils/optimin/README.md index 3036b449..c6f2af06 100644 --- a/utils/optimin/README.md +++ b/utils/optimin/README.md @@ -40,7 +40,7 @@ General options: -m megs - Memory limit for child process (default=none) -o dir - Output directory -p - Display progress bar - -t msec - Run time limit for child process (default=none) + -t msec - Run time limit for child process (default=5000) -w csv - Weights file Generic Options: diff --git a/utils/optimin/src/OptiMin.cpp b/utils/optimin/src/OptiMin.cpp index 84179022..4f88b897 100644 --- a/utils/optimin/src/OptiMin.cpp +++ b/utils/optimin/src/OptiMin.cpp @@ -142,7 +142,7 @@ static cl::opt MemLimit( cl::value_desc("megs"), cl::init("none")); static cl::opt Timeout( "t", cl::desc("Run time limit for child process (default=5000)"), - cl::value_desc("msec"), cl::init("4000")); + cl::value_desc("msec"), cl::init("5000")); static cl::opt CrashMode( "C", cl::desc("Keep crashing inputs, reject everything else")); -- cgit 1.4.1 From cc933bcc0c63df22df981ef991582170d40fbb95 Mon Sep 17 00:00:00 2001 From: Adrian Herrera Date: Mon, 26 Jul 2021 02:11:21 +0000 Subject: optimin: support afl-cmin crash mode --- utils/optimin/src/OptiMin.cpp | 20 ++++++++++++++------ 1 file changed, 14 insertions(+), 6 deletions(-) (limited to 'utils/optimin/src/OptiMin.cpp') diff --git a/utils/optimin/src/OptiMin.cpp b/utils/optimin/src/OptiMin.cpp index 4f88b897..276f6b0f 100644 --- a/utils/optimin/src/OptiMin.cpp +++ b/utils/optimin/src/OptiMin.cpp @@ -212,9 +212,14 @@ static Error getAFLCoverage(const StringRef Seed, AFLCoverageVector &Cov, // Run afl-showmap const int RC = sys::ExecuteAndWait(AFLShowmapPath, AFLShowmapArgs, /*env=*/None, Redirects); - if (RC) + if (RC && !CrashMode) { + + ErrMsg() << "Exit code " << RC << " != 0 received from afl-showmap"; + sys::fs::remove(TracePath); return createStringError(inconvertibleErrorCode(), "afl-showmap failed"); + } + // Parse afl-showmap output const auto CovOrErr = MemoryBuffer::getFile(TracePath); if (const auto EC = CovOrErr.getError()) { @@ -287,6 +292,8 @@ int main(int argc, char *argv[]) { SkipBinCheck = !!std::getenv("AFL_SKIP_BIN_CHECK"); const auto AFLPath = std::getenv("AFL_PATH"); + if (CrashMode) ::setenv("AFL_CMIN_CRASHES_ONLY", "1", /*overwrite=*/true); + for (const auto &Arg : TargetArgs) if (Arg == "@@") TargetArgsHasAtAt = true; @@ -354,8 +361,8 @@ int main(int argc, char *argv[]) { if ((EC = sys::fs::create_directories(TraceDir))) { - ErrMsg() << "Failed to create output directory `" << OutputDir << "`: " - << EC.message() << '\n'; + ErrMsg() << "Failed to create output directory `" << OutputDir + << "`: " << EC.message() << '\n'; return 1; } @@ -421,7 +428,8 @@ int main(int argc, char *argv[]) { if (auto Err = getAFLCoverage(SeedFiles.front(), Cov, /*BinCheck=*/true)) { - ErrMsg() << "No instrumentation output detected (perhaps crash or timeout)"; + ErrMsg() + << "No instrumentation output detected (perhaps crash or timeout)"; return 1; } @@ -456,8 +464,8 @@ int main(int argc, char *argv[]) { Cov.clear(); if (auto Err = getAFLCoverage(SeedFile, Cov)) { - ErrMsg() << "Failed to get coverage for seed `" << SeedFile << "`: " - << Err << '\n'; + ErrMsg() << "Failed to get coverage for seed `" << SeedFile + << "`: " << Err << '\n'; return 1; } -- cgit 1.4.1 From 036282185b0f06df7b7ba354f41159d682853008 Mon Sep 17 00:00:00 2001 From: Adrian Herrera Date: Wed, 28 Jul 2021 04:40:43 +0000 Subject: optimin: run showmap with an entire input directory Rather than executing with one seed at a time. This required disabling the target binary test (for now). --- utils/optimin/src/OptiMin.cpp | 295 +++++++++++++++++++++++++----------------- 1 file changed, 179 insertions(+), 116 deletions(-) (limited to 'utils/optimin/src/OptiMin.cpp') diff --git a/utils/optimin/src/OptiMin.cpp b/utils/optimin/src/OptiMin.cpp index 276f6b0f..242eeea9 100644 --- a/utils/optimin/src/OptiMin.cpp +++ b/utils/optimin/src/OptiMin.cpp @@ -64,7 +64,10 @@ using AFLTuple = std::pair; /// Coverage for a given seed file using AFLCoverageVector = std::vector; -/// Maps seed file paths to a weight +/// Map seed file paths to its coverage vector +using AFLCoverageMap = StringMap; + +/// Map seed file paths to a weight using WeightsMap = StringMap; /// A seed identifier in the MaxSAT solver @@ -88,6 +91,9 @@ using MaxSATCoverageMap = DenseMap; // `afl-showmap.c` static constexpr uint32_t MAX_EDGE_FREQ = 8; +// The maximum number of failures allowed when parsing a weights file +static constexpr unsigned MAX_WEIGHT_FAILURES = 5; + static sys::TimePoint<> StartTime, EndTime; static std::chrono::seconds Duration; @@ -98,30 +104,30 @@ static bool SkipBinCheck = false; static const auto ErrMsg = [] { - return WithColor(errs(), HighlightColor::Error) << "[-] "; + return WithColor(errs(), raw_ostream::RED, /*Bold=*/true) << "[-] "; }; static const auto WarnMsg = [] { - return WithColor(errs(), HighlightColor::Warning) << "[-] "; + return WithColor(errs(), raw_ostream::MAGENTA, /*Bold=*/true) << "[-] "; }; static const auto SuccMsg = [] { - return WithColor(outs(), HighlightColor::String) << "[+] "; + return WithColor(outs(), raw_ostream::GREEN, /*Bold=*/true) << "[+] "; }; static const auto StatMsg = [] { - return WithColor(outs(), HighlightColor::Remark) << "[*] "; + return WithColor(outs(), raw_ostream::BLUE, /*Bold=*/true) << "[*] "; }; -static cl::opt CorpusDir("i", cl::desc("Input directory"), - cl::value_desc("dir"), cl::Required); +static cl::opt InputDir("i", cl::desc("Input directory"), + cl::value_desc("dir"), cl::Required); static cl::opt OutputDir("o", cl::desc("Output directory"), cl::value_desc("dir"), cl::Required); @@ -164,6 +170,7 @@ static void GetWeights(const MemoryBuffer &MB, WeightsMap &Weights) { SmallVector Lines; MB.getBuffer().trim().split(Lines, '\n'); + unsigned FailureCount = 0; unsigned Weight = 0; for (const auto &Line : Lines) { @@ -176,7 +183,13 @@ static void GetWeights(const MemoryBuffer &MB, WeightsMap &Weights) { } else { - WarnMsg() << "Failed to read weight for `" << Seed << "`. Skipping...\n"; + if (FailureCount >= MAX_WEIGHT_FAILURES) { + ErrMsg() << "Too many failures. Aborting\n"; + std::exit(1); + } + + WarnMsg() << "Failed to read weight for '" << Seed << "'. Skipping...\n"; + FailureCount++; } @@ -184,23 +197,52 @@ static void GetWeights(const MemoryBuffer &MB, WeightsMap &Weights) { } -static Error getAFLCoverage(const StringRef Seed, AFLCoverageVector &Cov, - bool BinCheck = false) { +static std::error_code readCov(const StringRef Trace, AFLCoverageVector &Cov) { + + const auto CovOrErr = MemoryBuffer::getFile(Trace); + if (const auto EC = CovOrErr.getError()) return EC; + + SmallVector Lines; + CovOrErr.get()->getBuffer().trim().split(Lines, '\n'); + + AFLTupleID Edge = 0; + unsigned Freq = 0; + + for (const auto &Line : Lines) { + + const auto &[EdgeStr, FreqStr] = Line.split(':'); + + to_integer(EdgeStr, Edge, 10); + to_integer(FreqStr, Freq, 10); + Cov.push_back({Edge, Freq}); + + } + + return std::error_code(); + +} + +static Error runShowmap(AFLCoverageMap &CovMap) { Optional Redirects[] = {None, None, None}; - SmallString<32> TracePath{OutputDir}; - StringRef TraceName = BinCheck ? ".run_test" : sys::path::filename(Seed); - sys::path::append(TracePath, ".traces", TraceName); + SmallString<32> TraceDir{OutputDir}; + sys::path::append(TraceDir, ".traces"); + + SmallString<32> StdinFile{TraceDir}; + sys::path::append(StdinFile, ".cur_input"); // Prepare afl-showmap arguments SmallVector AFLShowmapArgs{ - AFLShowmapPath, "-m", MemLimit, "-t", Timeout, "-q", "-o", TracePath}; + AFLShowmapPath, "-m", MemLimit, "-t", Timeout, + "-q", "-i", InputDir, "-o", TraceDir}; - if (TargetArgsHasAtAt) - AFLShowmapArgs.append({"-A", Seed}); - else - Redirects[/* stdin */ 0] = Seed; + if (TargetArgsHasAtAt) { + + AFLShowmapArgs.append({"-A", StdinFile}); + Redirects[/* stdin */ 0] = "/dev/null"; + + } if (FridaMode) AFLShowmapArgs.push_back("-O"); if (QemuMode) AFLShowmapArgs.push_back("-Q"); @@ -214,39 +256,48 @@ static Error getAFLCoverage(const StringRef Seed, AFLCoverageVector &Cov, /*env=*/None, Redirects); if (RC && !CrashMode) { - ErrMsg() << "Exit code " << RC << " != 0 received from afl-showmap"; - sys::fs::remove(TracePath); + ErrMsg() << "Exit code " << RC << " != 0 received from afl-showmap\n"; + sys::fs::remove_directories(TraceDir); return createStringError(inconvertibleErrorCode(), "afl-showmap failed"); } // Parse afl-showmap output - const auto CovOrErr = MemoryBuffer::getFile(TracePath); - if (const auto EC = CovOrErr.getError()) { + AFLCoverageVector Cov; + std::error_code EC; + sys::fs::file_status Status; - sys::fs::remove(TracePath); - return createStringError(EC, "Failed to read afl-showmap output file `%s`", - TracePath.c_str()); + for (sys::fs::recursive_directory_iterator Dir(TraceDir, EC), DirEnd; + Dir != DirEnd && !EC; Dir.increment(EC)) { - } + if (EC) return errorCodeToError(EC); - SmallVector Lines; - CovOrErr.get()->getBuffer().trim().split(Lines, '\n'); + const auto &Path = Dir->path(); + if ((EC = sys::fs::status(Path, Status))) return errorCodeToError(EC); - AFLTupleID Edge = 0; - unsigned Freq = 0; + switch (Status.type()) { - for (const auto &Line : Lines) { + case sys::fs::file_type::regular_file: + case sys::fs::file_type::symlink_file: + case sys::fs::file_type::type_unknown: + Cov.clear(); + if ((EC = readCov(Path, Cov))) { - const auto &[EdgeStr, FreqStr] = Line.split(':'); + sys::fs::remove(Path); + return errorCodeToError(EC); - to_integer(EdgeStr, Edge, 10); - to_integer(FreqStr, Freq, 10); - Cov.push_back({Edge, Freq}); + } + + CovMap.try_emplace(sys::path::filename(Path), Cov); + default: + /* Ignore */ + break; + + } } - if (!KeepTraces || BinCheck) sys::fs::remove(TracePath); + if (!KeepTraces) sys::fs::remove_directories(TraceDir); return Error::success(); } @@ -326,14 +377,14 @@ int main(int argc, char *argv[]) { if (WeightsFile != "") { - StatMsg() << "Reading weights from `" << WeightsFile << "`... "; + StatMsg() << "Reading weights from '" << WeightsFile << "'... "; StartTimer(/*ShowProgBar=*/false); const auto WeightsOrErr = MemoryBuffer::getFile(WeightsFile); if ((EC = WeightsOrErr.getError())) { - ErrMsg() << "Failed to read weights from `" << WeightsFile - << "`: " << EC.message() << '\n'; + ErrMsg() << "Failed to read weights from '" << WeightsFile + << "': " << EC.message() << '\n'; return 1; } @@ -345,47 +396,33 @@ int main(int argc, char *argv[]) { } // ------------------------------------------------------------------------ // - // Setup output directory + // Traverse input directory + // + // Find the seed files inside this directory (and subdirectories). // ------------------------------------------------------------------------ // - SmallString<32> TraceDir{OutputDir}; - sys::path::append(TraceDir, ".traces"); - - if ((EC = sys::fs::remove_directories(TraceDir))) { - - ErrMsg() << "Failed to remove existing trace directory in `" << OutputDir - << "`: " << EC.message() << '\n'; - return 1; - - } + StatMsg() << "Locating seeds in '" << InputDir << "'... "; + StartTimer(/*ShowProgBar=*/false); - if ((EC = sys::fs::create_directories(TraceDir))) { + bool IsDirResult; + if ((EC = sys::fs::is_directory(InputDir, IsDirResult))) { - ErrMsg() << "Failed to create output directory `" << OutputDir - << "`: " << EC.message() << '\n'; + ErrMsg() << "Invalid input directory '" << InputDir + << "': " << EC.message() << '\n'; return 1; } - // ------------------------------------------------------------------------ // - // Traverse corpus directory - // - // Find the seed files inside this directory. - // ------------------------------------------------------------------------ // - - StatMsg() << "Locating seeds in `" << CorpusDir << "`... "; - StartTimer(/*ShowProgBar=*/false); - - std::vector SeedFiles; - sys::fs::file_status Status; + sys::fs::file_status Status; + StringMap SeedFiles; - for (sys::fs::recursive_directory_iterator Dir(CorpusDir, EC), DirEnd; + for (sys::fs::recursive_directory_iterator Dir(InputDir, EC), DirEnd; Dir != DirEnd && !EC; Dir.increment(EC)) { if (EC) { - ErrMsg() << "Failed to traverse corpus directory `" << CorpusDir - << "`: " << EC.message() << '\n'; + ErrMsg() << "Failed to traverse input directory '" << InputDir + << "': " << EC.message() << '\n'; return 1; } @@ -393,21 +430,22 @@ int main(int argc, char *argv[]) { const auto &Path = Dir->path(); if ((EC = sys::fs::status(Path, Status))) { - WarnMsg() << "Failed to access seed file `" << Path - << "`: " << EC.message() << ". Skipping...\n"; - continue; + ErrMsg() << "Failed to access '" << Path << "': " << EC.message() + << '\n'; + return 1; } switch (Status.type()) { - case sys::fs::file_type::regular_file: - case sys::fs::file_type::symlink_file: - case sys::fs::file_type::type_unknown: - SeedFiles.push_back(Path); - default: - /* Ignore */ - break; + case sys::fs::file_type::regular_file: + case sys::fs::file_type::symlink_file: + case sys::fs::file_type::type_unknown: + SeedFiles.try_emplace(sys::path::filename(Path), + sys::path::parent_path(Path)); + default: + /* Ignore */ + break; } @@ -415,24 +453,55 @@ int main(int argc, char *argv[]) { EndTimer(/*ShowProgBar=*/false); + if (SeedFiles.empty()) { + + ErrMsg() << "Failed to find any seed files in '" << InputDir << "'\n"; + return 1; + + } + + // ------------------------------------------------------------------------ // + // Setup output directory + // ------------------------------------------------------------------------ // + + SmallString<32> TraceDir{OutputDir}; + sys::path::append(TraceDir, ".traces"); + + if ((EC = sys::fs::remove_directories(TraceDir))) { + + ErrMsg() << "Failed to remove existing trace directory in '" << OutputDir + << "': " << EC.message() << '\n'; + return 1; + + } + + if ((EC = sys::fs::create_directories(TraceDir))) { + + ErrMsg() << "Failed to create output directory '" << OutputDir + << "': " << EC.message() << '\n'; + return 1; + + } + // ------------------------------------------------------------------------ // // Test the target binary // ------------------------------------------------------------------------ // AFLCoverageVector Cov; - if (!SkipBinCheck && SeedFiles.size() > 0) { + if (!SkipBinCheck) { StatMsg() << "Testing the target binary... "; StartTimer(/*ShowProgBar=*/false); - if (auto Err = getAFLCoverage(SeedFiles.front(), Cov, /*BinCheck=*/true)) { - - ErrMsg() - << "No instrumentation output detected (perhaps crash or timeout)"; - return 1; - - } + // if (auto Err = runShowmap(TestSeed, Cov, /*BinCheck=*/true)) { + // + // ErrMsg() + // << "No instrumentation output detected (perhaps crash or + // timeout)"; + // return 1; + // + // } EndTimer(/*ShowProgBar=*/false); SuccMsg() << "OK, " << Cov.size() << " tuples recorded\n"; @@ -447,35 +516,28 @@ int main(int argc, char *argv[]) { // then store this coverage information in the appropriate data structures. // ------------------------------------------------------------------------ // - size_t SeedCount = 0; - const size_t NumSeeds = SeedFiles.size(); - - if (!ShowProgBar) - StatMsg() << "Generating coverage for " << NumSeeds << " seeds... "; - StartTimer(ShowProgBar); + StatMsg() << "Running afl-showmap on " << SeedFiles.size() << " seeds...\n"; + StartTimer(/*ShowProgBar=*/false); - EvalMaxSAT Solver(/*nbMinimizeThread=*/0); + AFLCoverageMap CovMap; MaxSATSeeds SeedVars; MaxSATCoverageMap SeedCoverage; + EvalMaxSAT Solver(/*nbMinimizeThread=*/0); - for (const auto &SeedFile : SeedFiles) { - - // Execute seed - Cov.clear(); - if (auto Err = getAFLCoverage(SeedFile, Cov)) { + if (auto Err = runShowmap(CovMap)) { - ErrMsg() << "Failed to get coverage for seed `" << SeedFile - << "`: " << Err << '\n'; - return 1; + ErrMsg() << "Failed to generate coverage: " << Err << '\n'; + return 1; - } + } + for (const auto &SeedCov : CovMap) { // Create a variable to represent the seed const SeedID Var = Solver.newVar(); - SeedVars.push_back({Var, SeedFile}); + SeedVars.emplace_back(Var, SeedCov.first()); // Record the set of seeds that cover a particular edge - for (const auto &[Edge, Freq] : Cov) { + for (auto &[Edge, Freq] : SeedCov.second) { if (EdgesOnly) { @@ -492,12 +554,10 @@ int main(int argc, char *argv[]) { } - if ((++SeedCount % 10 == 0) && ShowProgBar) - ProgBar.update(SeedCount * 100 / NumSeeds, "Generating seed coverage"); - } - EndTimer(ShowProgBar); + SuccMsg() << "afl-showmap completed in "; + EndTimer(/*ShowProgBar=*/false); // ------------------------------------------------------------------------ // // Set the hard and soft constraints in the solver @@ -506,7 +566,7 @@ int main(int argc, char *argv[]) { if (!ShowProgBar) StatMsg() << "Generating constraints... "; StartTimer(ShowProgBar); - SeedCount = 0; + size_t SeedCount = 0; // Ensure that at least one seed is selected that covers a particular edge // (hard constraint) @@ -552,7 +612,7 @@ int main(int argc, char *argv[]) { // ------------------------------------------------------------------------ // SmallVector Solution; - SmallString<32> OutputSeed; + SmallString<32> InputSeed, OutputSeed; if (Solved) { @@ -561,28 +621,31 @@ int main(int argc, char *argv[]) { } else { - ErrMsg() << "Failed to find an optimal solution for `" << CorpusDir - << "`\n"; + ErrMsg() << "Failed to find an optimal solution for '" << InputDir << "'\n"; return 1; } SuccMsg() << "Minimized corpus size: " << Solution.size() << " seeds\n"; - if (!ShowProgBar) StatMsg() << "Copying to `" << OutputDir << "`... "; + if (!ShowProgBar) StatMsg() << "Copying to '" << OutputDir << "'... "; StartTimer(ShowProgBar); SeedCount = 0; for (const auto &Seed : Solution) { + InputSeed = SeedFiles[Seed]; + sys::path::append(InputSeed, Seed); + OutputSeed = OutputDir; - sys::path::append(OutputSeed, sys::path::filename(Seed)); + sys::path::append(OutputSeed, Seed); - if ((EC = sys::fs::copy_file(Seed, OutputSeed))) { + if ((EC = sys::fs::copy_file(InputSeed, OutputSeed))) { - WarnMsg() << "Failed to copy `" << Seed << "` to `" << OutputDir - << "`: " << EC.message() << '\n'; + ErrMsg() << "Failed to copy '" << Seed << "' to '" << OutputDir + << "': " << EC.message() << '\n'; + return 1; } -- cgit 1.4.1 From 0dbd37a20b90284201a117192d126cbeee998663 Mon Sep 17 00:00:00 2001 From: Adrian Herrera Date: Wed, 28 Jul 2021 04:47:20 +0000 Subject: optimin: removed progress bar It was hard to get working with showmap's directory mode --- utils/optimin/src/OptiMin.cpp | 62 ++++++++++++++--------------------- utils/optimin/src/ProgressBar.h | 71 ----------------------------------------- 2 files changed, 24 insertions(+), 109 deletions(-) delete mode 100644 utils/optimin/src/ProgressBar.h (limited to 'utils/optimin/src/OptiMin.cpp') diff --git a/utils/optimin/src/OptiMin.cpp b/utils/optimin/src/OptiMin.cpp index 242eeea9..4f0cb4f7 100644 --- a/utils/optimin/src/OptiMin.cpp +++ b/utils/optimin/src/OptiMin.cpp @@ -23,7 +23,6 @@ #include #include "EvalMaxSAT.h" -#include "ProgressBar.h" using namespace llvm; @@ -131,7 +130,6 @@ static cl::opt InputDir("i", cl::desc("Input directory"), static cl::opt OutputDir("o", cl::desc("Output directory"), cl::value_desc("dir"), cl::Required); -static cl::opt ShowProgBar("p", cl::desc("Display progress bar")); static cl::opt EdgesOnly("f", cl::desc("Include edge hit counts"), cl::init(true)); static cl::opt WeightsFile("w", cl::desc("Weights file"), @@ -302,22 +300,19 @@ static Error runShowmap(AFLCoverageMap &CovMap) { } -static inline void StartTimer(bool ShowProgBar) { +static inline void StartTimer() { StartTime = std::chrono::system_clock::now(); } -static inline void EndTimer(bool ShowProgBar) { +static inline void EndTimer() { EndTime = std::chrono::system_clock::now(); Duration = std::chrono::duration_cast(EndTime - StartTime); - if (ShowProgBar) - outs() << '\n'; - else - outs() << Duration.count() << "s\n"; + SuccMsg() << " Completed in " << Duration.count() << " s\n"; } @@ -328,7 +323,6 @@ static inline void EndTimer(bool ShowProgBar) { int main(int argc, char *argv[]) { WeightsMap Weights; - ProgressBar ProgBar; std::error_code EC; // ------------------------------------------------------------------------ // @@ -377,8 +371,8 @@ int main(int argc, char *argv[]) { if (WeightsFile != "") { - StatMsg() << "Reading weights from '" << WeightsFile << "'... "; - StartTimer(/*ShowProgBar=*/false); + StatMsg() << "Reading weights from '" << WeightsFile << "'...\n"; + StartTimer(); const auto WeightsOrErr = MemoryBuffer::getFile(WeightsFile); if ((EC = WeightsOrErr.getError())) { @@ -391,7 +385,7 @@ int main(int argc, char *argv[]) { GetWeights(*WeightsOrErr.get(), Weights); - EndTimer(/*ShowProgBar=*/false); + EndTimer(); } @@ -401,8 +395,8 @@ int main(int argc, char *argv[]) { // Find the seed files inside this directory (and subdirectories). // ------------------------------------------------------------------------ // - StatMsg() << "Locating seeds in '" << InputDir << "'... "; - StartTimer(/*ShowProgBar=*/false); + StatMsg() << "Locating seeds in '" << InputDir << "'...\n"; + StartTimer(); bool IsDirResult; if ((EC = sys::fs::is_directory(InputDir, IsDirResult))) { @@ -451,7 +445,7 @@ int main(int argc, char *argv[]) { } - EndTimer(/*ShowProgBar=*/false); + EndTimer(); if (SeedFiles.empty()) { @@ -491,8 +485,8 @@ int main(int argc, char *argv[]) { if (!SkipBinCheck) { - StatMsg() << "Testing the target binary... "; - StartTimer(/*ShowProgBar=*/false); + StatMsg() << "Testing the target binary...\n"; + StartTimer(); // if (auto Err = runShowmap(TestSeed, Cov, /*BinCheck=*/true)) { // @@ -503,7 +497,7 @@ int main(int argc, char *argv[]) { // // } - EndTimer(/*ShowProgBar=*/false); + EndTimer(); SuccMsg() << "OK, " << Cov.size() << " tuples recorded\n"; } @@ -517,7 +511,7 @@ int main(int argc, char *argv[]) { // ------------------------------------------------------------------------ // StatMsg() << "Running afl-showmap on " << SeedFiles.size() << " seeds...\n"; - StartTimer(/*ShowProgBar=*/false); + StartTimer(); AFLCoverageMap CovMap; MaxSATSeeds SeedVars; @@ -557,14 +551,14 @@ int main(int argc, char *argv[]) { } SuccMsg() << "afl-showmap completed in "; - EndTimer(/*ShowProgBar=*/false); + EndTimer(); // ------------------------------------------------------------------------ // // Set the hard and soft constraints in the solver // ------------------------------------------------------------------------ // - if (!ShowProgBar) StatMsg() << "Generating constraints... "; - StartTimer(ShowProgBar); + StatMsg() << "Generating constraints...\n"; + StartTimer(); size_t SeedCount = 0; @@ -581,10 +575,6 @@ int main(int argc, char *argv[]) { Solver.addClause(Clauses); - if ((++SeedCount % 10 == 0) && ShowProgBar) - ProgBar.update(SeedCount * 100 / SeedCoverage.size(), - "Generating clauses"); - } // Select the minimum number of seeds that cover a particular set of edges @@ -592,18 +582,18 @@ int main(int argc, char *argv[]) { for (const auto &[Var, Seed] : SeedVars) Solver.addWeightedClause({-Var}, Weights[sys::path::filename(Seed)]); - EndTimer(ShowProgBar); + EndTimer(); // ------------------------------------------------------------------------ // // Generate a solution // ------------------------------------------------------------------------ // - StatMsg() << "Solving... "; - StartTimer(/*ShowProgBar=*/false); + StatMsg() << "Solving...\n"; + StartTimer(); const bool Solved = Solver.solve(); - EndTimer(/*ShowProgBar=*/false); + EndTimer(); // ------------------------------------------------------------------------ // // Save the solution @@ -626,10 +616,9 @@ int main(int argc, char *argv[]) { } - SuccMsg() << "Minimized corpus size: " << Solution.size() << " seeds\n"; - - if (!ShowProgBar) StatMsg() << "Copying to '" << OutputDir << "'... "; - StartTimer(ShowProgBar); + StatMsg() << "Copying " << Solution.size() << " seeds to '" + << OutputDir << "'...\n"; + StartTimer(); SeedCount = 0; @@ -649,12 +638,9 @@ int main(int argc, char *argv[]) { } - if ((++SeedCount % 10 == 0) && ShowProgBar) - ProgBar.update(SeedCount * 100 / Solution.size(), "Copying seeds"); - } - EndTimer(ShowProgBar); + EndTimer(); SuccMsg() << "Done!\n"; return 0; diff --git a/utils/optimin/src/ProgressBar.h b/utils/optimin/src/ProgressBar.h deleted file mode 100644 index 9b75594b..00000000 --- a/utils/optimin/src/ProgressBar.h +++ /dev/null @@ -1,71 +0,0 @@ -/** - * Progress bar. - * - * Adapted from https://www.bfilipek.com/2020/02/inidicators.html - */ - -#pragma once - -#include -#include - -/// Display a progress bar in the terminal -class ProgressBar { - - private: - const size_t BarWidth; - const std::string Fill; - const std::string Remainder; - - public: - ProgressBar() : ProgressBar(60, "#", " ") { - - } - - ProgressBar(size_t Width, const llvm::StringRef F, const llvm::StringRef R) - : BarWidth(Width), Fill(F), Remainder(R) { - - } - - void update(float Progress, const llvm::StringRef Status = "", - llvm::raw_ostream &OS = llvm::outs()) { - - // No need to write once progress is 100% - if (Progress > 100.0f) return; - - // Move cursor to the first position on the same line and flush - OS << '\r'; - OS.flush(); - - // Start bar - OS << '['; - - const auto Completed = - static_cast(Progress * static_cast(BarWidth) / 100.0); - for (size_t I = 0; I < BarWidth; ++I) { - - if (I <= Completed) { - - OS << Fill; - - } else { - - OS << Remainder; - - } - - } - - // End bar - OS << ']'; - - // Write progress percentage - OS << ' ' << std::min(static_cast(Progress), size_t(100)) << '%'; - - // Write status text - OS << " " << Status; - - } - -}; - -- cgit 1.4.1 From 0bd6fda98b4d376d40a87d5efb8cc4fa761a453e Mon Sep 17 00:00:00 2001 From: Adrian Herrera Date: Wed, 28 Jul 2021 06:44:25 +0000 Subject: optimin: add target test back in --- utils/optimin/src/OptiMin.cpp | 183 +++++++++++++++++++++++++++--------------- 1 file changed, 118 insertions(+), 65 deletions(-) (limited to 'utils/optimin/src/OptiMin.cpp') diff --git a/utils/optimin/src/OptiMin.cpp b/utils/optimin/src/OptiMin.cpp index 4f0cb4f7..ce1fb850 100644 --- a/utils/optimin/src/OptiMin.cpp +++ b/utils/optimin/src/OptiMin.cpp @@ -96,7 +96,7 @@ static constexpr unsigned MAX_WEIGHT_FAILURES = 5; static sys::TimePoint<> StartTime, EndTime; static std::chrono::seconds Duration; -static std::string AFLShowmapPath; +static std::string ShowmapPath; static bool TargetArgsHasAtAt = false; static bool KeepTraces = false; static bool SkipBinCheck = false; @@ -220,82 +220,133 @@ static std::error_code readCov(const StringRef Trace, AFLCoverageVector &Cov) { } -static Error runShowmap(AFLCoverageMap &CovMap) { +static Error runShowmap(AFLCoverageMap &CovMap, const StringRef Input, + bool BinCheck = false) { + const bool InputIsFile = !sys::fs::is_directory(Input); Optional Redirects[] = {None, None, None}; SmallString<32> TraceDir{OutputDir}; sys::path::append(TraceDir, ".traces"); + SmallString<32> Output{TraceDir}; SmallString<32> StdinFile{TraceDir}; - sys::path::append(StdinFile, ".cur_input"); + // ------------------------------------------------------------------------ // // Prepare afl-showmap arguments - SmallVector AFLShowmapArgs{ - AFLShowmapPath, "-m", MemLimit, "-t", Timeout, - "-q", "-i", InputDir, "-o", TraceDir}; + // + // If the given input is a file, then feed this directly into stdin. + // Otherwise, if it is a directory, specify this on the afl-showmap command + // line. + // ------------------------------------------------------------------------ // + + SmallVector ShowmapArgs{ShowmapPath, "-q", + "-m", MemLimit, + "-t", Timeout}; + + if (InputIsFile) { + + StdinFile = Input; + sys::path::append(Output, + BinCheck ? ".run_test" : sys::path::filename(Input)); + + } else { + + sys::path::append(StdinFile, ".cur_input"); + ShowmapArgs.append({"-i", Input}); + + } + if (TargetArgsHasAtAt) { - AFLShowmapArgs.append({"-A", StdinFile}); + ShowmapArgs.append({"-A", StdinFile}); Redirects[/* stdin */ 0] = "/dev/null"; + } else if (InputIsFile) { + + Redirects[/* stdin */ 0] = Input; + } - if (FridaMode) AFLShowmapArgs.push_back("-O"); - if (QemuMode) AFLShowmapArgs.push_back("-Q"); - if (UnicornMode) AFLShowmapArgs.push_back("-U"); + if (FridaMode) ShowmapArgs.push_back("-O"); + if (QemuMode) ShowmapArgs.push_back("-Q"); + if (UnicornMode) ShowmapArgs.push_back("-U"); - AFLShowmapArgs.append({"--", TargetProg}); - AFLShowmapArgs.append(TargetArgs.begin(), TargetArgs.end()); + ShowmapArgs.append({"-o", Output, "--", TargetProg}); + ShowmapArgs.append(TargetArgs.begin(), TargetArgs.end()); + // ------------------------------------------------------------------------ // // Run afl-showmap - const int RC = sys::ExecuteAndWait(AFLShowmapPath, AFLShowmapArgs, + // ------------------------------------------------------------------------ // + + const int RC = sys::ExecuteAndWait(ShowmapPath, ShowmapArgs, /*env=*/None, Redirects); if (RC && !CrashMode) { ErrMsg() << "Exit code " << RC << " != 0 received from afl-showmap\n"; - sys::fs::remove_directories(TraceDir); return createStringError(inconvertibleErrorCode(), "afl-showmap failed"); } + // ------------------------------------------------------------------------ // // Parse afl-showmap output + // ------------------------------------------------------------------------ // + AFLCoverageVector Cov; std::error_code EC; sys::fs::file_status Status; - for (sys::fs::recursive_directory_iterator Dir(TraceDir, EC), DirEnd; - Dir != DirEnd && !EC; Dir.increment(EC)) { + if (InputIsFile) { - if (EC) return errorCodeToError(EC); + // Read a single output coverage file + if ((EC = readCov(Output, Cov))) { - const auto &Path = Dir->path(); - if ((EC = sys::fs::status(Path, Status))) return errorCodeToError(EC); + sys::fs::remove(Output); + return errorCodeToError(EC); - switch (Status.type()) { + } - case sys::fs::file_type::regular_file: - case sys::fs::file_type::symlink_file: - case sys::fs::file_type::type_unknown: - Cov.clear(); - if ((EC = readCov(Path, Cov))) { + CovMap.try_emplace(sys::path::filename(Input), Cov); + if (!KeepTraces) sys::fs::remove(Output); - sys::fs::remove(Path); - return errorCodeToError(EC); + } else { - } + // Read a directory of output coverage files + for (sys::fs::recursive_directory_iterator Dir(TraceDir, EC), DirEnd; + Dir != DirEnd && !EC; Dir.increment(EC)) { - CovMap.try_emplace(sys::path::filename(Path), Cov); - default: - /* Ignore */ - break; + if (EC) return errorCodeToError(EC); + + const auto &Path = Dir->path(); + if ((EC = sys::fs::status(Path, Status))) return errorCodeToError(EC); + + switch (Status.type()) { + + case sys::fs::file_type::regular_file: + case sys::fs::file_type::symlink_file: + case sys::fs::file_type::type_unknown: + Cov.clear(); + if ((EC = readCov(Path, Cov))) { + + sys::fs::remove(Path); + return errorCodeToError(EC); + + } + + CovMap.try_emplace(sys::path::filename(Path), Cov); + default: + // Ignore + break; + + } } + if (!KeepTraces) sys::fs::remove_directories(TraceDir); + } - if (!KeepTraces) sys::fs::remove_directories(TraceDir); return Error::success(); } @@ -352,15 +403,15 @@ int main(int argc, char *argv[]) { SplitString(PathEnv, EnvPaths, ":"); if (AFLPath) EnvPaths.push_back(AFLPath); - const auto AFLShowmapOrErr = sys::findProgramByName("afl-showmap", EnvPaths); - if (AFLShowmapOrErr.getError()) { + const auto ShowmapOrErr = sys::findProgramByName("afl-showmap", EnvPaths); + if (ShowmapOrErr.getError()) { ErrMsg() << "Failed to find afl-showmap. Check your PATH\n"; return 1; } - AFLShowmapPath = *AFLShowmapOrErr; + ShowmapPath = *ShowmapOrErr; // ------------------------------------------------------------------------ // // Parse weights @@ -401,13 +452,13 @@ int main(int argc, char *argv[]) { bool IsDirResult; if ((EC = sys::fs::is_directory(InputDir, IsDirResult))) { - ErrMsg() << "Invalid input directory '" << InputDir - << "': " << EC.message() << '\n'; + ErrMsg() << "Invalid input directory '" << InputDir << "': " << EC.message() + << '\n'; return 1; } - sys::fs::file_status Status; + sys::fs::file_status Status; StringMap SeedFiles; for (sys::fs::recursive_directory_iterator Dir(InputDir, EC), DirEnd; @@ -424,22 +475,21 @@ int main(int argc, char *argv[]) { const auto &Path = Dir->path(); if ((EC = sys::fs::status(Path, Status))) { - ErrMsg() << "Failed to access '" << Path << "': " << EC.message() - << '\n'; + ErrMsg() << "Failed to access '" << Path << "': " << EC.message() << '\n'; return 1; } switch (Status.type()) { - case sys::fs::file_type::regular_file: - case sys::fs::file_type::symlink_file: - case sys::fs::file_type::type_unknown: - SeedFiles.try_emplace(sys::path::filename(Path), - sys::path::parent_path(Path)); - default: - /* Ignore */ - break; + case sys::fs::file_type::regular_file: + case sys::fs::file_type::symlink_file: + case sys::fs::file_type::type_unknown: + SeedFiles.try_emplace(sys::path::filename(Path), + sys::path::parent_path(Path)); + default: + /* Ignore */ + break; } @@ -481,24 +531,27 @@ int main(int argc, char *argv[]) { // Test the target binary // ------------------------------------------------------------------------ // - AFLCoverageVector Cov; + AFLCoverageMap CovMap; if (!SkipBinCheck) { - StatMsg() << "Testing the target binary...\n"; + const auto It = SeedFiles.begin(); + SmallString<32> TestSeed{It->second}; + sys::path::append(TestSeed, It->first()); + + StatMsg() << "Testing the target binary with '" << TestSeed << "`...\n"; StartTimer(); - // if (auto Err = runShowmap(TestSeed, Cov, /*BinCheck=*/true)) { - // - // ErrMsg() - // << "No instrumentation output detected (perhaps crash or - // timeout)"; - // return 1; - // - // } + if (auto Err = runShowmap(CovMap, TestSeed, /*BinCheck=*/true)) { + + ErrMsg() << "No instrumentation output detected \n"; + return 1; + + } EndTimer(); - SuccMsg() << "OK, " << Cov.size() << " tuples recorded\n"; + SuccMsg() << "OK, " << CovMap.begin()->second.size() + << " tuples recorded\n"; } @@ -513,12 +566,12 @@ int main(int argc, char *argv[]) { StatMsg() << "Running afl-showmap on " << SeedFiles.size() << " seeds...\n"; StartTimer(); - AFLCoverageMap CovMap; MaxSATSeeds SeedVars; MaxSATCoverageMap SeedCoverage; EvalMaxSAT Solver(/*nbMinimizeThread=*/0); - if (auto Err = runShowmap(CovMap)) { + CovMap.clear(); + if (auto Err = runShowmap(CovMap, InputDir)) { ErrMsg() << "Failed to generate coverage: " << Err << '\n'; return 1; @@ -526,6 +579,7 @@ int main(int argc, char *argv[]) { } for (const auto &SeedCov : CovMap) { + // Create a variable to represent the seed const SeedID Var = Solver.newVar(); SeedVars.emplace_back(Var, SeedCov.first()); @@ -550,7 +604,6 @@ int main(int argc, char *argv[]) { } - SuccMsg() << "afl-showmap completed in "; EndTimer(); // ------------------------------------------------------------------------ // @@ -616,8 +669,8 @@ int main(int argc, char *argv[]) { } - StatMsg() << "Copying " << Solution.size() << " seeds to '" - << OutputDir << "'...\n"; + StatMsg() << "Copying " << Solution.size() << " seeds to '" << OutputDir + << "'...\n"; StartTimer(); SeedCount = 0; -- cgit 1.4.1