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/Instrumentation.h"
23#include <fstream>
24#include <mutex>
25#include <sstream>
26
27using namespace llvm;
28#define DEBUG_TYPE "instrorderfile"
29
30static cl::opt<std::string> ClOrderFileWriteMapping(
31 "orderfile-write-mapping", cl::init(Val: ""),
32 cl::desc(
33 "Dump functions and their MD5 hash to deobfuscate profile data"),
34 cl::Hidden);
35
36namespace {
37
38// We need a global bitmap to tell if a function is executed. We also
39// need a global variable to save the order of functions. We can use a
40// fixed-size buffer that saves the MD5 hash of the function. We need
41// a global variable to save the index into the buffer.
42
43std::mutex MappingMutex;
44
45struct InstrOrderFile {
46private:
47 GlobalVariable *OrderFileBuffer;
48 GlobalVariable *BufferIdx;
49 GlobalVariable *BitMap;
50 ArrayType *BufferTy;
51 ArrayType *MapTy;
52
53public:
54 InstrOrderFile() = default;
55
56 void createOrderFileData(Module &M) {
57 LLVMContext &Ctx = M.getContext();
58 int NumFunctions = 0;
59 for (Function &F : M) {
60 if (!F.isDeclaration())
61 NumFunctions++;
62 }
63
64 BufferTy =
65 ArrayType::get(ElementType: Type::getInt64Ty(C&: Ctx), INSTR_ORDER_FILE_BUFFER_SIZE);
66 Type *IdxTy = Type::getInt32Ty(C&: Ctx);
67 MapTy = ArrayType::get(ElementType: Type::getInt8Ty(C&: Ctx), NumElements: NumFunctions);
68
69 // Create the global variables.
70 std::string SymbolName = INSTR_PROF_ORDERFILE_BUFFER_NAME_STR;
71 OrderFileBuffer = new GlobalVariable(M, BufferTy, false, GlobalValue::LinkOnceODRLinkage,
72 Constant::getNullValue(Ty: BufferTy), SymbolName);
73 Triple TT = Triple(M.getTargetTriple());
74 OrderFileBuffer->setSection(
75 getInstrProfSectionName(IPSK: IPSK_orderfile, OF: TT.getObjectFormat()));
76
77 std::string IndexName = INSTR_PROF_ORDERFILE_BUFFER_IDX_NAME_STR;
78 BufferIdx = new GlobalVariable(M, IdxTy, false, GlobalValue::LinkOnceODRLinkage,
79 Constant::getNullValue(Ty: IdxTy), IndexName);
80
81 std::string BitMapName = "bitmap_0";
82 BitMap = new GlobalVariable(M, MapTy, false, GlobalValue::PrivateLinkage,
83 Constant::getNullValue(Ty: MapTy), BitMapName);
84 }
85
86 // Generate the code sequence in the entry block of each function to
87 // update the buffer.
88 void generateCodeSequence(Module &M, Function &F, int FuncId) {
89 if (!ClOrderFileWriteMapping.empty()) {
90 std::lock_guard<std::mutex> LogLock(MappingMutex);
91 std::error_code EC;
92 llvm::raw_fd_ostream OS(ClOrderFileWriteMapping, EC,
93 llvm::sys::fs::OF_Append);
94 if (EC) {
95 report_fatal_error(reason: Twine("Failed to open ") + ClOrderFileWriteMapping +
96 " to save mapping file for order file instrumentation\n");
97 } else {
98 std::stringstream stream;
99 stream << std::hex << MD5Hash(Str: F.getName());
100 std::string singleLine = "MD5 " + stream.str() + " " +
101 std::string(F.getName()) + '\n';
102 OS << singleLine;
103 }
104 }
105
106 BasicBlock *OrigEntry = &F.getEntryBlock();
107
108 LLVMContext &Ctx = M.getContext();
109 IntegerType *Int32Ty = Type::getInt32Ty(C&: Ctx);
110 IntegerType *Int8Ty = Type::getInt8Ty(C&: Ctx);
111
112 // Create a new entry block for instrumentation. We will check the bitmap
113 // in this basic block.
114 BasicBlock *NewEntry =
115 BasicBlock::Create(Context&: M.getContext(), Name: "order_file_entry", Parent: &F, InsertBefore: OrigEntry);
116 IRBuilder<> entryB(NewEntry);
117 // Create a basic block for updating the circular buffer.
118 BasicBlock *UpdateOrderFileBB =
119 BasicBlock::Create(Context&: M.getContext(), Name: "order_file_set", Parent: &F, InsertBefore: OrigEntry);
120 IRBuilder<> updateB(UpdateOrderFileBB);
121
122 // Check the bitmap, if it is already 1, do nothing.
123 // Otherwise, set the bit, grab the index, update the buffer.
124 Value *IdxFlags[] = {ConstantInt::get(Ty: Int32Ty, V: 0),
125 ConstantInt::get(Ty: Int32Ty, V: FuncId)};
126 Value *MapAddr = entryB.CreateGEP(Ty: MapTy, Ptr: BitMap, IdxList: IdxFlags, Name: "");
127 LoadInst *loadBitMap = entryB.CreateLoad(Ty: Int8Ty, Ptr: MapAddr, Name: "");
128 entryB.CreateStore(Val: ConstantInt::get(Ty: Int8Ty, V: 1), Ptr: MapAddr);
129 Value *IsNotExecuted =
130 entryB.CreateICmpEQ(LHS: loadBitMap, RHS: ConstantInt::get(Ty: Int8Ty, V: 0));
131 entryB.CreateCondBr(Cond: IsNotExecuted, True: UpdateOrderFileBB, False: OrigEntry);
132
133 // Fill up UpdateOrderFileBB: grab the index, update the buffer!
134 Value *IdxVal = updateB.CreateAtomicRMW(
135 Op: AtomicRMWInst::Add, Ptr: BufferIdx, Val: ConstantInt::get(Ty: Int32Ty, V: 1),
136 Align: MaybeAlign(), Ordering: AtomicOrdering::SequentiallyConsistent);
137 // We need to wrap around the index to fit it inside the buffer.
138 Value *WrappedIdx = updateB.CreateAnd(
139 LHS: IdxVal, RHS: ConstantInt::get(Ty: Int32Ty, INSTR_ORDER_FILE_BUFFER_MASK));
140 Value *BufferGEPIdx[] = {ConstantInt::get(Ty: Int32Ty, V: 0), WrappedIdx};
141 Value *BufferAddr =
142 updateB.CreateGEP(Ty: BufferTy, Ptr: OrderFileBuffer, IdxList: BufferGEPIdx, Name: "");
143 updateB.CreateStore(Val: ConstantInt::get(Ty: Type::getInt64Ty(C&: Ctx), V: MD5Hash(Str: F.getName())),
144 Ptr: BufferAddr);
145 updateB.CreateBr(Dest: OrigEntry);
146 }
147
148 bool run(Module &M) {
149 createOrderFileData(M);
150
151 int FuncId = 0;
152 for (Function &F : M) {
153 if (F.isDeclaration())
154 continue;
155 generateCodeSequence(M, F, FuncId);
156 ++FuncId;
157 }
158
159 return true;
160 }
161
162}; // End of InstrOrderFile struct
163} // End anonymous namespace
164
165PreservedAnalyses
166InstrOrderFilePass::run(Module &M, ModuleAnalysisManager &AM) {
167 if (InstrOrderFile().run(M))
168 return PreservedAnalyses::none();
169 return PreservedAnalyses::all();
170}
171