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