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