1 files changed, 25 insertions, 12 deletions

diff --git a/src/afl-cc.c b/src/afl-cc.c
index c3910e6d..9cf02059 100644
--- a/src/afl-cc.c
+++ b/src/afl-cc.c
@@ -73,7 +73,8 @@ enum {
   INSTRUMENT_GCC = 6,
   INSTRUMENT_CLANG = 7,
   INSTRUMENT_OPT_CTX = 8,
-  INSTRUMENT_OPT_NGRAM = 16
+  INSTRUMENT_OPT_NGRAM = 16,
+  INSTRUMENT_OPT_CALLER = 32,
 
 };
 
@@ -1273,7 +1274,8 @@ int main(int argc, char **argv, char **envp) {
   }
 
   if (getenv("AFL_LLVM_CTX")) instrument_opt_mode |= INSTRUMENT_OPT_CTX;
-
+  if (getenv("AFL_LLVM_CALLER")) instrument_opt_mode |= INSTRUMENT_OPT_CALLER;
+  
   if (getenv("AFL_LLVM_NGRAM_SIZE")) {
 
     instrument_opt_mode |= INSTRUMENT_OPT_NGRAM;
@@ -1387,6 +1389,13 @@ int main(int argc, char **argv, char **envp) {
         setenv("AFL_LLVM_CTX", "1", 1);
 
       }
+      
+      if (strncasecmp(ptr2, "caller", strlen("caller")) == 0) {
+
+        instrument_opt_mode |= INSTRUMENT_OPT_CALLER;
+        setenv("AFL_LLVM_CALLER", "1", 1);
+
+      }
 
       if (strncasecmp(ptr2, "ngram", strlen("ngram")) == 0) {
 
@@ -1420,6 +1429,11 @@ int main(int argc, char **argv, char **envp) {
     }
 
   }
+  
+  if ((instrument_opt_mode & INSTRUMENT_OPT_CTX) &&
+      (instrument_opt_mode & INSTRUMENT_OPT_CALLER)) {
+    FATAL("you cannot set CTX and CALLER together");
+  }
 
   if (instrument_opt_mode && instrument_mode == INSTRUMENT_DEFAULT &&
       (compiler_mode == LLVM || compiler_mode == UNSET)) {
@@ -1770,7 +1784,7 @@ int main(int argc, char **argv, char **envp) {
   }
 
   if (instrument_opt_mode && compiler_mode != LLVM)
-    FATAL("CTX and NGRAM can only be used in LLVM mode");
+    FATAL("CTX, CALLER and NGRAM can only be used in LLVM mode");
 
   if (!instrument_opt_mode) {
 
@@ -1780,15 +1794,14 @@ int main(int argc, char **argv, char **envp) {
 
   } else {
 
-    if (instrument_opt_mode == INSTRUMENT_OPT_CTX)
-
-      ptr = alloc_printf("%s + CTX", instrument_mode_string[instrument_mode]);
-    else if (instrument_opt_mode == INSTRUMENT_OPT_NGRAM)
-      ptr = alloc_printf("%s + NGRAM-%u",
-                         instrument_mode_string[instrument_mode], ngram_size);
-    else
-      ptr = alloc_printf("%s + CTX + NGRAM-%u",
-                         instrument_mode_string[instrument_mode], ngram_size);
+    char *ptr2 = alloc_printf(" + NGRAM-%u", ngram_size);
+    ptr = alloc_printf("%s%s%s%s", instrument_mode_string[instrument_mode],
+        (instrument_opt_mode & INSTRUMENT_OPT_CTX) ? " + CTX" : "",
+        (instrument_opt_mode & INSTRUMENT_OPT_CALLER) ? " + CALLER" : "",
+        (instrument_opt_mode & INSTRUMENT_OPT_NGRAM) ? ptr2 : ""
+    );
+    
+    ck_free(ptr2);
 
   }