Add function call level granularity for coverage accounting (#552)

* Add func call level granularity for coverage accounting

* code linting

fuzzers/libfuzzer_libpng_accounting/src/bin/libafl_cc.rs

@@ -1,6 +1,8 @@
 use libafl_cc::{ClangWrapper, CompilerWrapper, LLVMPasses};
 use std::env;
 
+const GRANULARITY: &str = "FUNC";
+
 pub fn main() {
     let args: Vec<String> = env::args().collect();
     if args.len() > 1 {
@@ -25,6 +27,7 @@ pub fn main() {
             .link_staticlib(&dir, "libfuzzer_libpng")
             .add_arg("-fsanitize-coverage=trace-pc-guard")
             .add_pass(LLVMPasses::CoverageAccounting)
+            .add_passes_arg(format!("-granularity={}", GRANULARITY))
             .run()
             .expect("Failed to run the wrapped compiler")
         {

fuzzers/libfuzzer_libpng_ctx/src/bin/libafl_cc.rs

@@ -23,8 +23,7 @@ pub fn main() {
             .parse_args(&args)
             .expect("Failed to parse the command line")
             .add_pass(LLVMPasses::AFLCoverage)
-            .add_arg("-mllvm")
-            .add_arg("-ctx") // Context sensitive coverage
+            .add_passes_arg("-ctx") // Context sensitive coverage
             .link_staticlib(&dir, "libfuzzer_libpng")
             .run()
             .expect("Failed to run the wrapped compiler")

libafl_cc/src/clang.rs

@@ -77,6 +77,7 @@ pub struct ClangWrapper {
     cc_args: Vec<String>,
     link_args: Vec<String>,
     passes: Vec<LLVMPasses>,
+    passes_args: Vec<String>,
 }
 
 #[allow(clippy::match_same_arms)] // for the linking = false wip for "shared"
@@ -264,6 +265,10 @@ impl CompilerWrapper for ClangWrapper {
             args.push("-Xclang".into());
             args.push(pass.path().into_os_string().into_string().unwrap());
         }
+        for passes_arg in &self.passes_args {
+            args.push("-mllvm".into());
+            args.push(passes_arg.into());
+        }
         if self.linking {
             if self.x_set {
                 args.push("-x".into());
@@ -325,6 +330,7 @@ impl ClangWrapper {
             cc_args: vec![],
             link_args: vec![],
             passes: vec![],
+            passes_args: vec![],
             is_silent: false,
         }
     }
@@ -359,6 +365,15 @@ impl ClangWrapper {
         self
     }
 
+    /// Add LLVM pass arguments
+    pub fn add_passes_arg<S>(&mut self, arg: S) -> &'_ mut Self
+    where
+        S: AsRef<str>,
+    {
+        self.passes_args.push(arg.as_ref().to_string());
+        self
+    }
+
     /// Set if linking
     pub fn linking(&mut self, value: bool) -> &'_ mut Self {
         self.linking = value;

libafl_cc/src/coverage-accounting-pass.cc

@@ -54,14 +54,61 @@ typedef uint32_t prev_loc_t;
 
 #define MAP_SIZE LIBAFL_ACCOUNTING_MAP_SIZE
 
+
+#define SECURITY_SENSITIVE_FUNCS(CF) static CF securitySensitiveFunctions[] = {\
+  CF("memcpy"), \
+  CF("strlen"), \
+  CF("ReadImage"), \
+  CF("memmove"), \
+  CF("free"), \
+  CF("memset"), \
+  CF("delete"), \
+  CF("memcmp"), \
+  CF("getString"), \
+  CF("vsprintf"), \
+  CF("GET_COLOR"), \
+  CF("read"), \
+  CF("load_bmp"), \
+  CF("huffcode"), \
+  CF("strcmp"), \
+  CF("new"), \
+  CF("getName"), \
+  CF("strncat"), \
+  CF("png_load"), \
+};
+
+
 using namespace llvm;
 
+enum AccountingGranularity {
+    BB_GRAN,
+    FUNC_GRAN,
+    // LOOP,
+    UKNOWN_GRAN
+};
+
 static cl::opt<bool> Debug("debug", cl::desc("Debug prints"), cl::init(false), cl::NotHidden);
+static cl::opt<std::string> GranularityStr("granularity", cl::desc("Granularity of accounting (BB, FUNC)"), cl::init(std::string("BB")), cl::NotHidden);
 static cl::opt<uint32_t> InstRatio("inst_ratio", cl::desc("Instrumentation ratio in percentage"), cl::init(100), cl::NotHidden);
 static cl::opt<bool> ThreadSafe("thread_safe", cl::desc("Use the thread safe instrumentation"), cl::init(false), cl::NotHidden);
 
 namespace {
 
+SECURITY_SENSITIVE_FUNCS(StringRef)
+
+bool isSecuritySensitiveFunction(Function* F) {
+    if (!F) return 0;
+    auto func_name = F->getName();
+    for (auto name : securitySensitiveFunctions) {
+        if (func_name.contains(name)) {
+            if (Debug)
+                fprintf(stderr, "Counted %s as security sensitive", func_name.str().c_str());
+            return 1;
+        }
+    }
+    return 0;
+}
+
 #ifdef USE_NEW_PM
 class AFLCoverage : public PassInfoMixin<AFLCoverage> {
  public:
@@ -72,7 +119,10 @@ class AFLCoverage : public ModulePass {
   static char ID;
   AFLCoverage() : ModulePass(ID) {
 #endif
-
+    granularity = StringSwitch<AccountingGranularity>(GranularityStr)
+        .Case("BB", BB_GRAN)
+        .Case("FUNC", FUNC_GRAN)
+        .Default(UKNOWN_GRAN);
     // initInstrumentList();
 
   }
@@ -86,6 +136,7 @@ class AFLCoverage : public ModulePass {
  protected:
   uint32_t    map_size = MAP_SIZE;
   uint32_t    function_minimum_size = 1;
+  AccountingGranularity granularity;
 
 };
 
@@ -196,12 +247,20 @@ bool AFLCoverage::runOnModule(Module &M) {
 
       // Start with 1 to implicitly track edge coverage too
       uint32_t MemCnt = 1;
-
       for (auto &I : BB) {
-          if (I.mayReadFromMemory() || I.mayWriteToMemory())
+        switch (granularity) {
+          case BB_GRAN:
+            if (I.mayReadFromMemory() || I.mayWriteToMemory())
               ++MemCnt;
+            break;
+          case FUNC_GRAN:
+            if (auto *C = dyn_cast<CallInst>(&I)) {
+              auto F = C->getCalledFunction();
+              MemCnt += isSecuritySensitiveFunction(F);
+            }
+            break;
+        }
       }
-      
       /* Make up cur_loc */
 
       cur_loc = RandBelow(map_size);