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