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