diff options
Diffstat (limited to 'llvm_mode/split-switches-pass.so.cc')
| -rw-r--r-- | llvm_mode/split-switches-pass.so.cc | 199 |
1 files changed, 193 insertions, 6 deletions
diff --git a/llvm_mode/split-switches-pass.so.cc b/llvm_mode/split-switches-pass.so.cc index 2743a71a..9101dc26 100644 --- a/llvm_mode/split-switches-pass.so.cc +++ b/llvm_mode/split-switches-pass.so.cc @@ -18,6 +18,13 @@ #include <stdlib.h> #include <unistd.h> +#include <list> +#include <string> +#include <fstream> +#include <sys/time.h> + +#include "llvm/Config/llvm-config.h" + #include "llvm/ADT/Statistic.h" #include "llvm/IR/IRBuilder.h" #include "llvm/IR/LegacyPassManager.h" @@ -26,10 +33,20 @@ #include "llvm/Support/raw_ostream.h" #include "llvm/Transforms/IPO/PassManagerBuilder.h" #include "llvm/Transforms/Utils/BasicBlockUtils.h" -#include "llvm/IR/Verifier.h" #include "llvm/Pass.h" #include "llvm/Analysis/ValueTracking.h" +#include "llvm/IR/IRBuilder.h" +#if LLVM_VERSION_MAJOR > 3 || \ + (LLVM_VERSION_MAJOR == 3 && LLVM_VERSION_MINOR > 4) +#include "llvm/IR/Verifier.h" +#include "llvm/IR/DebugInfo.h" +#else +#include "llvm/Analysis/Verifier.h" +#include "llvm/DebugInfo.h" +#define nullptr 0 +#endif + #include <set> using namespace llvm; @@ -42,6 +59,41 @@ class SplitSwitchesTransform : public ModulePass { static char ID; SplitSwitchesTransform() : ModulePass(ID) { + char *instWhiteListFilename = getenv("AFL_LLVM_WHITELIST"); + if (instWhiteListFilename) { + + std::string line; + std::ifstream fileStream; + fileStream.open(instWhiteListFilename); + if (!fileStream) report_fatal_error("Unable to open AFL_LLVM_WHITELIST"); + getline(fileStream, line); + while (fileStream) { + + myWhitelist.push_back(line); + getline(fileStream, line); + + } + + } + + } + + static bool isBlacklisted(const Function *F) { + + static const char *Blacklist[] = { + + "asan.", "llvm.", "sancov.", "__ubsan_handle_", "ign." + + }; + + for (auto const &BlacklistFunc : Blacklist) { + + if (F->getName().startswith(BlacklistFunc)) { return true; } + + } + + return false; + } bool runOnModule(Module &M) override; @@ -71,6 +123,9 @@ class SplitSwitchesTransform : public ModulePass { typedef std::vector<CaseExpr> CaseVector; + protected: + std::list<std::string> myWhitelist; + private: bool splitSwitches(Module &M); bool transformCmps(Module &M, const bool processStrcmp, @@ -96,7 +151,7 @@ BasicBlock *SplitSwitchesTransform::switchConvert( IntegerType * ByteType = IntegerType::get(OrigBlock->getContext(), 8); unsigned BytesInValue = bytesChecked.size(); std::vector<uint8_t> setSizes; - std::vector<std::set<uint8_t>> byteSets(BytesInValue, std::set<uint8_t>()); + std::vector<std::set<uint8_t> > byteSets(BytesInValue, std::set<uint8_t>()); assert(ValTypeBitWidth >= 8 && ValTypeBitWidth <= 64); @@ -169,8 +224,25 @@ BasicBlock *SplitSwitchesTransform::switchConvert( NewNode->getInstList().push_back(Comp); bytesChecked[smallestIndex] = true; - if (std::all_of(bytesChecked.begin(), bytesChecked.end(), - [](bool b) { return b; })) { + bool allBytesAreChecked = true; + + for (std::vector<bool>::iterator BCI = bytesChecked.begin(), + E = bytesChecked.end(); + BCI != E; ++BCI) { + + if (!*BCI) { + + allBytesAreChecked = false; + break; + + } + + } + + // if (std::all_of(bytesChecked.begin(), bytesChecked.end(), + // [](bool b) { return b; })) { + + if (allBytesAreChecked) { assert(Cases.size() == 1); BranchInst::Create(Cases[0].BB, NewDefault, Comp, NewNode); @@ -262,16 +334,132 @@ BasicBlock *SplitSwitchesTransform::switchConvert( bool SplitSwitchesTransform::splitSwitches(Module &M) { +#if (LLVM_VERSION_MAJOR == 3 && LLVM_VERSION_MINOR < 7) + LLVMContext &C = M.getContext(); +#endif + std::vector<SwitchInst *> switches; /* iterate over all functions, bbs and instruction and add * all switches to switches vector for later processing */ for (auto &F : M) { + if (isBlacklisted(&F)) continue; + for (auto &BB : F) { SwitchInst *switchInst = nullptr; + if (!myWhitelist.empty()) { + + bool instrumentBlock = false; + BasicBlock::iterator IP = BB.getFirstInsertionPt(); + + /* Get the current location using debug information. + * For now, just instrument the block if we are not able + * to determine our location. */ + DebugLoc Loc = IP->getDebugLoc(); +#if LLVM_VERSION_MAJOR >= 4 || \ + (LLVM_VERSION_MAJOR == 3 && LLVM_VERSION_MINOR >= 7) + if (Loc) { + + DILocation *cDILoc = dyn_cast<DILocation>(Loc.getAsMDNode()); + + unsigned int instLine = cDILoc->getLine(); + StringRef instFilename = cDILoc->getFilename(); + + if (instFilename.str().empty()) { + + /* If the original location is empty, try using the inlined location + */ + DILocation *oDILoc = cDILoc->getInlinedAt(); + if (oDILoc) { + + instFilename = oDILoc->getFilename(); + instLine = oDILoc->getLine(); + + } + + } + + (void)instLine; + + /* Continue only if we know where we actually are */ + if (!instFilename.str().empty()) { + + for (std::list<std::string>::iterator it = myWhitelist.begin(); + it != myWhitelist.end(); ++it) { + + /* We don't check for filename equality here because + * filenames might actually be full paths. Instead we + * check that the actual filename ends in the filename + * specified in the list. */ + if (instFilename.str().length() >= it->length()) { + + if (instFilename.str().compare( + instFilename.str().length() - it->length(), + it->length(), *it) == 0) { + + instrumentBlock = true; + break; + + } + + } + + } + + } + + } + +#else + if (!Loc.isUnknown()) { + + DILocation cDILoc(Loc.getAsMDNode(C)); + + unsigned int instLine = cDILoc.getLineNumber(); + StringRef instFilename = cDILoc.getFilename(); + + (void)instLine; + + /* Continue only if we know where we actually are */ + if (!instFilename.str().empty()) { + + for (std::list<std::string>::iterator it = myWhitelist.begin(); + it != myWhitelist.end(); ++it) { + + /* We don't check for filename equality here because + * filenames might actually be full paths. Instead we + * check that the actual filename ends in the filename + * specified in the list. */ + if (instFilename.str().length() >= it->length()) { + + if (instFilename.str().compare( + instFilename.str().length() - it->length(), + it->length(), *it) == 0) { + + instrumentBlock = true; + break; + + } + + } + + } + + } + + } + +#endif + + /* Either we couldn't figure out our location or the location is + * not whitelisted, so we skip instrumentation. */ + if (!instrumentBlock) continue; + + } + if ((switchInst = dyn_cast<SwitchInst>(BB.getTerminator()))) { if (switchInst->getNumCases() < 1) continue; @@ -313,8 +501,7 @@ bool SplitSwitchesTransform::splitSwitches(Module &M) { * if the default block is set as an unreachable we avoid creating one * because will never be a valid target.*/ BasicBlock *NewDefault = nullptr; - NewDefault = BasicBlock::Create(SI->getContext(), "NewDefault"); - NewDefault->insertInto(F, Default); + NewDefault = BasicBlock::Create(SI->getContext(), "NewDefault", F, Default); BranchInst::Create(Default, NewDefault); /* Prepare cases vector. */ |