Merge pull request #1107 from devnexen/llvm_pass_update

LLVM passes making slightly more C++

3 files changed, 24 insertions, 29 deletions

diff --git a/instrumentation/SanitizerCoverageLTO.so.cc b/instrumentation/SanitizerCoverageLTO.so.cc
index e06f8b93..eb0f06b2 100644
--- a/instrumentation/SanitizerCoverageLTO.so.cc
+++ b/instrumentation/SanitizerCoverageLTO.so.cc
@@ -250,7 +250,7 @@ class ModuleSanitizerCoverage {
   Module *                         Mo = NULL;
   GlobalVariable *                 AFLMapPtr = NULL;
   Value *                          MapPtrFixed = NULL;
-  FILE *                           documentFile = NULL;
+  std::ofstream                    dFile;
   size_t                           found = 0;
   // afl++ END
 
@@ -446,7 +446,8 @@ bool ModuleSanitizerCoverage::instrumentModule(
 
   if ((ptr = getenv("AFL_LLVM_DOCUMENT_IDS")) != NULL) {
 
-    if ((documentFile = fopen(ptr, "a")) == NULL)
+    dFile.open(ptr, std::ofstream::out | std::ofstream::app);
+    if (dFile.is_open())
       WARNF("Cannot access document file %s", ptr);
 
   }
@@ -1003,12 +1004,7 @@ bool ModuleSanitizerCoverage::instrumentModule(
     instrumentFunction(F, DTCallback, PDTCallback);
 
   // afl++ START
-  if (documentFile) {
-
-    fclose(documentFile);
-    documentFile = NULL;
-
-  }
+  if (dFile.is_open()) dFile.close();
 
   if (!getenv("AFL_LLVM_LTO_DONTWRITEID") || dictionary.size() || map_addr) {
 
@@ -1509,12 +1505,11 @@ void ModuleSanitizerCoverage::InjectCoverageAtBlock(Function &F, BasicBlock &BB,
     // afl++ START
     ++afl_global_id;
 
-    if (documentFile) {
+    if (dFile.is_open()) {
 
       unsigned long long int moduleID =
           (((unsigned long long int)(rand() & 0xffffffff)) << 32) | getpid();
-      fprintf(documentFile, "ModuleID=%llu Function=%s edgeID=%u\n", moduleID,
-              F.getName().str().c_str(), afl_global_id);
+      dFile << "ModuleID=" << moduleID << " Function=" << F.getName().str() << " edgeID=" << afl_global_id << "\n";
 
     }
 
diff --git a/instrumentation/afl-llvm-dict2file.so.cc b/instrumentation/afl-llvm-dict2file.so.cc
index 4622e488..c4ad1783 100644
--- a/instrumentation/afl-llvm-dict2file.so.cc
+++ b/instrumentation/afl-llvm-dict2file.so.cc
@@ -65,7 +65,8 @@ using namespace llvm;
 namespace {
 
 class AFLdict2filePass : public ModulePass {
-
+  std::ofstream of;
+  void dict2file(u8 *, u32);
  public:
   static char ID;
 
@@ -81,7 +82,7 @@ class AFLdict2filePass : public ModulePass {
 
 }  // namespace
 
-void dict2file(int fd, u8 *mem, u32 len) {
+void AFLdict2filePass::dict2file(u8 *mem, u32 len) {
 
   u32  i, j, binary = 0;
   char line[MAX_AUTO_EXTRA * 8], tmp[8];
@@ -113,9 +114,8 @@ void dict2file(int fd, u8 *mem, u32 len) {
 
   line[j] = 0;
   strcat(line, "\"\n");
-  if (write(fd, line, strlen(line)) <= 0)
-    PFATAL("Could not write to dictionary file");
-  fsync(fd);
+  of << line;
+  of.flush();
 
   if (!be_quiet) fprintf(stderr, "Found dictionary token: %s", line);
 
@@ -125,7 +125,7 @@ bool AFLdict2filePass::runOnModule(Module &M) {
 
   DenseMap<Value *, std::string *> valueMap;
   char *                           ptr;
-  int                              fd, found = 0;
+  int                              found = 0;
 
   /* Show a banner */
   setvbuf(stdout, NULL, _IONBF, 0);
@@ -146,7 +146,8 @@ bool AFLdict2filePass::runOnModule(Module &M) {
   if (!ptr || *ptr != '/')
     FATAL("AFL_LLVM_DICT2FILE is not set to an absolute path: %s", ptr);
 
-  if ((fd = open(ptr, O_WRONLY | O_APPEND | O_CREAT | O_DSYNC, 0644)) < 0)
+  of.open(ptr, std::ofstream::out | std::ofstream::app);
+  if (!of.is_open())
     PFATAL("Could not open/create %s.", ptr);
 
   /* Instrument all the things! */
@@ -264,11 +265,11 @@ bool AFLdict2filePass::runOnModule(Module &M) {
 
               }
 
-              dict2file(fd, (u8 *)&val, len);
+              dict2file((u8 *)&val, len);
               found++;
               if (val2) {
 
-                dict2file(fd, (u8 *)&val2, len);
+                dict2file((u8 *)&val2, len);
                 found++;
 
               }
@@ -630,7 +631,7 @@ bool AFLdict2filePass::runOnModule(Module &M) {
 
           ptr = (char *)thestring.c_str();
 
-          dict2file(fd, (u8 *)ptr, optLen);
+          dict2file((u8 *)ptr, optLen);
           found++;
 
         }
@@ -641,7 +642,7 @@ bool AFLdict2filePass::runOnModule(Module &M) {
 
   }
 
-  close(fd);
+  of.close();
 
   /* Say something nice. */
 
diff --git a/instrumentation/afl-llvm-lto-instrumentation.so.cc b/instrumentation/afl-llvm-lto-instrumentation.so.cc
index 4a5738de..73e41f60 100644
--- a/instrumentation/afl-llvm-lto-instrumentation.so.cc
+++ b/instrumentation/afl-llvm-lto-instrumentation.so.cc
@@ -108,8 +108,8 @@ bool AFLLTOPass::runOnModule(Module &M) {
   //  std::vector<CallInst *>          calls;
   DenseMap<Value *, std::string *> valueMap;
   std::vector<BasicBlock *>        BlockList;
+  std::ofstream                    dFile;
   char *                           ptr;
-  FILE *                           documentFile = NULL;
   size_t                           found = 0;
 
   srand((unsigned int)time(NULL));
@@ -137,7 +137,8 @@ bool AFLLTOPass::runOnModule(Module &M) {
 
   if ((ptr = getenv("AFL_LLVM_DOCUMENT_IDS")) != NULL) {
 
-    if ((documentFile = fopen(ptr, "a")) == NULL)
+    dFile.open(ptr, std::ofstream::out | std::ofstream::app);
+    if (!dFile.is_open())
       WARNF("Cannot access document file %s", ptr);
 
   }
@@ -845,10 +846,9 @@ bool AFLLTOPass::runOnModule(Module &M) {
 
           }
 
-          if (documentFile) {
+          if (dFile.is_open()) {
 
-            fprintf(documentFile, "ModuleID=%llu Function=%s edgeID=%u\n",
-                    moduleID, F.getName().str().c_str(), afl_global_id);
+             dFile << "ModuleID=" << moduleID << " Function=" << F.getName().str() << " edgeID=" << afl_global_id << "\n";
 
           }
 
@@ -920,8 +920,7 @@ bool AFLLTOPass::runOnModule(Module &M) {
 
   }
 
-  if (documentFile) fclose(documentFile);
-  documentFile = NULL;
+  if (dFile.is_open()) dFile.close();
 
   // save highest location ID to global variable
   // do this after each function to fail faster