Parsing tables state transition coverage

libafl_cc/build.rs

@@ -397,6 +397,7 @@ pub const LIBAFL_CC_LLVM_VERSION: Option<usize> = None;
         "afl-coverage-pass.cc",
         "autotokens-pass.cc",
         "coverage-accounting-pass.cc",
+        "tables-pass.cc",
     ] {
         build_pass(
             bindir_path,

libafl_cc/src/clang.rs

@@ -39,6 +39,8 @@ pub enum LLVMPasses {
     AutoTokens,
     /// The Coverage Accouting (BB metric) pass
     CoverageAccounting,
+    /// Coverage from parsing tables
+    Tables,
     /// The dump cfg pass
     DumpCfg,
 }
@@ -57,6 +59,9 @@ impl LLVMPasses {
             }
             LLVMPasses::CoverageAccounting => PathBuf::from(env!("OUT_DIR"))
                 .join(format!("coverage-accounting-pass.{}", dll_extension())),
+            LLVMPasses::Tables => {
+                PathBuf::from(env!("OUT_DIR")).join(format!("tabes-pass.{}", dll_extension()))
+            }
             LLVMPasses::DumpCfg => {
                 PathBuf::from(env!("OUT_DIR")).join(format!("dump-cfg-pass.{}", dll_extension()))
             }

libafl_cc/src/no-link-rt.c

@@ -11,6 +11,11 @@ void __libafl_targets_cmplog_instructions(uintptr_t k, uint8_t shape,
   (void)arg2;
 }
 
+void __libafl_tables_transition(uint32_t arg1, uint32_t arg2) {
+  (void)arg1;
+  (void)arg2;
+}
+
 void __cmplog_rtn_hook(uint8_t *ptr1, uint8_t *ptr2) {
   (void)ptr1;
   (void)ptr2;

libafl_cc/src/tables-pass.cc

@@ -0,0 +1,229 @@
+/*
+   LibAFL - Coverage accounting LLVM pass
+   --------------------------------------------------
+
+   Written by Andrea Fioraldi <[email protected]>
+
+   Copyright 2023 AFLplusplus Project. All rights reserved.
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at:
+
+     http://www.apache.org/licenses/LICENSE-2.0
+
+*/
+
+#include "common-llvm.h"
+
+#include <time.h>
+
+#include <unordered_set>
+#include <list>
+#include <string>
+#include <fstream>
+
+#include "llvm/Support/CommandLine.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/BasicBlock.h"
+#include "llvm/IR/Module.h"
+#include "llvm/Support/Debug.h"
+#include "llvm/Support/MathExtras.h"
+
+// Without this, Can't build with llvm-14 & old PM
+#if LLVM_VERSION_MAJOR >= 14 && !defined(USE_NEW_PM)
+  #include "llvm/Pass.h"
+#endif
+
+#if LLVM_VERSION_MAJOR > 3 || \
+    (LLVM_VERSION_MAJOR == 3 && LLVM_VERSION_MINOR > 4)
+  #include "llvm/IR/DebugInfo.h"
+  #include "llvm/IR/CFG.h"
+#else
+  #include "llvm/DebugInfo.h"
+  #include "llvm/Support/CFG.h"
+#endif
+
+using namespace llvm;
+
+static cl::opt<bool> Debug("debug", cl::desc("Debug prints"), cl::init(false),
+                           cl::NotHidden);
+
+namespace {
+
+Value *recurseCast(Value *V) {
+  CastInst *CI;
+  if ((CI = dyn_cast<CastInst>(V))) { return recurseCast(CI->getOperand(0)); }
+  return V;
+}
+
+#ifdef USE_NEW_PM
+class TamingParsingTables : public PassInfoMixin<TamingParsingTables> {
+ public:
+  TamingParsingTables() {
+#else
+class TamingParsingTables : public ModulePass {
+ public:
+  static char ID;
+  TamingParsingTables() : ModulePass(ID) {
+#endif
+    // initInstrumentList();
+  }
+
+#ifdef USE_NEW_PM
+  PreservedAnalyses run(Module &M, ModuleAnalysisManager &MAM);
+#else
+  bool runOnModule(Module &M) override;
+#endif
+
+ protected:
+  uint32_t function_minimum_size = 1;
+};
+
+}  // namespace
+
+#ifdef USE_NEW_PM
+extern "C" ::llvm::PassPluginLibraryInfo LLVM_ATTRIBUTE_WEAK
+llvmGetPassPluginInfo() {
+  return {LLVM_PLUGIN_API_VERSION, "TamingParsingTables", "v0.1",
+          /* lambda to insert our pass into the pass pipeline. */
+          [](PassBuilder &PB) {
+  #if 1
+    #if LLVM_VERSION_MAJOR <= 13
+            using OptimizationLevel = typename PassBuilder::OptimizationLevel;
+    #endif
+            PB.registerOptimizerLastEPCallback(
+                [](ModulePassManager &MPM, OptimizationLevel OL) {
+                  MPM.addPass(TamingParsingTables());
+                });
+  /* TODO LTO registration */
+  #else
+            using PipelineElement = typename PassBuilder::PipelineElement;
+            PB.registerPipelineParsingCallback([](StringRef          Name,
+                                                  ModulePassManager &MPM,
+                                                  ArrayRef<PipelineElement>) {
+              if (Name == "TamingParsingTables") {
+                MPM.addPass(TamingParsingTables());
+                return true;
+              } else {
+                return false;
+              }
+            });
+  #endif
+          }};
+}
+#else
+
+char TamingParsingTables::ID = 0;
+#endif
+
+#ifdef USE_NEW_PM
+PreservedAnalyses TamingParsingTables::run(Module &M, ModuleAnalysisManager &MAM) {
+#else
+bool TamingParsingTables::runOnModule(Module &M) {
+#endif
+
+  LLVMContext &C = M.getContext();
+
+  IntegerType *Int64Ty = IntegerType::getInt64Ty(C);
+  IntegerType *Int32Ty = IntegerType::getInt32Ty(C);
+  IntegerType *Int16Ty = IntegerType::getInt16Ty(C);
+  IntegerType *Int8Ty = IntegerType::getInt8Ty(C);
+  IntegerType *Int1Ty = IntegerType::getInt1Ty(C);
+  Type *VoidTy = Type::getVoidTy(C);
+
+  FunctionCallee LogFunc = M.getOrInsertFunction("__libafl_tables_transition", VoidTy, Int32Ty, Int32Ty);
+
+#ifdef USE_NEW_PM
+  auto PA = PreservedAnalyses::all();
+#endif
+
+  /* Instrument all the things! */
+
+  for (auto &F : M) {
+    int has_calls = 0;
+
+    // if (!isInInstrumentList(&F)) { continue; }
+
+    if (F.size() < function_minimum_size) { continue; }
+
+    std::unordered_set<LoadInst *> loads;
+    std::unordered_set<Value *>    geps;
+
+    for (auto &BB : F) {
+      for (auto &I : BB) {
+        GetElementPtrInst *GEP;
+        StoreInst         *ST;
+        LoadInst          *LI;
+        if ((LI = dyn_cast<LoadInst>(&I))) {
+          loads.insert(LI);
+        } else if ((GEP = dyn_cast<GetElementPtrInst>(&I))) {
+          if (!GEP->hasIndices() || GEP->hasAllConstantIndices()) continue;
+          // TODO handle multiple idxs
+          Value *IDX = *GEP->idx_begin();
+          IDX = recurseCast(IDX);
+
+          if ((LI = dyn_cast<LoadInst>(IDX)) && loads.contains(LI)) {
+            geps.insert(GEP);
+          }
+        } else if ((ST = dyn_cast<StoreInst>(&I))) {
+          Value    *PTR = ST->getPointerOperand();  // must be in a prev load
+          Value    *VAL = recurseCast(ST->getValueOperand());
+          LoadInst *GL = nullptr;
+          Value    *V = nullptr;
+
+          if ((GL = dyn_cast<LoadInst>(VAL))) {
+            V = GL->getPointerOperand();
+            if (V == nullptr || !geps.contains(V)) { continue; }
+          } else {
+            continue;
+          }
+
+          // the value comes from a load in which the ptr is obtained with a gep
+          GEP = (GetElementPtrInst *)V;
+          Value *IDX = *GEP->idx_begin();
+          IDX = recurseCast(IDX);
+          LI = dyn_cast<LoadInst>(IDX);
+
+          if (LI == nullptr) continue;
+          if (recurseCast(LI->getPointerOperand()) !=
+              recurseCast(ST->getPointerOperand()))
+            continue;
+
+          std::string location = std::string("UNKNOWN");
+          if (DILocation *Loc = GEP->getDebugLoc().get()) {
+            location = std::string(Loc->getFilename().data()) +
+                       std::string(":") + std::to_string(Loc->getLine());
+          }
+
+          errs() << "FOUND " << location << "\n\t" << *LI << "\n\t" << *GEP
+                 << "\n\t" << *ST << "\n\n";
+
+          IRBuilder<> IRB(ST);
+          Value *A1 = IRB.CreateIntCast(LI, Int32Ty, false);
+          Value *A2 = IRB.CreateIntCast(ST->getValueOperand(), Int32Ty, false);
+          IRB.CreateCall(LogFunc, {A1, A2});
+        }
+      }
+    }
+  }
+
+#ifdef USE_NEW_PM
+  return PA;
+#else
+  return true;
+#endif
+}
+
+#ifndef USE_NEW_PM
+static void registerAFLPass(const PassManagerBuilder &,
+                            legacy::PassManagerBase &PM) {
+  PM.add(new TamingParsingTables());
+}
+
+static RegisterStandardPasses RegisterAFLPass(
+    PassManagerBuilder::EP_OptimizerLast, registerAFLPass);
+
+static RegisterStandardPasses RegisterAFLPass0(
+    PassManagerBuilder::EP_EnabledOnOptLevel0, registerAFLPass);
+#endif

libafl_targets/build.rs

@@ -27,6 +27,9 @@ fn main() {
     let acc_map_size: usize = option_env!("LIBAFL_ACCOUNTING_MAP_SIZE")
         .map_or(Ok(65536), str::parse)
         .expect("Could not parse LIBAFL_ACCOUNTING_MAP_SIZE");
+    let tables_map_size: usize = option_env!("LIBAFL_TABLES_MAP_SIZE")
+        .map_or(Ok(16384), str::parse)
+        .expect("Could not parse LIBAFL_TABLES_MAP_SIZE");
 
     write!(
         constants_file,
@@ -42,6 +45,8 @@ fn main() {
         pub const CMPLOG_MAP_H: usize = {cmplog_map_h};
         /// The size of the accounting maps
         pub const ACCOUNTING_MAP_SIZE: usize = {acc_map_size};
+        /// The size of the tables map
+        pub const TABLES_MAP_SIZE: usize = {tables_map_size};
 "
     )
     .expect("Could not write file");
@@ -51,6 +56,7 @@ fn main() {
     println!("cargo:rerun-if-env-changed=LIBAFL_CMPLOG_MAP_W");
     println!("cargo:rerun-if-env-changed=LIBAFL_CMPLOG_MAP_H");
     println!("cargo:rerun-if-env-changed=LIBAFL_ACCOUNTING_MAP_SIZE");
+    println!("cargo:rerun-if-env-changed=LIBAFL_TABLES_MAP_SIZE");
 
     //std::env::set_var("CC", "clang");
     //std::env::set_var("CXX", "clang++");

libafl_targets/src/lib.rs

@@ -98,6 +98,9 @@ pub use value_profile::*;
 pub mod cmplog;
 pub use cmplog::*;
 
+pub mod tables;
+pub use tables::*;
+
 #[cfg(feature = "std")]
 pub mod drcov;
 

libafl_targets/src/tables.rs

@@ -0,0 +1,32 @@
+//! Tables state transition pass runtime for `LibAFL`.
+
+use crate::TABLES_MAP_SIZE;
+
+/// Map with tables transitions
+pub static mut TABLES_MAP: [u8; TABLES_MAP_SIZE] = [0; TABLES_MAP_SIZE];
+
+fn merge_u32(a: u32, b: u32) -> u64 {
+    (a as u64) << 32 + (b as u64)
+}
+
+/// From https://sair.synerise.com/efficient-integer-pairs-hashing/
+fn splitmix64(target: u64, seed: u64) -> u64 {
+    let sp_step = 0x9E3779B97F4A7C15_u64;
+    let mut out = target;
+    out = out + seed * sp_step;
+    out ^= out >> 30;
+    out *= 0xBF58476D1CE4E5B9_u64;
+    out ^= out >> 27;
+    out *= 0x94D049BB133111EB_u64;
+    out ^= out >> 31;
+    out
+}
+
+#[no_mangle]
+/// Log tables transitions and insert them in the map
+pub extern "C" fn __libafl_tables_transition(cur: u32, next: u32) {
+    let hash = splitmix64(merge_u32(cur, next), 52) as usize % TABLES_MAP_SIZE;
+    unsafe {
+        TABLES_MAP[hash] += 1;
+    }
+}