add lto caller instrumentation

4 files changed, 35 insertions, 5 deletions

diff --git a/docs/Changelog.md b/docs/Changelog.md
index 2f0fba33..e5169daf 100644
--- a/docs/Changelog.md
+++ b/docs/Changelog.md
@@ -7,6 +7,11 @@
   - afl-fuzz:
     - the new deterministic fuzzing feature is now activated by default,
       deactivate with -z. Parameters -d and -D are ignored.
+  - afl-cc:
+    - added collision free caller instrumentation to LTO mode. activate with
+      `AFL_LLVM_LTO_CALLER=1`. You can set a max depth to go through single
+      block functions with `AFL_LLVM_LTO_CALLER_DEPTH` (default 0)
+
 
 ### Version ++4.10c (release)
   - afl-fuzz:
diff --git a/docs/env_variables.md b/docs/env_variables.md
index a972b6da..1e4fc7ba 100644
--- a/docs/env_variables.md
+++ b/docs/env_variables.md
@@ -248,6 +248,9 @@ use (which only ever the author of this LTO implementation will use). These are
 used if several separated instrumentations are performed which are then later
 combined.
 
+  - `AFL_LLVM_LTO_CALLER` activates collision free CALLER instrumentation
+  - `AFL_LLVM_LTO_CALLER` sets the maximum mumber of single block functions
+    to dig deeper into a real function. Default 0.
   - `AFL_LLVM_DOCUMENT_IDS=file` will document to a file which edge ID was given
     to which function. This helps to identify functions with variable bytes or
     which functions were touched by an input.
diff --git a/instrumentation/SanitizerCoverageLTO.so.cc b/instrumentation/SanitizerCoverageLTO.so.cc
index 65602109..b93b72bf 100644
--- a/instrumentation/SanitizerCoverageLTO.so.cc
+++ b/instrumentation/SanitizerCoverageLTO.so.cc
@@ -251,6 +251,7 @@ class ModuleSanitizerCoverageLTO
   uint32_t                         unhandled = 0;
   uint32_t                         select_cnt = 0;
   uint32_t                         instrument_ctx = 0;
+  uint32_t                         instrument_ctx_max_depth = 0;
   uint32_t                         extra_ctx_inst = 0;
   uint64_t                         map_addr = 0;
   const char                      *skip_nozero = NULL;
@@ -428,12 +429,31 @@ bool ModuleSanitizerCoverageLTO::instrumentModule(
   setvbuf(stdout, NULL, _IONBF, 0);
   if (getenv("AFL_DEBUG")) { debug = 1; }
   if (getenv("AFL_LLVM_DICT2FILE_NO_MAIN")) { autodictionary_no_main = 1; }
-  if (getenv("AFL_LLVM_CALLER") || getenv("AFL_LLVM_CTX")) {
+  if (getenv("AFL_LLVM_CALLER") || getenv("AFL_LLVM_CTX") ||
+      getenv("AFL_LLVM_LTO_CALLER") || getenv("AFL_LLVM_LTO_CTX")) {
 
     instrument_ctx = 1;
 
   }
 
+  if (getenv("AFL_LLVM_LTO_CALLER_DEPTH")) {
+
+    instrument_ctx_max_depth = atoi(getenv("AFL_LLVM_LTO_CALLER_DEPTH"));
+
+  } else if (getenv("AFL_LLVM_LTO_CTX_DEPTH")) {
+
+    instrument_ctx_max_depth = atoi(getenv("AFL_LLVM_LTO_CTX_DEPTH"));
+
+  } else if (getenv("AFL_LLVM_CALLER_DEPTH")) {
+
+    instrument_ctx_max_depth = atoi(getenv("AFL_LLVM_CALLER_DEPTH"));
+
+  } else if (getenv("AFL_LLVM_CTX_DEPTH")) {
+
+    instrument_ctx_max_depth = atoi(getenv("AFL_LLVM_CTX_DEPTH"));
+
+  }
+
   if ((isatty(2) && !getenv("AFL_QUIET")) || debug) {
 
     SAYF(cCYA "afl-llvm-lto" VERSION cRST
@@ -1406,11 +1426,12 @@ void ModuleSanitizerCoverageLTO::instrumentFunction(
     call_counter = countCallers(caller);
     Function *callee = caller;
 
-    if (call_counter == 1) {
+    if (call_counter == 1 && instrument_ctx_max_depth) {
 
       ++call_depth;
 
-      while (((caller = returnOnlyCaller(callee)) || 1 == 1) &&
+      while (instrument_ctx_max_depth >= call_depth &&
+             ((caller = returnOnlyCaller(callee)) || 1 == 1) &&
              (call_counter = countCallers(callee)) == 1) {
 
         if (debug && caller && callee)
diff --git a/src/afl-cc.c b/src/afl-cc.c
index 7d33b9f5..4d586ce8 100644
--- a/src/afl-cc.c
+++ b/src/afl-cc.c
@@ -2921,11 +2921,12 @@ static void maybe_usage(aflcc_state_t *aflcc, int argc, char **argv) {
             "  AFL_LLVM_DOCUMENT_IDS: write all edge IDs and the corresponding "
             "functions\n"
             "    into this file (LTO mode)\n"
+            "  AFL_LLVM_LTO_CALLER: activate CALLER/CTX instrumentation\n"
+            "  AFL_LLVM_LTO_CALLER_DEPTH: skip how many empty functions\n"
             "  AFL_LLVM_LTO_DONTWRITEID: don't write the highest ID used to a "
             "global var\n"
             "  AFL_LLVM_LTO_STARTID: from which ID to start counting from for "
-            "a "
-            "bb\n"
+            "a bb\n"
             "  AFL_REAL_LD: use this lld linker instead of the compiled in "
             "path\n"
             "  AFL_LLVM_LTO_SKIPINIT: don't inject initialization code "