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).
MODIFY
utils/optimin/src/OptiMin.cpp
utils/optimin/src/OptiMin.cpp
@@ -64,7 +64,10 @@ using AFLTuple = std::pair<AFLTupleID, /* Frequency */ unsigned>;
/// Coverage for a given seed file
using AFLCoverageVector = std::vector<AFLTuple>;
-/// Maps seed file paths to a weight
+/// Map seed file paths to its coverage vector
+using AFLCoverageMap = StringMap<AFLCoverageVector>;
+
+/// Map seed file paths to a weight
using WeightsMap = StringMap<Weight>;
/// A seed identifier in the MaxSAT solver
@@ -88,6 +91,9 @@ using MaxSATCoverageMap = DenseMap<AFLTupleID, MaxSATSeedSet>;
// `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,29 +104,29 @@ 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<std::string> CorpusDir("i", cl::desc("Input directory"),
+static cl::opt<std::string> InputDir("i", cl::desc("Input directory"),
cl::value_desc("dir"), cl::Required);
static cl::opt<std::string> 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<StringRef, 0> 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<StringRef, 0> 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<StringRef> 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<StringRef, 12> 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<StringRef, 0> 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<std::string> SeedFiles;
sys::fs::file_status Status;
+ StringMap<std::string> 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,9 +430,9 @@ 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;
}
@@ -404,7 +441,8 @@ int main(int argc, char *argv[]) {
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);
+ SeedFiles.try_emplace(sys::path::filename(Path),
+ sys::path::parent_path(Path));
default:
/* Ignore */
break;
@@ -415,25 +453,56 @@ int main(int argc, char *argv[]) {
EndTimer(/*ShowProgBar=*/false);
+ if (SeedFiles.empty()) {
+
+ ErrMsg() << "Failed to find any seed files in '" << InputDir << "'\n";
+ return 1;
+
+ }
+
// ------------------------------------------------------------------------ //
- // Test the target binary
+ // Setup output directory
// ------------------------------------------------------------------------ //
- AFLCoverageVector Cov;
+ SmallString<32> TraceDir{OutputDir};
+ sys::path::append(TraceDir, ".traces");
- if (!SkipBinCheck && SeedFiles.size() > 0) {
+ if ((EC = sys::fs::remove_directories(TraceDir))) {
- StatMsg() << "Testing the target binary... ";
- StartTimer(/*ShowProgBar=*/false);
+ ErrMsg() << "Failed to remove existing trace directory in '" << OutputDir
+ << "': " << EC.message() << '\n';
+ return 1;
+
+ }
- if (auto Err = getAFLCoverage(SeedFiles.front(), Cov, /*BinCheck=*/true)) {
+ if ((EC = sys::fs::create_directories(TraceDir))) {
- ErrMsg()
- << "No instrumentation output detected (perhaps crash or timeout)";
+ ErrMsg() << "Failed to create output directory '" << OutputDir
+ << "': " << EC.message() << '\n';
return 1;
}
+ // ------------------------------------------------------------------------ //
+ // Test the target binary
+ // ------------------------------------------------------------------------ //
+
+ AFLCoverageVector Cov;
+
+ if (!SkipBinCheck) {
+
+ StatMsg() << "Testing the target binary... ";
+ StartTimer(/*ShowProgBar=*/false);
+
+ // 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';
+ 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<StringRef, 64> 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;
}

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).