1 files changed, 54 insertions, 57 deletions

diff --git a/instrumentation/afl-llvm-dict2file.so.cc b/instrumentation/afl-llvm-dict2file.so.cc
index 37cd8ad0..79cdf491 100644
--- a/instrumentation/afl-llvm-dict2file.so.cc
+++ b/instrumentation/afl-llvm-dict2file.so.cc
@@ -45,7 +45,6 @@
   #include "llvm/IR/PassManager.h"
 #else
   #include "llvm/IR/LegacyPassManager.h"
-  #include "llvm/Transforms/IPO/PassManagerBuilder.h"
 #endif
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Module.h"
@@ -67,47 +66,6 @@
   #define O_DSYNC O_SYNC
 #endif
 
-std::ofstream of;
-
-void dict2file(u8 *mem, u32 len) {
-
-  u32  i, j, binary = 0;
-  char line[MAX_AUTO_EXTRA * 8], tmp[8];
-
-  strcpy(line, "\"");
-  j = 1;
-  for (i = 0; i < len; i++) {
-
-    if (isprint(mem[i]) && mem[i] != '\\' && mem[i] != '"') {
-
-      line[j++] = mem[i];
-
-    } else {
-
-      if (i + 1 != len || mem[i] != 0 || binary || len == 4 || len == 8) {
-
-        line[j] = 0;
-        sprintf(tmp, "\\x%02x", (u8)mem[i]);
-        strcat(line, tmp);
-        j = strlen(line);
-
-      }
-
-      binary = 1;
-
-    }
-
-  }
-
-  line[j] = 0;
-  strcat(line, "\"\n");
-  of << line;
-  of.flush();
-
-  if (!be_quiet) fprintf(stderr, "Found dictionary token: %s", line);
-
-}
-
 using namespace llvm;
 
 namespace {
@@ -115,14 +73,22 @@ namespace {
 #if LLVM_VERSION_MAJOR >= 11                        /* use new pass manager */
 class AFLdict2filePass : public PassInfoMixin<AFLdict2filePass> {
 
+  std::ofstream of;
+  void          dict2file(u8 *, u32);
+
  public:
   AFLdict2filePass() {
 
 #else
+
 class AFLdict2filePass : public ModulePass {
 
+  std::ofstream of;
+  void          dict2file(u8 *, u32);
+
  public:
-  bool runOnModule(Module &M) override;
+  static char ID;
+
   AFLdict2filePass() : ModulePass(ID) {
 
 #endif
@@ -133,6 +99,8 @@ class AFLdict2filePass : public ModulePass {
 
 #if LLVM_VERSION_MAJOR >= 11                        /* use new pass manager */
   PreservedAnalyses run(Module &M, ModuleAnalysisManager &MAM);
+#else
+  bool runOnModule(Module &M) override;
 #endif
 
 };
@@ -143,7 +111,7 @@ class AFLdict2filePass : public ModulePass {
 extern "C" ::llvm::PassPluginLibraryInfo LLVM_ATTRIBUTE_WEAK
 llvmGetPassPluginInfo() {
 
-  return {LLVM_PLUGIN_API_VERSION, "dict2file", "v0.1",
+  return {LLVM_PLUGIN_API_VERSION, "AFLdict2filePass", "v0.1",
           /* lambda to insert our pass into the pass pipeline. */
           [](PassBuilder &PB) {
 
@@ -162,16 +130,52 @@ llvmGetPassPluginInfo() {
 }
 
 #else
-
 char AFLdict2filePass::ID = 0;
-
 #endif
 
+void AFLdict2filePass::dict2file(u8 *mem, u32 len) {
+
+  u32  i, j, binary = 0;
+  char line[MAX_AUTO_EXTRA * 8], tmp[8];
+
+  strcpy(line, "\"");
+  j = 1;
+  for (i = 0; i < len; i++) {
+
+    if (isprint(mem[i]) && mem[i] != '\\' && mem[i] != '"') {
+
+      line[j++] = mem[i];
+
+    } else {
+
+      if (i + 1 != len || mem[i] != 0 || binary || len == 4 || len == 8) {
+
+        line[j] = 0;
+        sprintf(tmp, "\\x%02x", (u8)mem[i]);
+        strcat(line, tmp);
+        j = strlen(line);
+
+      }
+
+      binary = 1;
+
+    }
+
+  }
+
+  line[j] = 0;
+  strcat(line, "\"\n");
+  of << line;
+  of.flush();
+
+  if (!be_quiet) fprintf(stderr, "Found dictionary token: %s", line);
+
+}
+
 #if LLVM_VERSION_MAJOR >= 11                        /* use new pass manager */
 PreservedAnalyses AFLdict2filePass::run(Module &M, ModuleAnalysisManager &MAM) {
 
 #else
-
 bool AFLdict2filePass::runOnModule(Module &M) {
 
 #endif
@@ -640,6 +644,7 @@ bool AFLdict2filePass::runOnModule(Module &M) {
 
               if (optLen < 2) { continue; }
               if (literalLength + 1 == optLen) {  // add null byte
+
                 thestring.append("\0", 1);
 
               }
@@ -662,17 +667,11 @@ bool AFLdict2filePass::runOnModule(Module &M) {
                         }
 
             */
-
-            if (!isStdString) {
+            if (!isStdString && thestring.find('\0', 0) != std::string::npos) {
 
               // ensure we do not have garbage
               size_t offset = thestring.find('\0', 0);
-              if (offset && offset < optLen && offset + 1 < optLen) {
-
-                optLen = offset + 1;
-
-              }
-
+              if (offset + 1 < optLen) optLen = offset + 1;
               thestring = thestring.substr(0, optLen);
 
             }
@@ -716,8 +715,6 @@ bool AFLdict2filePass::runOnModule(Module &M) {
 
 #if LLVM_VERSION_MAJOR >= 11                        /* use new pass manager */
   auto PA = PreservedAnalyses::all();
-#endif
-#if LLVM_VERSION_MAJOR >= 11                        /* use new pass manager */
   return PA;
 #else
   return true;