llvm/lib/Transforms/Instrumentation/InstrOrderFile.cpp

   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 }