Autotokens New PM (#605)

* autotokens newpm

* typo

* fmt

* clp

* fix

* fix

* include &fmt

* include

* fmt

* llvm14 & clippy fix

* fix

libafl_cc/build.rs

@@ -206,6 +206,7 @@ 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())))

libafl_cc/src/autotokens-pass.cc

@@ -35,9 +35,15 @@
 #include "llvm/ADT/Statistic.h"
 #include "llvm/IR/IRBuilder.h"
 
-#include "llvm/IR/LegacyPassManager.h"
-#include "llvm/Transforms/IPO/PassManagerBuilder.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/Transforms/IPO/PassManagerBuilder.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/DebugInfo.h"
@@ -134,14 +140,24 @@ bool isIgnoreFunction(const llvm::Function *F) {
   return false;
 }
 
+#if USE_NEW_PM
+class AutoTokensPass : public PassInfoMixin<AutoTokensPass> {
+ public:
+  AutoTokensPass() {
+#else
 class AutoTokensPass : public ModulePass {
  public:
   static char ID;
 
   AutoTokensPass() : ModulePass(ID) {
+#endif
   }
 
+#if USE_NEW_PM
+  PreservedAnalyses run(Module &M, ModuleAnalysisManager &MAM);
+#else
   bool runOnModule(Module &M) override;
+#endif
 
  protected:
  private:
@@ -150,7 +166,25 @@ class AutoTokensPass : public ModulePass {
 
 }  // namespace
 
+#if USE_NEW_PM
+extern "C" ::llvm::PassPluginLibraryInfo LLVM_ATTRIBUTE_WEAK
+llvmGetPassPluginInfo() {
+  return {LLVM_PLUGIN_API_VERSION, "AutoTokensPass", "v0.1",
+          /* lambda to insert our pass into the pass pipeline. */
+          [](PassBuilder &PB) {
+
+  #if LLVM_VERSION_MAJOR <= 13
+            using OptimizationLevel = typename PassBuilder::OptimizationLevel;
+  #endif
+            PB.registerOptimizerLastEPCallback(
+                [](ModulePassManager &MPM, OptimizationLevel OL) {
+                  MPM.addPass(AutoTokensPass());
+                });
+          }};
+}
+#else
 char AutoTokensPass::ID = 0;
+#endif
 
 void dict2file(int fd, uint8_t *mem, uint32_t len) {
   uint32_t i, j, binary = 0;
@@ -177,11 +211,17 @@ void dict2file(int fd, uint8_t *mem, uint32_t len) {
   line[j] = 0;
   strcat(line, "\"\n");
   if (write(fd, line, strlen(line)) <= 0)
-    FATAL("Could not write to dictionary file");
+    FATAL("Could not write to the dictionary file");
   fsync(fd);
 }
 
+#if USE_NEW_PM
+PreservedAnalyses AutoTokensPass::run(Module &M, ModuleAnalysisManager &MAM) {
+#else
 bool AutoTokensPass::runOnModule(Module &M) {
+
+#endif
+
   DenseMap<Value *, std::string *> valueMap;
   char                            *ptr;
   int                              fd, found = 0;
@@ -426,7 +466,7 @@ bool AutoTokensPass::runOnModule(Module &M) {
           // we handle the 2nd parameter first because of llvm memcpy
           if (!HasStr2) {
             auto *Ptr = dyn_cast<ConstantExpr>(Str2P);
-            if (Ptr && Ptr->isGEPWithNoNotionalOverIndexing()) {
+            if (Ptr && Ptr->getOpcode() == Instruction::GetElementPtr) {
               if (auto *Var = dyn_cast<GlobalVariable>(Ptr->getOperand(0))) {
                 if (Var->hasInitializer()) {
                   if (auto *Array =
@@ -475,7 +515,7 @@ bool AutoTokensPass::runOnModule(Module &M) {
           if (!HasStr1) {
             auto Ptr = dyn_cast<ConstantExpr>(Str1P);
 
-            if (Ptr && Ptr->isGEPWithNoNotionalOverIndexing()) {
+            if (Ptr && Ptr->getOpcode() == Instruction::GetElementPtr) {
               if (auto *Var = dyn_cast<GlobalVariable>(Ptr->getOperand(0))) {
                 if (Var->hasInitializer()) {
                   if (auto *Array =
@@ -568,7 +608,12 @@ bool AutoTokensPass::runOnModule(Module &M) {
 
   if (use_file) {
     close(fd);
+#if USE_NEW_PM
+    auto PA = PreservedAnalyses::all();
+    return PA;
+#else
     return true;
+#endif
   }
 
   LLVMContext &Ctx = M.getContext();
@@ -608,9 +653,17 @@ bool AutoTokensPass::runOnModule(Module &M) {
       "libafl_dictionary_" + M.getName());
   dict->setSection("libafl_token");
 
+#if USE_NEW_PM
+  auto PA = PreservedAnalyses::all();
+  return PA;
+#else
   return true;
+#endif
 }
 
+#if USE_NEW_PM
+
+#else
 static void registerAutoTokensPass(const PassManagerBuilder &,
                                    legacy::PassManagerBase &PM) {
   PM.add(new AutoTokensPass());
@@ -625,3 +678,4 @@ static RegisterStandardPasses RegisterAutoTokensPass(
 
 static RegisterStandardPasses RegisterAutoTokensPass0(
     PassManagerBuilder::EP_EnabledOnOptLevel0, registerAutoTokensPass);
+#endif

libafl_cc/src/clang.rs

@@ -72,6 +72,7 @@ pub struct ClangWrapper {
     bit_mode: u32,
     need_libafl_arg: bool,
     has_libafl_arg: bool,
+    use_new_pm: bool,
 
     parse_args_called: bool,
     base_args: Vec<String>,
@@ -281,14 +282,23 @@ impl CompilerWrapper for ClangWrapper {
             return Ok(args);
         }
 
-        if !self.passes.is_empty() {
-            args.push("-fno-experimental-new-pass-manager".into());
+        if self.use_new_pm {
+            args.push("-fexperimental-new-pass-manager".into());
+        } else {
+            args.push("-flegacy-pass-manager".into());
         }
         for pass in &self.passes {
-            args.push("-Xclang".into());
-            args.push("-load".into());
-            args.push("-Xclang".into());
-            args.push(pass.path().into_os_string().into_string().unwrap());
+            if self.use_new_pm {
+                args.push(format!(
+                    "-fpass-plugin={}",
+                    pass.path().into_os_string().into_string().unwrap()
+                ));
+            } else {
+                args.push("-Xclang".into());
+                args.push("-load".into());
+                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());
@@ -351,6 +361,7 @@ impl ClangWrapper {
             bit_mode: 0,
             need_libafl_arg: false,
             has_libafl_arg: false,
+            use_new_pm: false,
             parse_args_called: false,
             base_args: vec![],
             cc_args: vec![],
@@ -411,6 +422,12 @@ impl ClangWrapper {
         self.need_libafl_arg = value;
         self
     }
+
+    /// Set if use new llvm pass manager.
+    pub fn use_new_pm(&mut self, value: bool) -> &'_ mut Self {
+        self.use_new_pm = value;
+        self
+    }
 }
 
 #[cfg(test)]

libafl_targets/src/coverage.rs

@@ -47,7 +47,7 @@ pub fn autotokens() -> Result<Tokens, Error> {
     unsafe {
         if __token_start.is_null() || __token_stop.is_null() {
             Err(Error::illegal_state(
-                "AutoTokens section not found, likely the targe is not compiled with AutoTokens",
+                "AutoTokens section not found, likely the target is not compiled with AutoTokens",
             ))
         } else {
             // we can safely unwrap