1
#include "llvm/IR/Function.h"
2
#include "llvm/Pass.h"
3
#include "llvm/IR/LegacyPassManager.h"
4
#include "llvm/Transforms/IPO/PassManagerBuilder.h"
5
#include "llvm/Support/raw_ostream.h"
6
#include "llvm/IR/Instructions.h"
7
#include "llvm/Transforms/Utils.h"
8
#include "llvm/Support/CommandLine.h"
9
#include "llvm/Transforms/Utils/Local.h"
10
#include "../include/SplitBasicBlock.h"
11
#include "../include/Utils.h"
12
#include <vector>
13
#include <cstdlib>
14
#include <ctime>
15
using namespace llvm;
16
using std::vector;
17

18
namespace
19
{
20
    class Flattening : public FunctionPass{
21
        public:
22
            static char ID;
23
            Flattening() : FunctionPass(ID) {
24
                srand(time(0));
25
            }
26

27
            // 对函数 F 进行平坦化
28
            void flatten(Function &F);
29

30
            bool runOnFunction(Function &F);
31
    };
32
} // namespace
33

34
// runOnFunction函数实现
35
bool Flattening :: runOnFunction(Function &F){
36
    INIT_CONTEXT(F);
37
    FunctionPass *pass = createSplitBasicBlockPass();
38
    pass->runOnFunction(F);
39
    flatten(F);
40
    return true;
41
}
42

43
void Flattening::flatten(Function &F){
44
    // 基本块数量不超过1则无需平坦化
45
    if(F.size() <= 1){
46
        return;
47
    }
48

49
    //将除入口块以外的基本块保存到 vector 容器中，方便后续处理
50
    // 首先保存所有基本块
51
    vector<BasicBlock*> origBB;
52
    for(BasicBlock &BB : F){
53
        origBB.push_back(&BB);
54
    }
55
    // 从 vector 中去除第一个基本块
56
    // 因为第一个基本块一定是入口块
57
    origBB.erase(origBB.begin());
58
    BasicBlock &entryBB = F.getEntryBlock();
59
    // 判断：如果第一个基本块的末尾是条件跳转 -> 单独分离
60
    // dyn_cast：用来类型判断和类型转换，如果给定参数（entryBB.getRerminator()）符合想要类型（BranchInst），进行转换，转换成（BranchInst *br），如果不符合类型，返回一个 nullptr 空指针
61
    if(BranchInst *br = dyn_cast<BranchInst>(entryBB.getTerminator())){
62
        if(br->isConditional()){
63
            // 分离跳转指令作为新基本快，插入到 vector 容器最前端
64
            BasicBlock *newBB = entryBB.splitBasicBlock(br, "newBB");
65
            origBB.insert(origBB.begin(), newBB);
66
        }
67
    }
68

69
    // 创建分发块和返回块
70
    BasicBlock *dispatchBB = BasicBlock::Create(*CONTEXT, "dispatchBB", &F, &entryBB);
71
    BasicBlock *returnBB = BasicBlock::Create(*CONTEXT, "returnBB", &F, &entryBB);
72
    // 此时的块顺序如下
73
    /*-----------------------
74
        分发块
75
        返回块
76
        入口块
77
    -----------------------*/
78
    // 建立分发块和返回块的链接
79
    BranchInst::Create(dispatchBB, returnBB);
80
    // 把入口块移到最前面，分发块前面
81
    entryBB.moveBefore(dispatchBB);
82
    // 此时的块顺序如下
83
    /*-----------------------
84
        入口块
85
        分发块
86
        返回块
87
    -----------------------*/
88
    // 去除入口块末尾的跳转
89
    entryBB.getTerminator()->eraseFromParent();
90
    // 使入口块跳转到 dispatchBB，brDispatchBB 也就是入口块的终结指令
91
    BranchInst *brDispatchBB = BranchInst::Create(dispatchBB, &entryBB);
92

93
    // 在入口块插入 alloca 和 store 指令创建并初始化 switch 变量，初始值为随机值
94
    int randNumCase = rand();
95
    // 在入口块跳转分发块指令（brDispatchBB）之前插入一条申请分配内存的指令，使用 AllocaInst 的构造函数
96
    AllocaInst *swVarPtr = new AllocaInst(TYPE_I32, 0, "swVar.ptr", brDispatchBB);
97
    // 在入口块跳转分发块指令（brDispatchBB）之前插入一条内存存储的指令
98
    new StoreInst(CONST_I32(randNumCase), swVarPtr, brDispatchBB);
99
    // 在分发块插入 load 指令读取 switch 变量
100
    LoadInst *swVar = new LoadInst(TYPE_I32, swVarPtr, "swVar", false, dispatchBB);
101
    // 在返回块前创建 switch 默认跳转的基本块
102
    BasicBlock *swDefault = BasicBlock::Create(*CONTEXT, "seDefault", &F, returnBB);
103
    // switch 默认跳转的基本块默认跳转到返回块
104
    BranchInst::Create(returnBB, swDefault);
105
    // 在分发块末尾的 switch 块中添加 switch 指令
106
    SwitchInst *swInst = SwitchInst::Create(swVar, swDefault, 0, dispatchBB);
107
    // 将原基本块插入到返回块之前，并分配 case 值
108
    for(BasicBlock *BB : origBB){
109
        BB->moveBefore(returnBB);
110
        swInst->addCase(CONST_I32(randNumCase), BB);
111
        randNumCase = rand();
112
    }
113

114
    // 在每个基本块最后添加修改 switch 变量的指令和跳转到返回块的指令
115
    for(BasicBlock *BB : origBB){
116
        // retn BB
117
        if(BB->getTerminator()->getNumSuccessors() == 0){
118
            continue;
119
        }
120
        // 非条件跳转
121
        else if(BB->getTerminator()->getNumSuccessors() == 1){
122
            // 获取后继块
123
            BasicBlock *sucBB = BB->getTerminator()->getSuccessor(0);
124
            // 删除基本块最后一条跳转到后继块的指令
125
            BB->getTerminator()->eraseFromParent();
126
            // 获取后继块的 case 值，（常量类型）
127
            ConstantInt *numCase = swInst->findCaseDest(sucBB);
128
            // 把 switch 变量修改成它的后继块的 case 值，下一轮进入分发块就根据这个值跳转到后继块
129
            new StoreInst(numCase, swVarPtr, BB);
130
            // 基本块跳转到返回块
131
            BranchInst::Create(returnBB, BB);
132
        }
133
        // 条件跳转
134
        else if(BB->getTerminator()->getNumSuccessors() == 2){
135
            ConstantInt *numCaseTrue = swInst->findCaseDest(BB->getTerminator()->getSuccessor(0));
136
            ConstantInt *numCaseFalse = swInst->findCaseDest(BB->getTerminator()->getSuccessor(1));
137
            BranchInst *br = cast<BranchInst>(BB->getTerminator());
138
            // 创建 select 指令判断需要修改哪一个 case 值
139
            SelectInst *sel = SelectInst::Create(br->getCondition(), numCaseTrue, numCaseFalse, "", BB->getTerminator());
140
            // 删除基本块最后一条跳转到后继块的指令
141
            BB->getTerminator()->eraseFromParent();
142
            // 把 switch 变量修改成它的后继块的 case 值，下一轮进入分发块就根据这个值跳转到后继块
143
            new StoreInst(sel, swVarPtr, BB);
144
            // 基本块跳转到返回块
145
            BranchInst::Create(returnBB, BB);
146
        }
147
    }
148
    fixStack(F);
149
}
150

