added support for __afl_coverage_interesting

6 files changed, 76 insertions, 13 deletions

diff --git a/docs/Changelog.md b/docs/Changelog.md
index 6e59961b..99bc8b47 100644
--- a/docs/Changelog.md
+++ b/docs/Changelog.md
@@ -36,6 +36,9 @@ sending a mail to <afl-users+subscribe@googlegroups.com>.
     - cmplog/redqueen now also tracks floating point, _ExtInt() + 128bit
     - cmplog/redqueen can now process basic libc++ and libstdc++
       std::string comparisons (though no position or length type variants)
+    - added support for __afl_coverage_interesting() for LTO and
+      and our own PCGUARD (llvm 10.0.1+), read more about this function
+      and selective coverage in instrumentation/README.instrument_list.md
     - added AFL_LLVM_INSTRUMENT option NATIVE for native clang pc-guard
       support (less performant than our own), GCC for old afl-gcc and
       CLANG for old afl-clang
diff --git a/instrumentation/README.instrument_list.md b/instrumentation/README.instrument_list.md
index b47b50f6..b7dfb40c 100644
--- a/instrumentation/README.instrument_list.md
+++ b/instrumentation/README.instrument_list.md
@@ -43,6 +43,13 @@ in any function where you want:
   * `__AFL_COVERAGE_DISCARD();` - reset all coverage gathered until this point
   * `__AFL_COVERAGE_SKIP();` - mark this test case as unimportant. Whatever happens, afl-fuzz will ignore it.
 
+A special function is `__afl_coverage_interesting`.
+To use this, you must define `void __afl_coverage_interesting(u8 val, u32 id);`.
+Then you can use this function globally, where the `val` parameter can be set
+by you, the `id` parameter is for afl-fuzz and will be overwritten.
+Note that useful parameters are for `val` are: 1, 2, 3, 4, 8, 16, 32, 64, 128.
+A value of e.g. 33 will be seen as 32 for coverage purposes.
+
 ## 3) Selective instrumenation with AFL_LLVM_ALLOWLIST/AFL_LLVM_DENYLIST
 
 This feature is equivalent to llvm 12 sancov feature and allows to specify
