diff options
Diffstat (limited to 'instrumentation/afl-llvm-lto-instrumentation.so.cc')
| -rw-r--r-- | instrumentation/afl-llvm-lto-instrumentation.so.cc | 30 |
1 files changed, 12 insertions, 18 deletions
diff --git a/instrumentation/afl-llvm-lto-instrumentation.so.cc b/instrumentation/afl-llvm-lto-instrumentation.so.cc index e300044c..73e41f60 100644 --- a/instrumentation/afl-llvm-lto-instrumentation.so.cc +++ b/instrumentation/afl-llvm-lto-instrumentation.so.cc @@ -28,6 +28,7 @@ #include <sys/time.h> #include <list> +#include <memory> #include <string> #include <fstream> #include <set> @@ -107,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)); @@ -136,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); } @@ -844,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"; } @@ -919,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 @@ -1015,13 +1015,7 @@ bool AFLLTOPass::runOnModule(Module &M) { if (count) { - if ((ptr = (char *)malloc(memlen + count)) == NULL) { - - fprintf(stderr, "Error: malloc for %zu bytes failed!\n", - memlen + count); - exit(-1); - - } + auto ptrhld = std::unique_ptr<char []>(new char[memlen + count]); count = 0; @@ -1030,8 +1024,8 @@ bool AFLLTOPass::runOnModule(Module &M) { if (offset + token.length() < 0xfffff0 && count < MAX_AUTO_EXTRAS) { - ptr[offset++] = (uint8_t)token.length(); - memcpy(ptr + offset, token.c_str(), token.length()); + ptrhld.get()[offset++] = (uint8_t)token.length(); + memcpy(ptrhld.get() + offset, token.c_str(), token.length()); offset += token.length(); count++; @@ -1051,10 +1045,10 @@ bool AFLLTOPass::runOnModule(Module &M) { GlobalVariable *AFLInternalDictionary = new GlobalVariable( M, ArrayTy, true, GlobalValue::ExternalLinkage, ConstantDataArray::get(C, - *(new ArrayRef<char>((char *)ptr, offset))), + *(new ArrayRef<char>(ptrhld.get(), offset))), "__afl_internal_dictionary"); AFLInternalDictionary->setInitializer(ConstantDataArray::get( - C, *(new ArrayRef<char>((char *)ptr, offset)))); + C, *(new ArrayRef<char>(ptrhld.get(), offset)))); AFLInternalDictionary->setConstant(true); GlobalVariable *AFLDictionary = new GlobalVariable( |