LLVM passes making slightly more C++

1 files changed, 12 insertions, 11 deletions

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. */