diff --git a/instrumentation/SanitizerCoverageLTO.so.cc b/instrumentation/SanitizerCoverageLTO.so.cc
index 016ac71f..e3490847 100644
--- a/instrumentation/SanitizerCoverageLTO.so.cc
+++ b/instrumentation/SanitizerCoverageLTO.so.cc
@@ -1237,6 +1237,25 @@ void ModuleSanitizerCoverage::instrumentFunction(
 
   for (auto &BB : F) {
 
+    for (auto &IN : BB) {
+
+      CallInst *callInst = nullptr;
+
+      if ((callInst = dyn_cast<CallInst>(&IN))) {
+
+        Function *Callee = callInst->getCalledFunction();
+        if (!Callee) continue;
+        if (callInst->getCallingConv() != llvm::CallingConv::C) continue;
+        StringRef FuncName = Callee->getName();
+        if (FuncName.compare(StringRef("__afl_coverage_interesting"))) continue;
+
+        Value *val = ConstantInt::get(Int32Ty, ++afl_global_id);
+        callInst->setOperand(1, val);
+
+      }
+
+    }
+
     if (shouldInstrumentBlock(F, &BB, DT, PDT, Options))
       BlocksToInstrument.push_back(&BB);
     for (auto &Inst : BB) {
@@ -1338,6 +1357,7 @@ bool ModuleSanitizerCoverage::InjectCoverage(Function &             F,
 
   if (AllBlocks.empty()) return false;
   CreateFunctionLocalArrays(F, AllBlocks);
+
   for (size_t i = 0, N = AllBlocks.size(); i < N; i++) {
 
     // afl++ START
diff --git a/instrumentation/SanitizerCoveragePCGUARD.so.cc b/instrumentation/SanitizerCoveragePCGUARD.so.cc
index ecd6bc9b..5b274770 100644
--- a/instrumentation/SanitizerCoveragePCGUARD.so.cc
+++ b/instrumentation/SanitizerCoveragePCGUARD.so.cc
@@ -311,7 +311,8 @@ class ModuleSanitizerCoverage {
                                                     Function &F, Type *Ty,
                                                     const char *Section);
   GlobalVariable *CreatePCArray(Function &F, ArrayRef<BasicBlock *> AllBlocks);
-  void CreateFunctionLocalArrays(Function &F, ArrayRef<BasicBlock *> AllBlocks);
+  void CreateFunctionLocalArrays(Function &F, ArrayRef<BasicBlock *> AllBlocks,
+                                 uint32_t special);
   void InjectCoverageAtBlock(Function &F, BasicBlock &BB, size_t Idx,
                              bool IsLeafFunc = true);
   Function *CreateInitCallsForSections(Module &M, const char *CtorName,
@@ -970,11 +971,11 @@ GlobalVariable *ModuleSanitizerCoverage::CreatePCArray(
 }
 
 void ModuleSanitizerCoverage::CreateFunctionLocalArrays(
-    Function &F, ArrayRef<BasicBlock *> AllBlocks) {
+    Function &F, ArrayRef<BasicBlock *> AllBlocks, uint32_t special) {
 
   if (Options.TracePCGuard)
     FunctionGuardArray = CreateFunctionLocalArrayInSection(
-        AllBlocks.size(), F, Int32Ty, SanCovGuardsSectionName);
+        AllBlocks.size() + special, F, Int32Ty, SanCovGuardsSectionName);
 
   if (Options.Inline8bitCounters)
     Function8bitCounterArray = CreateFunctionLocalArrayInSection(
@@ -993,9 +994,38 @@ bool ModuleSanitizerCoverage::InjectCoverage(Function &             F,
                                              bool IsLeafFunc) {
 
   if (AllBlocks.empty()) return false;
-  CreateFunctionLocalArrays(F, AllBlocks);
+
+  uint32_t special = 0;
+  for (auto &BB : F) {
+
+    for (auto &IN : BB) {
+
+      CallInst *callInst = nullptr;
+
+      if ((callInst = dyn_cast<CallInst>(&IN))) {
+
+        Function *Callee = callInst->getCalledFunction();
+        StringRef FuncName = Callee->getName();
+        if (!Callee) continue;
+        if (callInst->getCallingConv() != llvm::CallingConv::C) continue;
+        if (FuncName.compare(StringRef("__afl_coverage_interesting"))) continue;
+
+        uint32_t id = 1 + instr + (uint32_t)AllBlocks.size() + special++;
+        Value *  val = ConstantInt::get(Int32Ty, id);
+        callInst->setOperand(1, val);
+
+      }
+
+    }
+
+  }
+
+  CreateFunctionLocalArrays(F, AllBlocks, special);
   for (size_t i = 0, N = AllBlocks.size(); i < N; i++)
     InjectCoverageAtBlock(F, *AllBlocks[i], i, IsLeafFunc);
+
+  instr += special;
+
   return true;
 
 }
diff --git a/src/afl-cc.c b/src/afl-cc.c
index b5dcb632..f513764a 100644
--- a/src/afl-cc.c
+++ b/src/afl-cc.c
@@ -572,7 +572,8 @@ static void edit_params(u32 argc, char **argv, char **envp) {
 
       cc_params[cc_par_cnt++] = "-Wl,--allow-multiple-definition";
 
-      if (instrument_mode == INSTRUMENT_CFG)
+      if (instrument_mode == INSTRUMENT_CFG ||
+          instrument_mode == INSTRUMENT_PCGUARD)
         cc_params[cc_par_cnt++] = alloc_printf(
             "-Wl,-mllvm=-load=%s/SanitizerCoverageLTO.so", obj_path);
       else
@@ -1670,15 +1671,16 @@ int main(int argc, char **argv, char **envp) {
   if (compiler_mode == LTO) {
 
     if (instrument_mode == 0 || instrument_mode == INSTRUMENT_LTO ||
-        instrument_mode == INSTRUMENT_CFG) {
+        instrument_mode == INSTRUMENT_CFG ||
+        instrument_mode == INSTRUMENT_PCGUARD) {
 
       lto_mode = 1;
-      if (!instrument_mode) {
+      // force CFG
+      // if (!instrument_mode) {
 
-        instrument_mode = INSTRUMENT_CFG;
-        // ptr = instrument_mode_string[instrument_mode];
-
-      }
+      instrument_mode = INSTRUMENT_PCGUARD;
+      // ptr = instrument_mode_string[instrument_mode];
+      // }
 
     } else if (instrument_mode == INSTRUMENT_LTO ||
 
diff --git a/src/afl-fuzz.c b/src/afl-fuzz.c
index 647a665e..82eff61f 100644
--- a/src/afl-fuzz.c
+++ b/src/afl-fuzz.c
@@ -145,7 +145,8 @@ static void usage(u8 *argv0, int more_help) {
 
       "Other stuff:\n"
       "  -M/-S id      - distributed mode (see docs/parallel_fuzzing.md)\n"
-      "                  -M auto-sets -D, -Z (use -d to disable -D) and no trimming\n"
+      "                  -M auto-sets -D, -Z (use -d to disable -D) and no "
+      "trimming\n"
       "  -F path       - sync to a foreign fuzzer queue directory (requires "
       "-M, can\n"
       "                  be specified up to %u times)\n"
@@ -502,7 +503,7 @@ int main(int argc, char **argv_orig, char **envp) {
         afl->sync_id = ck_strdup(optarg);
         afl->skip_deterministic = 0;  // force deterministic fuzzing
         afl->old_seed_selection = 1;  // force old queue walking seed selection
-        afl->disable_trim = 1;  // disable trimming
+        afl->disable_trim = 1;        // disable trimming
 
         if ((c = strchr(afl->sync_id, ':'))) {