151
// https://www.leadroyal.cn/p/1072/
152
// 方案一： 注册阶段，将LowerSwitchPass放到Flatten前，执行时按顺序执行即可。
153
#if LLVM_VERSION_MAJOR >= 9
154
    // FIXME: Pass has not been inserted into a PassManager object!
155
    // >=9.0, LowerSwitchPass depends on LazyValueInfoWrapperPass, which cause AssertError.
156
    // So I move LowerSwitchPass into register function, just before FlatteningPass.
157
#else
158
    FunctionPass *lower = createLowerSwitchPass();
159
    lower->runOnFunction(*F);
160
#endif
161

162
char Flattening :: ID = 0;
163
// 注册该 Flattening Pass
164
static RegisterPass<Flattening> x("fla", "My control flow flattening obfuscation.");

1
#include "../include/Utils.h"
2
#include <vector>
3
#include "llvm/IR/Instructions.h"
4
#include "llvm/Transforms/Utils/Local.h"
5
#include "llvm/Transforms/Utils/ValueMapper.h"
6
#include "llvm/Transforms/Utils/Cloning.h"
7

8
using std::vector;
9
using namespace llvm;
10

11
LLVMContext *CONTEXT = nullptr;
12

13
void llvm::fixStack(Function &F) {
14
    vector<PHINode*> origPHI;
15
    vector<Instruction*> origReg;
16
    BasicBlock &entryBB = F.getEntryBlock();
17
    for(BasicBlock &BB : F){
18
        for(Instruction &I : BB){
19
            if(PHINode *PN = dyn_cast<PHINode>(&I)){
20
                origPHI.push_back(PN);
21
            }else if(!(isa<AllocaInst>(&I) && I.getParent() == &entryBB)
22
                && I.isUsedOutsideOfBlock(&BB)){
23
                origReg.push_back(&I);
24
            }
25
        }
26
    }
27
    for(PHINode *PN : origPHI){
28
        DemotePHIToStack(PN, entryBB.getTerminator());
29
    }
30
    for(Instruction *I : origReg){
31
        DemoteRegToStack(*I, entryBB.getTerminator());
32
    }
33
}

1
#ifndef _UTILS_H_
2
#define _UTILS_H_
3

4
#include "llvm/IR/Function.h"
5
// 获取当前函数 context
6
#define INIT_CONTEXT(F) CONTEXT=&F.getContext()
7
#define TYPE_I32 Type::getInt32Ty(*CONTEXT)
8
// 创建常量
9
#define CONST_I32(V) ConstantInt::get(TYPE_I32, V, false)
10
#define CONST(T, V) ConstantInt::get(T, V)
11

12
extern llvm::LLVMContext *CONTEXT;
13

14
// 对 PHI 指令和逃逸变量进行修复函数，添加到命名空间 llvm 里
15
namespace llvm{
16
    void fixStack(Function &F);
17
}
18

19
#endif