1 //===- InstrOrderFile.cpp ---- Late IR instrumentation for order file ----===// 2 // 3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. 4 // See https://llvm.org/LICENSE.txt for license information. 5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception 6 // 7 //===----------------------------------------------------------------------===// 8 // 9 //===----------------------------------------------------------------------===// 10 11 #include "llvm/Transforms/Instrumentation/InstrOrderFile.h" 12 #include "llvm/IR/Constants.h" 13 #include "llvm/IR/Function.h" 14 #include "llvm/IR/GlobalValue.h" 15 #include "llvm/IR/IRBuilder.h" 16 #include "llvm/IR/Instructions.h" 17 #include "llvm/IR/Module.h" 18 #include "llvm/InitializePasses.h" 19 #include "llvm/Pass.h" 20 #include "llvm/PassRegistry.h" 21 #include "llvm/ProfileData/InstrProf.h" 22 #include "llvm/Support/CommandLine.h" 23 #include "llvm/Support/FileSystem.h" 24 #include "llvm/Support/raw_ostream.h" 25 #include "llvm/Transforms/Instrumentation.h" 26 #include <fstream> 27 #include <mutex> 28 #include <sstream> 29 30 using namespace llvm; 31 #define DEBUG_TYPE "instrorderfile" 32 33 static cl::opt<std::string> ClOrderFileWriteMapping( 34 "orderfile-write-mapping", cl::init(""), 35 cl::desc( 36 "Dump functions and their MD5 hash to deobfuscate profile data"), 37 cl::Hidden); 38 39 namespace { 40 41 // We need a global bitmap to tell if a function is executed. We also 42 // need a global variable to save the order of functions. We can use a 43 // fixed-size buffer that saves the MD5 hash of the function. We need 44 // a global variable to save the index into the buffer. 45 46 std::mutex MappingMutex; 47 48 struct InstrOrderFile { 49 private: 50 GlobalVariable *OrderFileBuffer; 51 GlobalVariable *BufferIdx; 52 GlobalVariable *BitMap; 53 ArrayType *BufferTy; 54 ArrayType *MapTy; 55 56 public: 57 InstrOrderFile() = default; 58 59 void createOrderFileData(Module &M) { 60 LLVMContext &Ctx = M.getContext(); 61 int NumFunctions = 0; 62 for (Function &F : M) { 63 if (!F.isDeclaration()) 64 NumFunctions++; 65 } 66 67 BufferTy = 68 ArrayType::get(Type::getInt64Ty(Ctx), INSTR_ORDER_FILE_BUFFER_SIZE); 69 Type *IdxTy = Type::getInt32Ty(Ctx); 70 MapTy = ArrayType::get(Type::getInt8Ty(Ctx), NumFunctions); 71 72 // Create the global variables. 73 std::string SymbolName = INSTR_PROF_ORDERFILE_BUFFER_NAME_STR; 74 OrderFileBuffer = new GlobalVariable(M, BufferTy, false, GlobalValue::LinkOnceODRLinkage, 75 Constant::getNullValue(BufferTy), SymbolName); 76 Triple TT = Triple(M.getTargetTriple()); 77 OrderFileBuffer->setSection( 78 getInstrProfSectionName(IPSK_orderfile, TT.getObjectFormat())); 79 80 std::string IndexName = INSTR_PROF_ORDERFILE_BUFFER_IDX_NAME_STR; 81 BufferIdx = new GlobalVariable(M, IdxTy, false, GlobalValue::LinkOnceODRLinkage, 82 Constant::getNullValue(IdxTy), IndexName); 83 84 std::string BitMapName = "bitmap_0"; 85 BitMap = new GlobalVariable(M, MapTy, false, GlobalValue::PrivateLinkage, 86 Constant::getNullValue(MapTy), BitMapName); 87 } 88 89 // Generate the code sequence in the entry block of each function to 90 // update the buffer. 91 void generateCodeSequence(Module &M, Function &F, int FuncId) { 92 if (!ClOrderFileWriteMapping.empty()) { 93 std::lock_guard<std::mutex> LogLock(MappingMutex); 94 std::error_code EC; 95 llvm::raw_fd_ostream OS(ClOrderFileWriteMapping, EC, 96 llvm::sys::fs::OF_Append); 97 if (EC) { 98 report_fatal_error(Twine("Failed to open ") + ClOrderFileWriteMapping + 99 " to save mapping file for order file instrumentation\n"); 100 } else { 101 std::stringstream stream; 102 stream << std::hex << MD5Hash(F.getName()); 103 std::string singleLine = "MD5 " + stream.str() + " " + 104 std::string(F.getName()) + '\n'; 105 OS << singleLine; 106 } 107 } 108 109 BasicBlock *OrigEntry = &F.getEntryBlock(); 110 111 LLVMContext &Ctx = M.getContext(); 112 IntegerType *Int32Ty = Type::getInt32Ty(Ctx); 113 IntegerType *Int8Ty = Type::getInt8Ty(Ctx); 114 115 // Create a new entry block for instrumentation. We will check the bitmap 116 // in this basic block. 117 BasicBlock *NewEntry = 118 BasicBlock::Create(M.getContext(), "order_file_entry", &F, OrigEntry); 119 IRBuilder<> entryB(NewEntry); 120 // Create a basic block for updating the circular buffer. 121 BasicBlock *UpdateOrderFileBB = 122 BasicBlock::Create(M.getContext(), "order_file_set", &F, OrigEntry); 123 IRBuilder<> updateB(UpdateOrderFileBB); 124 125 // Check the bitmap, if it is already 1, do nothing. 126 // Otherwise, set the bit, grab the index, update the buffer. 127 Value *IdxFlags[] = {ConstantInt::get(Int32Ty, 0), 128 ConstantInt::get(Int32Ty, FuncId)}; 129 Value *MapAddr = entryB.CreateGEP(MapTy, BitMap, IdxFlags, ""); 130 LoadInst *loadBitMap = entryB.CreateLoad(Int8Ty, MapAddr, ""); 131 entryB.CreateStore(ConstantInt::get(Int8Ty, 1), MapAddr); 132 Value *IsNotExecuted = 133 entryB.CreateICmpEQ(loadBitMap, ConstantInt::get(Int8Ty, 0)); 134 entryB.CreateCondBr(IsNotExecuted, UpdateOrderFileBB, OrigEntry); 135 136 // Fill up UpdateOrderFileBB: grab the index, update the buffer! 137 Value *IdxVal = updateB.CreateAtomicRMW( 138 AtomicRMWInst::Add, BufferIdx, ConstantInt::get(Int32Ty, 1), 139 MaybeAlign(), AtomicOrdering::SequentiallyConsistent); 140 // We need to wrap around the index to fit it inside the buffer. 141 Value *WrappedIdx = updateB.CreateAnd( 142 IdxVal, ConstantInt::get(Int32Ty, INSTR_ORDER_FILE_BUFFER_MASK)); 143 Value *BufferGEPIdx[] = {ConstantInt::get(Int32Ty, 0), WrappedIdx}; 144 Value *BufferAddr = 145 updateB.CreateGEP(BufferTy, OrderFileBuffer, BufferGEPIdx, ""); 146 updateB.CreateStore(ConstantInt::get(Type::getInt64Ty(Ctx), MD5Hash(F.getName())), 147 BufferAddr); 148 updateB.CreateBr(OrigEntry); 149 } 150 151 bool run(Module &M) { 152 createOrderFileData(M); 153 154 int FuncId = 0; 155 for (Function &F : M) { 156 if (F.isDeclaration()) 157 continue; 158 generateCodeSequence(M, F, FuncId); 159 ++FuncId; 160 } 161 162 return true; 163 } 164 165 }; // End of InstrOrderFile struct 166 } // End anonymous namespace 167 168 PreservedAnalyses 169 InstrOrderFilePass::run(Module &M, ModuleAnalysisManager &AM) { 170 if (InstrOrderFile().run(M)) 171 return PreservedAnalyses::none(); 172 return PreservedAnalyses::all(); 173 } 174