Custom function call context (#55)

* improve custom function context

* minor fix doc

* minor fix

* typo

Author: spinpx

docs/environment_variables.md

@@ -3,8 +3,7 @@
 - `USE_FAST=1`: use fast mode to compile the program. It includes branch counting, getting the feedback of the fuzzing constraint (the output of its function).
 - `USE_TRACK=1`: use taint tracking and collect all constraints.
 - `USE_DFSAN=1`: use taint tracking.
-- `ANGORA_DISABLE_CONTEXT=1` : Disable function call based contexts in compiling.
-- `ANGORA_DIRECT_FN_CONTEXT=1` : Use only the last function call location as the context.
+- `ANGORA_CUSTOM_FN_CONTEXT=k` : Use only the last k ( 0 <= k <= 32) function call location as the context, e.g. `ANGORA_CUSTOM_FN_CONTEXT=8`. Angora disables context if k is 0.
 - `ANGORA_GEN_ID_RANDOM=1` : Generate ids for predicates randomly instead of the hash of their locations.
 - `ANGORA_OUTPUT_COND_LOC=1` : (Debug option) Output the location of each predicate during compiling.
 - `ANGORA_TAINT_CUSTOM_RULE=/path/to/object` : object contains those proxy function (how to propagate taints), e.g. `ANGORA_TAINT_CUSTOM_RULE=~/angora/bin/zlib-func.o` . You should add it as custom type in the file passed by `ANGORA_TAINT_RULE_LIST` first.

fuzzer/src/stats/chart.rs

@@ -126,7 +126,7 @@ impl ChartStats {
 impl fmt::Display for ChartStats {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         if self.density.0 > 10.0 {
-            warn!("Density is too large (> 10%). Please increase `MAP_SIZE_POW2` in `llvm_mode/config.h` and `MAP_LENGTH` in `common/src/config.rs`. Or disable function-call context by compiling with `ANGORA_DISABLE_CONTEXT=1` or `ANGORA_DIRECT_FN_CONTEXT=1` environment variable.");
+            warn!("Density is too large (> 10%). Please increase `MAP_SIZE_POW2` in `llvm_mode/config.h` and `MAP_LENGTH` in `common/src/config.rs`. Or disable function-call context(density > 50%) by compiling with `ANGORA_CUSTOM_FN_CONTEXT=k` (k is an integer and 0 <= k <= 32) environment variable. Angora disables context if k is 0.");
         }
 
         if self.search.multiple_inconsist() {

llvm_mode/angora-clang.c

@@ -431,12 +431,13 @@ int main(int argc, char **argv) {
   find_obj(argv[0]);
 
   edit_params(argc, argv);
-
+  /*
   for (int i = 0; i < cc_par_cnt; i++) {
     printf("%s ", cc_params[i]);
   }
   printf("\n");
-
+  */
+ 
   execvp(cc_params[0], (char **)cc_params);
 
   FATAL("Oops, failed to execute '%s' - check your PATH", cc_params[0]);

llvm_mode/angora-llvm-pass.so.cc

@@ -68,10 +68,9 @@ class AngoraLLVMPass : public ModulePass {
   DenseSet<u32> UniqCidSet;
 
   // Configurations
-  bool enable_ctx;
   bool gen_id_random;
   bool output_cond_loc;
-  bool direct_fn_ctx;
+  int num_fn_ctx;
 
   MDNode *ColdCallWeights;
 
@@ -333,17 +332,25 @@ void AngoraLLVMPass::initVariables(Module &M) {
                              ClExploitListFiles.end());
   ExploitList.set(SpecialCaseList::createOrDie(AllExploitListFiles));
 
-  enable_ctx = !getenv(DISABLE_CTX_VAR);
-  direct_fn_ctx = !!getenv(DIRECT_FN_CTX);
   gen_id_random = !!getenv(GEN_ID_RANDOM_VAR);
   output_cond_loc = !!getenv(OUTPUT_COND_LOC_VAR);
 
-  if (!enable_ctx) {
-    errs() << "disable context\n";
+  num_fn_ctx = -1;
+  char* custom_fn_ctx = getenv(CUSTOM_FN_CTX);
+  if (custom_fn_ctx) {
+    num_fn_ctx = atoi(custom_fn_ctx);
+    if (num_fn_ctx < 0 || num_fn_ctx > 32) {
+      errs() << "custom context should be: >= 0 && <=32 \n"; 
+      exit(1);
+    }
   }
 
-  if (direct_fn_ctx) {
-    errs() << "use direct function call context\n";
+  if (num_fn_ctx == 0) {
+    errs() << "disable context\n";
+  } 
+
+  if (num_fn_ctx > 0) {
+    errs() << "use custom function call context: " << num_fn_ctx << "\n";
   }
 
   if (gen_id_random) {
@@ -400,7 +407,7 @@ void AngoraLLVMPass::countEdge(Module &M, BasicBlock &BB) {
   IRB.CreateStore(IncRet, MapPtrIdx)->setMetadata(NoSanMetaId, NoneMetaNode);
 
   Value *NewPrevLoc = NULL;
-  if (enable_ctx) { // Call-based context
+  if (num_fn_ctx != 0) { // Call-based context
     // Load ctx
     LoadInst *CtxVal = IRB.CreateLoad(AngoraContext);
     setInsNonSan(CtxVal);
@@ -410,7 +417,7 @@ void AngoraLLVMPass::countEdge(Module &M, BasicBlock &BB) {
     // Udate PrevLoc
     NewPrevLoc =
         IRB.CreateXor(CtxValCasted, ConstantInt::get(Int32Ty, cur_loc >> 1));
-  } else {
+  } else { // disable context
     NewPrevLoc = ConstantInt::get(Int32Ty, cur_loc >> 1);
   }
   setValueNonSan(NewPrevLoc);
@@ -421,6 +428,9 @@ void AngoraLLVMPass::countEdge(Module &M, BasicBlock &BB) {
 
 
 void AngoraLLVMPass::addFnWrap(Function &F) {
+
+  if (num_fn_ctx == 0) return;
+
   // *** Pre Fn ***
   BasicBlock *BB = &F.getEntryBlock();
   Instruction *InsertPoint = &(*(BB->getFirstInsertionPt()));
@@ -439,8 +449,8 @@ void AngoraLLVMPass::addFnWrap(Function &F) {
   // by `xor` with the same value
   // Implementation of function context for AFL by heiko eissfeldt:
   // https://github.com/vanhauser-thc/afl-patches/blob/master/afl-fuzz-context_sensitive.diff
-  if (direct_fn_ctx) {
-    OriCtxVal = IRB.CreateLShr(OriCtxVal, 6);
+  if (num_fn_ctx > 0) {
+    OriCtxVal = IRB.CreateLShr(OriCtxVal, 32 / num_fn_ctx);
     setValueNonSan(OriCtxVal);
   }
 
@@ -471,11 +481,11 @@ void AngoraLLVMPass::processCall(Instruction *Inst) {
 
   //  if (ABIList.isIn(*Callee, "uninstrumented"))
   //  return;
-
-  IRBuilder<> IRB(Inst);
-  Constant* CallSite = ConstantInt::get(Int32Ty, getRandomContextId());
-  IRB.CreateStore(CallSite, AngoraCallSite)->setMetadata(NoSanMetaId, NoneMetaNode);
-
+  if (num_fn_ctx != 0) {
+    IRBuilder<> IRB(Inst);
+    Constant* CallSite = ConstantInt::get(Int32Ty, getRandomContextId());
+    IRB.CreateStore(CallSite, AngoraCallSite)->setMetadata(NoSanMetaId, NoneMetaNode);
+  }
 }
 
 void AngoraLLVMPass::visitCallInst(Instruction *Inst) {

llvm_mode/config.h

@@ -3,10 +3,9 @@
 
 #define MAP_SIZE_POW2 20
 #define MAP_SIZE (1 << MAP_SIZE_POW2)
-#define MAX_FUNCALL_LEVEL 25
 #define ENABLE_UNFOLD_BRANCH 1
 
-#define VERSION "1.10"
+#define VERSION "1.2.0"
 
 // Without taint tracking
 #define CLANG_FAST_TYPE 0
@@ -32,10 +31,7 @@
   } while (0)
 #endif
 
-#define SHM_ENV_VAR "ANGORA_BRANCHES_SHM_ID"
-#define ENABLE_FORKSRV "ANGORA_ENABLE_FORKSRV"
-#define DISABLE_CTX_VAR "ANGORA_DISABLE_CONTEXT"
-#define DIRECT_FN_CTX "ANGORA_DIRECT_FN_CONTEXT"
+#define CUSTOM_FN_CTX "ANGORA_CUSTOM_FN_CONTEXT"
 #define GEN_ID_RANDOM_VAR "ANGORA_GEN_ID_RANDOM"
 #define OUTPUT_COND_LOC_VAR "ANGORA_OUTPUT_COND_LOC"
 #define TAINT_CUSTOM_RULE_VAR "ANGORA_TAINT_CUSTOM_RULE"

tests/test.sh

@@ -41,6 +41,8 @@ target=${name}/${name}
 
 rm -f ${target}.fast ${target}.cmp ${target}.taint 
 
+# export ANGORA_CUSTOM_FN_CONTEXT=0
+
 bin_dir=../bin/
 ANGORA_USE_ASAN=1 USE_FAST=1 ${bin_dir}/angora-clang ${target}.c -lz -o ${target}.fast
 USE_TRACK=1 ${bin_dir}/angora-clang ${target}.c -lz -o ${target}.taint