Cmplog New Pass Manager & LLVM 14 Fixes (#626)

* wip * more * match aflpp * llvm14 * fix * more llvm14 * check llvm version in libafl_cc * safe access * more * fmt * no windows * no windows
2022-05-17 15:45:48 +09:00 · 2022-05-17 15:45:48 +09:00 · afb32fb351
commit afb32fb351
parent 2ead2c398e
5 changed files with 144 additions and 21 deletions
--- a/libafl_cc/build.rs
+++ b/libafl_cc/build.rs
@ -112,6 +112,21 @@ fn main() {

    let llvm_config = find_llvm_config();

+    // Get LLVM version.
+    let llvm_version = match llvm_config.split('-').collect::<Vec<&str>>().get(2) {
+        Some(ver) => ver.parse::<usize>().ok(),
+        None => None,
+    };
+
+    match llvm_version {
+        Some(ver) => {
+            if ver >= 14 {
+                custom_flags.push("-DUSE_NEW_PM".to_string());
+            }
+        }
+        None => (),
+    }
+
    if let Ok(output) = Command::new(&llvm_config).args(&["--bindir"]).output() {
        let llvm_bindir = Path::new(
            str::from_utf8(&output.stdout)
@ -132,11 +147,15 @@ fn main() {

            /// The size of the accounting maps
            pub const ACCOUNTING_MAP_SIZE: usize = {};
+
+            /// The llvm version used to build llvm passes
+            pub const LIBAFL_CC_LLVM_VERSION: Option<usize> = {:?};
            ",
            llvm_bindir.join("clang"),
            llvm_bindir.join("clang++"),
            edges_map_size,
-            acc_map_size
+            acc_map_size,
+            llvm_version,
        )
        .expect("Could not write file");

@ -206,7 +225,6 @@ fn main() {
            .args(&cxxflags)
            .args(&custom_flags)
            .arg(src_dir.join("autotokens-pass.cc"))
-            //.arg("-DUSE_NEW_PM")
            .args(&ldflags)
            .args(&["-fPIC", "-shared", "-o"])
            .arg(out_dir.join(format!("autotokens-pass.{}", dll_extension())))
--- a/libafl_cc/src/afl-coverage-pass.cc
+++ b/libafl_cc/src/afl-coverage-pass.cc
@ -140,7 +140,9 @@ llvmGetPassPluginInfo() {
          /* 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(AFLCoverage());
@ -433,7 +435,11 @@ bool AFLCoverage::runOnModule(Module &M) {
      if (instrument_ctx && &BB == &F.getEntryBlock()) {
 #ifdef HAVE_VECTOR_INTRINSICS
        if (CtxK) {
-          PrevCaller = IRB.CreateLoad(AFLPrevCaller);
+          PrevCaller = IRB.CreateLoad(
+  #if LLVM_VERSION_MAJOR >= 14
+              PrevCallerTy,
+  #endif
+              AFLPrevCaller);
          PrevCaller->setMetadata(M.getMDKindID("nosanitize"),
                                  MDNode::get(C, None));
          PrevCtx =
@ -445,7 +451,12 @@ bool AFLCoverage::runOnModule(Module &M) {

          // load the context ID of the previous function and write to to a
          // local variable on the stack
-          LoadInst *PrevCtxLoad = IRB.CreateLoad(AFLContext);
+          LoadInst *PrevCtxLoad = IRB.CreateLoad(
+#if LLVM_VERSION_MAJOR >= 14
+            IRB.getInt32Ty(),
+#endif
+            AFLContext
+          );
          PrevCtxLoad->setMetadata(M.getMDKindID("nosanitize"),
                                   MDNode::get(C, None));
          PrevCtx = PrevCtxLoad;
@ -573,7 +584,11 @@ bool AFLCoverage::runOnModule(Module &M) {

      /* Load prev_loc */

-      LoadInst *PrevLoc = IRB.CreateLoad(AFLPrevLoc);
+      LoadInst *PrevLoc = IRB.CreateLoad(
+#if LLVM_VERSION_MAJOR >= 14
+          PrevLocTy,
+#endif
+          AFLPrevLoc);
      PrevLoc->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None));
      Value *PrevLocTrans;

@ -597,20 +612,31 @@ bool AFLCoverage::runOnModule(Module &M) {

      /* Load SHM pointer */

-      LoadInst *MapPtr = IRB.CreateLoad(AFLMapPtr);
+      LoadInst *MapPtr = IRB.CreateLoad(
+#if LLVM_VERSION_MAJOR >= 14
+          PointerType::get(Int8Ty, 0),
+#endif
+          AFLMapPtr);
      MapPtr->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None));

      Value *MapPtrIdx;
 #ifdef HAVE_VECTOR_INTRINSICS
      if (Ngram)
        MapPtrIdx = IRB.CreateGEP(
+  #if LLVM_VERSION_MAJOR >= 14
+            Int8Ty,
+  #endif
            MapPtr,
            IRB.CreateZExt(
                IRB.CreateXor(PrevLocTrans, IRB.CreateZExt(CurLoc, Int32Ty)),
                Int32Ty));
      else
 #endif
-        MapPtrIdx = IRB.CreateGEP(MapPtr, IRB.CreateXor(PrevLocTrans, CurLoc));
+        MapPtrIdx = IRB.CreateGEP(
+#if LLVM_VERSION_MAJOR >= 14
+            Int8Ty,
+#endif
+            MapPtr, IRB.CreateXor(PrevLocTrans, CurLoc));

      /* Update bitmap */

@ -643,7 +669,11 @@ bool AFLCoverage::runOnModule(Module &M) {
        */

      } else {
-        LoadInst *Counter = IRB.CreateLoad(MapPtrIdx);
+        LoadInst *Counter = IRB.CreateLoad(
+#if LLVM_VERSION_MAJOR >= 14
+            IRB.getInt8Ty(),
+#endif
+            MapPtrIdx);
        Counter->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None));

        Value *Incr = IRB.CreateAdd(Counter, One);
--- a/libafl_cc/src/clang.rs
+++ b/libafl_cc/src/clang.rs
@ -349,6 +349,14 @@ impl ClangWrapper {
    /// Create a new Clang Wrapper
    #[must_use]
    pub fn new() -> Self {
+        #[cfg(unix)]
+        let use_new_pm = match LIBAFL_CC_LLVM_VERSION {
+            Some(ver) => ver >= 14,
+            None => false,
+        };
+        #[cfg(not(unix))]
+        let use_new_pm = false;
+
        Self {
            optimize: true,
            wrapped_cc: CLANG_PATH.into(),
@ -361,7 +369,7 @@ impl ClangWrapper {
            bit_mode: 0,
            need_libafl_arg: false,
            has_libafl_arg: false,
-            use_new_pm: false,
+            use_new_pm,
            parse_args_called: false,
            base_args: vec![],
            cc_args: vec![],
--- a/libafl_cc/src/cmplog-routines-pass.cc
+++ b/libafl_cc/src/cmplog-routines-pass.cc
@ -25,9 +25,16 @@
 #include <sys/time.h>
 #include "llvm/Config/llvm-config.h"

+#if USE_NEW_PM
+  #include "llvm/Passes/PassPlugin.h"
+  #include "llvm/Passes/PassBuilder.h"
+  #include "llvm/IR/PassManager.h"
+#else
+  #include "llvm/IR/LegacyPassManager.h"
+#endif
+
 #include "llvm/ADT/Statistic.h"
 #include "llvm/IR/IRBuilder.h"
-#include "llvm/IR/LegacyPassManager.h"
 #include "llvm/IR/Module.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
@ -111,22 +118,33 @@ bool isIgnoreFunction(const llvm::Function *F) {
  return false;
 }

+#if USE_NEW_PM
+class CmpLogRoutines : public PassInfoMixin<CmpLogRoutines> {
+ public:
+  CmpLogRoutines() {
+#else
+
 class CmpLogRoutines : public ModulePass {
 public:
  static char ID;
  CmpLogRoutines() : ModulePass(ID) {
+#endif
  }

-  bool runOnModule(Module &M) override;
-
-#if LLVM_VERSION_MAJOR < 4
-  const char *getPassName() const override {
+#if USE_NEW_PM
+  PreservedAnalyses run(Module &M, ModuleAnalysisManager &MAM);
 #else
+  bool        runOnModule(Module &M) override;
+
+  #if LLVM_VERSION_MAJOR < 4
+  const char *getPassName() const override {
+  #else
  StringRef getPassName() const override {

-#endif
+  #endif
    return "cmplog routines";
  }
+#endif

 private:
  bool hookRtns(Module &M);
@ -134,7 +152,23 @@ class CmpLogRoutines : public ModulePass {

 }  // namespace

+#if USE_NEW_PM
+extern "C" ::llvm::PassPluginLibraryInfo LLVM_ATTRIBUTE_WEAK
+llvmGetPassPluginInfo() {
+  return {LLVM_PLUGIN_API_VERSION, "CmpLogRoutines", "v0.1",
+          [](PassBuilder &PB) {
+  #if LLVM_VERSION_MAJOR <= 13
+            using OptimizationLevel = typename PassBuilder::OptimizationLevel;
+  #endif
+            PB.registerOptimizerLastEPCallback(
+                [](ModulePassManager &MPM, OptimizationLevel OL) {
+                  MPM.addPass(CmpLogRoutines());
+                });
+          }};
+}
+#else
 char CmpLogRoutines::ID = 0;
+#endif

 bool CmpLogRoutines::hookRtns(Module &M) {
  std::vector<CallInst *> calls, llvmStdStd, llvmStdC, gccStdStd, gccStdC;
@ -407,13 +441,27 @@ bool CmpLogRoutines::hookRtns(Module &M) {
  return true;
 }

+#if USE_NEW_PM
+PreservedAnalyses CmpLogRoutines::run(Module &M, ModuleAnalysisManager &MAM) {
+#else
 bool CmpLogRoutines::runOnModule(Module &M) {
+#endif
  hookRtns(M);
+
+#if USE_NEW_PM
+  auto PA = PreservedAnalyses::all();
+#endif
  verifyModule(M);

+#if USE_NEW_PM
+  return PA;
+#else
  return true;
+#endif
 }

+#if USE_NEW_PM
+#else
 static void registerCmpLogRoutinesPass(const PassManagerBuilder &,
                                       legacy::PassManagerBase &PM) {
  auto p = new CmpLogRoutines();
@ -426,8 +474,10 @@ static RegisterStandardPasses RegisterCmpLogRoutinesPass(
 static RegisterStandardPasses RegisterCmpLogRoutinesPass0(
    PassManagerBuilder::EP_EnabledOnOptLevel0, registerCmpLogRoutinesPass);

-#if LLVM_VERSION_MAJOR >= 11
+  #if LLVM_VERSION_MAJOR >= 11
 static RegisterStandardPasses RegisterCmpLogRoutinesPassLTO(
    PassManagerBuilder::EP_FullLinkTimeOptimizationLast,
    registerCmpLogRoutinesPass);
+  #endif
+
 #endif
--- a/libafl_cc/src/coverage-accounting-pass.cc
+++ b/libafl_cc/src/coverage-accounting-pass.cc
@ -140,7 +140,9 @@ llvmGetPassPluginInfo() {
          /* 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(AFLCoverage());
@ -254,20 +256,35 @@ bool AFLCoverage::runOnModule(Module &M) {

      /* Load prev_loc */

-      LoadInst *PrevLoc = IRB.CreateLoad(AFLPrevLoc);
+      LoadInst *PrevLoc = IRB.CreateLoad(
+#if LLVM_VERSION_MAJOR >= 14
+          Int32Ty,
+#endif
+          AFLPrevLoc);
      PrevLoc->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None));

      /* Load SHM pointer */

-      LoadInst *MemReadPtr = IRB.CreateLoad(AFLMemOpPtr);
+      LoadInst *MemReadPtr = IRB.CreateLoad(
+#if LLVM_VERSION_MAJOR >= 14
+          PointerType::get(Int32Ty, 0),
+#endif
+          AFLMemOpPtr);
      MemReadPtr->setMetadata(M.getMDKindID("nosanitize"),
                              MDNode::get(C, None));
-      Value *MemReadPtrIdx =
-          IRB.CreateGEP(MemReadPtr, IRB.CreateXor(PrevLoc, CurLoc));
+      Value *MemReadPtrIdx = IRB.CreateGEP(
+#if LLVM_VERSION_MAJOR >= 14
+          Int32Ty,
+#endif
+          MemReadPtr, IRB.CreateXor(PrevLoc, CurLoc));

      /* Update bitmap */

-      LoadInst *MemReadCount = IRB.CreateLoad(MemReadPtrIdx);
+      LoadInst *MemReadCount = IRB.CreateLoad(
+#if LLVM_VERSION_MAJOR >= 14
+          Int32Ty,
+#endif
+          MemReadPtrIdx);
      MemReadCount->setMetadata(M.getMDKindID("nosanitize"),
                                MDNode::get(C, None));
      Value *MemReadIncr =