aboutsummaryrefslogtreecommitdiff
path: root/utils/optimin/src/OptiMin.cpp
diff options
context:
space:
mode:
authorvanhauser-thc <vh@thc.org>2021-07-21 09:55:22 +0200
committervanhauser-thc <vh@thc.org>2021-07-21 09:55:22 +0200
commit60cbe5b4bec445438718d27dc0cfda6a4d5de81e (patch)
tree43b0af8d301ae31c8874eb94b17fb33b7ba95170 /utils/optimin/src/OptiMin.cpp
parentfa2b164429f8488bf8a360c20933c7db238ed17c (diff)
downloadafl++-60cbe5b4bec445438718d27dc0cfda6a4d5de81e.tar.gz
optimin nits
Diffstat (limited to 'utils/optimin/src/OptiMin.cpp')
-rw-r--r--utils/optimin/src/OptiMin.cpp75
1 files changed, 75 insertions, 0 deletions
diff --git a/utils/optimin/src/OptiMin.cpp b/utils/optimin/src/OptiMin.cpp
index e02fcbe5..4fbf3416 100644
--- a/utils/optimin/src/OptiMin.cpp
+++ b/utils/optimin/src/OptiMin.cpp
@@ -33,16 +33,20 @@ namespace {
/// Ensure seed weights default to 1
class Weight {
+
public:
Weight() : Weight(1){};
Weight(uint32_t V) : Value(V){};
operator unsigned() const {
+
return Value;
+
}
private:
const unsigned Value;
+
};
// -------------------------------------------------------------------------- //
@@ -89,16 +93,27 @@ static std::string AFLShowmapPath;
static bool TargetArgsHasAtAt = false;
static const auto ErrMsg = [] {
+
return WithColor(errs(), HighlightColor::Error) << "[-] ";
+
};
+
static const auto WarnMsg = [] {
+
return WithColor(errs(), HighlightColor::Warning) << "[-] ";
+
};
+
static const auto SuccMsg = [] {
+
return WithColor(outs(), HighlightColor::String) << "[+] ";
+
};
+
static const auto StatMsg = [] {
+
return WithColor(outs(), HighlightColor::Remark) << "[*] ";
+
};
static cl::opt<std::string> CorpusDir("i", cl::desc("Input directory"),
@@ -124,6 +139,7 @@ static cl::opt<std::string> Timeout(
static cl::opt<bool> CrashMode(
"C", cl::desc("Keep crashing inputs, reject everything else"));
static cl::opt<bool> QemuMode("Q", cl::desc("Use binary-only instrumentation"));
+
} // anonymous namespace
// -------------------------------------------------------------------------- //
@@ -131,24 +147,33 @@ static cl::opt<bool> QemuMode("Q", cl::desc("Use binary-only instrumentation"));
// -------------------------------------------------------------------------- //
static void GetWeights(const MemoryBuffer &MB, WeightsMap &Weights) {
+
SmallVector<StringRef, 0> Lines;
MB.getBuffer().trim().split(Lines, '\n');
unsigned Weight = 0;
for (const auto &Line : Lines) {
+
const auto &[Seed, WeightStr] = Line.split(',');
if (to_integer(WeightStr, Weight, 10)) {
+
Weights.try_emplace(Seed, Weight);
+
} else {
+
WarnMsg() << "Failed to read weight for `" << Seed << "`. Skipping...\n";
+
}
+
}
+
}
[[nodiscard]] static std::error_code getAFLCoverage(const StringRef Seed,
AFLCoverageVector &Cov) {
+
Optional<StringRef> Redirects[] = {None, None, None};
std::error_code EC;
@@ -159,6 +184,7 @@ static void GetWeights(const MemoryBuffer &MB, WeightsMap &Weights) {
// Prepare afl-showmap arguments
SmallVector<StringRef, 12> AFLShowmapArgs{
+
AFLShowmapPath, "-m", MemLimit, "-t", Timeout, "-q", "-o", OutputPath};
if (TargetArgsHasAtAt)
@@ -180,8 +206,10 @@ static void GetWeights(const MemoryBuffer &MB, WeightsMap &Weights) {
// Parse afl-showmap output
const auto CovOrErr = MemoryBuffer::getFile(OutputPath);
if (EC = CovOrErr.getError()) {
+
sys::fs::remove(OutputPath);
return EC;
+
}
SmallVector<StringRef, 0> Lines;
@@ -191,21 +219,27 @@ static void GetWeights(const MemoryBuffer &MB, WeightsMap &Weights) {
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 sys::fs::remove(OutputPath);
+
}
static inline void StartTimer(bool ShowProgBar) {
+
StartTime = std::chrono::system_clock::now();
+
}
static inline void EndTimer(bool ShowProgBar) {
+
EndTime = std::chrono::system_clock::now();
Duration =
std::chrono::duration_cast<std::chrono::seconds>(EndTime - StartTime);
@@ -214,6 +248,7 @@ static inline void EndTimer(bool ShowProgBar) {
outs() << '\n';
else
outs() << Duration.count() << "s\n";
+
}
// -------------------------------------------------------------------------- //
@@ -221,6 +256,7 @@ static inline void EndTimer(bool ShowProgBar) {
// -------------------------------------------------------------------------- //
int main(int argc, char *argv[]) {
+
WeightsMap Weights;
ProgressBar ProgBar;
std::error_code EC;
@@ -234,8 +270,10 @@ int main(int argc, char *argv[]) {
cl::ParseCommandLineOptions(argc, argv, "Optimal corpus minimizer");
if (!sys::fs::is_directory(OutputDir)) {
+
ErrMsg() << "Invalid output directory `" << OutputDir << "`\n";
return 1;
+
}
for (const auto &Arg : TargetArgs)
@@ -247,9 +285,12 @@ int main(int argc, char *argv[]) {
const auto AFLShowmapOrErr = sys::findProgramByName("afl-showmap");
if (AFLShowmapOrErr.getError()) {
+
ErrMsg() << "Failed to find afl-showmap. Check your PATH\n";
return 1;
+
}
+
AFLShowmapPath = *AFLShowmapOrErr;
// ------------------------------------------------------------------------ //
@@ -260,19 +301,23 @@ int main(int argc, char *argv[]) {
// ------------------------------------------------------------------------ //
if (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';
return 1;
+
}
GetWeights(*WeightsOrErr.get(), Weights);
EndTimer(/*ShowProgBar=*/false);
+
}
// ------------------------------------------------------------------------ //
@@ -289,20 +334,26 @@ int main(int argc, char *argv[]) {
for (sys::fs::directory_iterator Dir(CorpusDir, EC), DirEnd;
Dir != DirEnd && !EC; Dir.increment(EC)) {
+
if (EC) {
+
ErrMsg() << "Failed to traverse corpus directory `" << CorpusDir
<< "`: " << EC.message() << '\n';
return 1;
+
}
const auto &Path = Dir->path();
if (EC = sys::fs::status(Path, Status)) {
+
WarnMsg() << "Failed to access seed file `" << Path
<< "`: " << EC.message() << ". Skipping...\n";
continue;
+
}
switch (Status.type()) {
+
case sys::fs::file_type::regular_file:
case sys::fs::file_type::symlink_file:
case sys::fs::file_type::type_unknown:
@@ -310,7 +361,9 @@ int main(int argc, char *argv[]) {
default:
/* Ignore */
break;
+
}
+
}
EndTimer(/*ShowProgBar=*/false);
@@ -336,12 +389,15 @@ int main(int argc, char *argv[]) {
AFLCoverageVector Cov;
for (const auto &SeedFile : SeedFiles) {
+
// Execute seed
Cov.clear();
if (EC = getAFLCoverage(SeedFile, Cov)) {
+
ErrMsg() << "Failed to get coverage for seed " << SeedFile << ": "
<< EC.message() << '\n';
return 1;
+
}
// Create a variable to represent the seed
@@ -350,18 +406,25 @@ int main(int argc, char *argv[]) {
// Record the set of seeds that cover a particular edge
for (const auto &[Edge, Freq] : Cov) {
+
if (EdgesOnly) {
+
// Ignore edge frequency
SeedCoverage[Edge].insert(Var);
+
} else {
+
// Executing edge `E` `N` times means that it was executed `N - 1` times
for (unsigned I = 0; I < Freq; ++I)
SeedCoverage[MAX_EDGE_FREQ * Edge + I].insert(Var);
+
}
+
}
if ((++SeedCount % 10 == 0) && ShowProgBar)
ProgBar.update(SeedCount * 100 / NumSeeds, "Generating seed coverage");
+
}
EndTimer(ShowProgBar);
@@ -379,6 +442,7 @@ int main(int argc, char *argv[]) {
// (hard constraint)
std::vector<SeedID> Clauses;
for (const auto &[_, Seeds] : SeedCoverage) {
+
if (Seeds.empty()) continue;
Clauses.clear();
@@ -390,6 +454,7 @@ int main(int argc, char *argv[]) {
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
@@ -420,12 +485,16 @@ int main(int argc, char *argv[]) {
SmallString<32> OutputSeed;
if (Solved) {
+
for (const auto &[Var, Seed] : SeedVars)
if (Solver.getValue(Var) > 0) Solution.push_back(Seed);
+
} else {
+
ErrMsg() << "Failed to find an optimal solution for `" << CorpusDir
<< "`\n";
return 1;
+
}
SuccMsg() << "Minimized corpus size: " << Solution.size() << " seeds\n";
@@ -436,20 +505,26 @@ int main(int argc, char *argv[]) {
SeedCount = 0;
for (const auto &Seed : Solution) {
+
OutputSeed = OutputDir;
sys::path::append(OutputSeed, sys::path::filename(Seed));
if (EC = sys::fs::copy_file(Seed, OutputSeed)) {
+
WarnMsg() << "Failed to copy `" << Seed << "` to `" << OutputDir
<< "`: " << EC.message() << '\n';
+
}
if ((++SeedCount % 10 == 0) && ShowProgBar)
ProgBar.update(SeedCount * 100 / Solution.size(), "Copying seeds");
+
}
EndTimer(ShowProgBar);
SuccMsg() << "Done!\n";
return 0;
+
}
+