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