diff options
| -rw-r--r-- | instrumentation/SanitizerCoverageLTO.so.cc | 17 | ||||
| -rw-r--r-- | instrumentation/afl-llvm-dict2file.so.cc | 23 | ||||
| -rw-r--r-- | instrumentation/afl-llvm-lto-instrumentation.so.cc | 13 |
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 e300044c..d685e76c 100644 --- a/instrumentation/afl-llvm-lto-instrumentation.so.cc +++ b/instrumentation/afl-llvm-lto-instrumentation.so.cc @@ -107,8 +107,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)); @@ -136,7 +136,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); } @@ -844,10 +845,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"; } @@ -919,8 +919,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 |