diff options
| -rw-r--r-- | include/envs.h | 1 | ||||
| -rw-r--r-- | instrumentation/afl-llvm-pass.so.cc | 4 | ||||
| -rw-r--r-- | src/afl-cc.c | 29 |
3 files changed, 28 insertions, 6 deletions
diff --git a/include/envs.h b/include/envs.h index 37748a56..4d4d6b0e 100644 --- a/include/envs.h +++ b/include/envs.h @@ -80,6 +80,7 @@ static char *afl_environment_variables[] = { "AFL_LLVM_BLOCKLIST", "AFL_LLVM_CMPLOG", "AFL_LLVM_INSTRIM", + "AFL_LLVM_CALLER", "AFL_LLVM_CTX", "AFL_LLVM_CTX_K", "AFL_LLVM_DICT2FILE", diff --git a/instrumentation/afl-llvm-pass.so.cc b/instrumentation/afl-llvm-pass.so.cc index f4717345..0f773aba 100644 --- a/instrumentation/afl-llvm-pass.so.cc +++ b/instrumentation/afl-llvm-pass.so.cc @@ -217,8 +217,8 @@ bool AFLCoverage::runOnModule(Module &M) { VectorType *PrevCallerTy = NULL; if (ctx_k_str) - if (sscanf(ctx_k_str, "%u", &ctx_k) != 1 || ctx_k < 2 || ctx_k > CTX_MAX_K) - FATAL("Bad value of AFL_CTX_K (must be between 2 and CTX_MAX_K (%u))", + if (sscanf(ctx_k_str, "%u", &ctx_k) != 1 || ctx_k < 1 || ctx_k > CTX_MAX_K) + FATAL("Bad value of AFL_CTX_K (must be between 1 and CTX_MAX_K (%u))", CTX_MAX_K); if (ctx_k == 1) { diff --git a/src/afl-cc.c b/src/afl-cc.c index a517124f..b9e0c101 100644 --- a/src/afl-cc.c +++ b/src/afl-cc.c @@ -1298,11 +1298,21 @@ int main(int argc, char **argv, char **envp) { if (getenv("AFL_LLVM_CTX_K")) { - instrument_opt_mode |= INSTRUMENT_OPT_CTX_K; ctx_k = atoi(getenv("AFL_LLVM_CTX_K")); if (ctx_k < 1 || ctx_k > CTX_MAX_K) FATAL("K-CTX instrumentation mode must be between 1 and CTX_MAX_K (%u)", CTX_MAX_K); + if (ctx_k == 1) { + + setenv("AFL_LLVM_CALLER", "1", 1); + unsetenv("AFL_LLVM_CTX_K"); + instrument_opt_mode |= INSTRUMENT_OPT_CALLER; + + } else { + + instrument_opt_mode |= INSTRUMENT_OPT_CTX_K; + + } } @@ -1422,9 +1432,20 @@ int main(int argc, char **argv, char **envp) { "K-CTX instrumentation option must be between 1 and CTX_MAX_K " "(%u)", CTX_MAX_K); - instrument_opt_mode |= (INSTRUMENT_OPT_CTX_K); - u8 *ptr4 = alloc_printf("%u", ctx_k); - setenv("AFL_LLVM_CTX_K", ptr4, 1); + + if (ctx_k == 1) { + + instrument_opt_mode |= INSTRUMENT_OPT_CALLER; + setenv("AFL_LLVM_CALLER", "1", 1); + unsetenv("AFL_LLVM_CTX_K"); + + } else { + + instrument_opt_mode |= (INSTRUMENT_OPT_CTX_K); + u8 *ptr4 = alloc_printf("%u", ctx_k); + setenv("AFL_LLVM_CTX_K", ptr4, 1); + + } } |