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
39 using namespace llvm;
40 #define DEBUG_TYPE "instrorderfile"
41
42 static 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
48 namespace {
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
55 std::mutex MappingMutex;
56
57 struct InstrOrderFile {
58 private:
59 GlobalVariable *OrderFileBuffer;
60 GlobalVariable *BufferIdx;
61 GlobalVariable *BitMap;
62 ArrayType *BufferTy;
63 ArrayType *MapTy;
64
65 public:
InstrOrderFile__anon496844480111::InstrOrderFile66 InstrOrderFile() {}
67
createOrderFileData__anon496844480111::InstrOrderFile68 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.
generateCodeSequence__anon496844480111::InstrOrderFile100 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
run__anon496844480111::InstrOrderFile160 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
176 class InstrOrderFileLegacyPass : public ModulePass {
177 public:
178 static char ID;
179
InstrOrderFileLegacyPass()180 InstrOrderFileLegacyPass() : ModulePass(ID) {
181 initializeInstrOrderFileLegacyPassPass(
182 *PassRegistry::getPassRegistry());
183 }
184
185 bool runOnModule(Module &M) override;
186 };
187
188 } // End anonymous namespace
189
runOnModule(Module & M)190 bool InstrOrderFileLegacyPass::runOnModule(Module &M) {
191 if (skipModule(M))
192 return false;
193
194 return InstrOrderFile().run(M);
195 }
196
197 PreservedAnalyses
run(Module & M,ModuleAnalysisManager & AM)198 InstrOrderFilePass::run(Module &M, ModuleAnalysisManager &AM) {
199 if (InstrOrderFile().run(M))
200 return PreservedAnalyses::none();
201 return PreservedAnalyses::all();
202 }
203
204 INITIALIZE_PASS_BEGIN(InstrOrderFileLegacyPass, "instrorderfile",
205 "Instrumentation for Order File", false, false)
206 INITIALIZE_PASS_END(InstrOrderFileLegacyPass, "instrorderfile",
207 "Instrumentation for Order File", false, false)
208
209 char InstrOrderFileLegacyPass::ID = 0;
210
createInstrOrderFilePass()211 ModulePass *llvm::createInstrOrderFilePass() {
212 return new InstrOrderFileLegacyPass();
213 }
214