OLLVM Pluto 之平坦化增强源码学习

Posted by Qfrost on 2022-07-27
Estimated Reading Time 11 Minutes
Words 2.4k In Total
Viewed Times

(抄)了一手Pluto的平坦化增强(FlatteningEnhanced),不知道为啥Pluto的这份代码和其他的代码风格差异极大,代码可读性很差,看了好久才看明白是什么意思。

Pluto平坦化增强(FlatteningEnhanced)的核心思想是:自实现真实块的分发方式,使用树结构取代标准平坦化的switch进行分发。在构建平坦化之前,生成一颗不规则的随机的二叉树,其叶子节点数目等于真实块数量,其非叶子节点出度均为2,故生成的树最终会有n个节点,n=2*真实块数量-1。生成树后,将真实块排列到叶子节点上(每个真实块与叶子节点一一对应),为每个非叶子节点建立一个分发块,而后从根节点开始向下,为每个节点赋予一个权值。根节点的权值为 [0, MAX),其两个孩子节点的权值为 [0, val)和 [val, MAX), val为0到MAX间的一个随机数(但这个随机数会保证其下的叶子节点均能赋到长度至少为1的权值)。以此类推,最终所有的叶子节点(真实块)会得到一个唯一的权值区间,且所有叶子节点的权值区间的并集为[0, MAX)。 对每个非叶子节点的分发块末尾做条件跳转,根据switchVar的值处于左孩子节点的权值区间还是右孩子节点的权值区间判断跳转目标,以此保证最终都能分发到真实块上。

建树后,其他就和标准平坦化一样了,处理真实块的末尾跳转,根据跳转条件,为switchVar变量赋值,而后跳转到loopEnd上,再由loopEnd跳转回loopEntry。画了一个草图说明这个算法的分发流程。

FlatteningEnhanced

下面做一个源码解析(代码经过一定程度的重构,但核心思想不变

先定义树结构体

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
struct TreeNode
{

unsigned int val=0,limit=0,l=0,r=0;
TreeNode *left=NULL,*right=NULL;

// 为当前节点申请两个孩子节点
void expandNode() {
TreeNode *newLeftNode=(TreeNode*)malloc(sizeof(TreeNode));
TreeNode *newRightNode=(TreeNode*)malloc(sizeof(TreeNode));
newLeftNode->left=newLeftNode->right=NULL;
newRightNode->left=newRightNode->right=NULL;

this->left=newLeftNode;
this->right=newRightNode;
}

// 遍历树节点 对该节点下的所有孩子节点设置limit值
int walk()
{
if(this->left!=NULL) //Traverse all branches
this->left->walk();
if(this->right!=NULL)
this->right->walk();
this->limit=0;
if(this->left==NULL && this->right==NULL) //Start to calculate node info
this->limit=1;
else
this->limit=this->left->limit + this->right->limit;

return this->limit;
}
};

先保存所有的基本块避免递归处理,并且如果块末尾为InvokeInst跳转则无法对该函数进行平坦化

1
2
3
4
5
6
7
8
9
10
11
bool invoke_found = false;
vector<BasicBlock*> origBB_list;
for(BasicBlock &BB : F){
origBB_list.push_back(&BB);
if (isa<InvokeInst>(BB.getTerminator())) {
errs() << "[-]Find InvokeInst At BasicBlock Terminator!\n";
invoke_found = true;
}
}
if (invoke_found)
continue;

同样的对第一个基本块处理,如果末尾是条件跳转就切开,然后要获得序言块之后的第一个块,后面为switchVar设置初值要用

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
// 获取第一个基本块   并将它从列表中删除
BasicBlock *firstBB=&*F.begin();
origBB_list.erase (origBB_list.begin());

// 序言块之后的第一个块
BasicBlock *startBB=firstBB->getTerminator()->getSuccessor(0);

// 判断第一个基本块末尾是否为条件跳转 若为条件跳转则切开将跳转作为一个单独的块
if( BranchInst *br = cast<BranchInst>(firstBB->getTerminator()) ) {
if ( (br->isConditional()) ||
firstBB->getTerminator()->getNumSuccessors() > 2) {

BasicBlock::iterator iter = firstBB->end();
--iter;

if (firstBB->size() > 1)
--iter;

BasicBlock *splited = firstBB->splitBasicBlock(iter, "splited");
startBB = splited;
origBB_list.insert(origBB_list.begin(), splited);
}
}
firstBB->getTerminator()->eraseFromParent(); // 删除第一个基本块的末尾跳转

动态计算要多少个loopEnd块

1
2
3
4
5
6
7
8
9
10
// 统计ret块的数量
unsigned int retBlockNum=0;
for(BasicBlock *BB : origBB_list){
if(BB->getTerminator()->getNumSuccessors()==0)
retBlockNum++;
}

// 动态计算loopEnd块数量
unsigned int loopEndNum=int( sqrt(origBB_list.size()-retBlockNum) );
errs() << "[+]loopEndNum : " << loopEndNum << '\n';

创建loopEntry、loopEnd块列表和switchVar,并调整块间位置

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
// 在第一个块firstBB前插入一个loopEntry块
BasicBlock *loopEntry=BasicBlock::Create(F.getContext(), "loopEntry", &F, firstBB);
// 创建loopEndNum数量的LoopEndBB 并对每个LoopEndBB创建到loopEntry的跳转
std::vector<BasicBlock*> loopEndList;
for(int i=0;i<loopEndNum;i++)
{
BasicBlock *loopEnd=BasicBlock::Create(F.getContext(),"LoopEnd", &F, firstBB);
loopEndList.push_back(loopEnd);
BranchInst::Create(loopEntry, loopEnd);
}

// adjust loopEntry firstBB place
firstBB->moveBefore(loopEntry);
BranchInst::Create(loopEntry, firstBB);

// Create Switch var
AllocaInst* switchVar = new AllocaInst(TYPE_I32, 0, "switchVar", firstBB->getTerminator());

上面部分都和正常平坦化差不多,下面开始建树

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
std::map<BasicBlock*,TreeNode*> bb_map;
std::map<BasicBlock*,unsigned int> nums_map;

// Spawn Random If
{
// Gen Random Tree
TreeNode tree;
genRandomTree(&tree, origBB_list.size());
tree.walk();
if( allocNode(&tree, 0, 0x7fffffff) ) { // 0x7fffffff
std::vector<BasicBlock*>::iterator iter= origBB_list.begin();
BasicBlock *head=createRandomBasicBlock(&tree, loopEntry->getParent(), switchVar, iter, &bb_map);
BranchInst::Create(head, loopEntry);
}
} // Spawn Random If


// 生成一颗不规则的二叉树 叶子节点的数量为基本块的数量
int llvm::FlatteningEnhanced::genRandomTree(TreeNode *node,int node_limit)
{
std::list<TreeNode*> son_list;
son_list.push_back(node);

// 当叶子节点数量达到node_limit时终止生成
while(son_list.size()!=node_limit)
{
// 从所有孩子节点中随机挑选一个孩子节点
std::list<TreeNode*>::iterator tmp = llvm::Utils::randomChooseElementFromList(&son_list);
TreeNode *node=*tmp;

node->expandNode();
son_list.push_back(node->left);
son_list.push_back(node->right);
son_list.erase(tmp);
}
return son_list.size();
}

// 在[l, r]的大区间上,为当前节点的所有子节点设定一段唯一的左闭右开的区间
bool llvm::FlatteningEnhanced::allocNode(TreeNode *node,unsigned int l,unsigned int r)
{
if(r-l<node->limit)
return false;
node->l=l;
node->r=r;
if(node->left==NULL && node->right==NULL)
return true;
unsigned int var;
if(r-l-node->limit==0)
var=0;
else
var=rand()%(r-l-node->limit);
unsigned int mid=l+node->left->limit+var;
if(!allocNode(node->left,l,mid) || !allocNode(node->right,mid,r))
return false;
return true;
}

// 把基本块排列到所有叶子节点上 并创建子分发块
BasicBlock * llvm::FlatteningEnhanced::createRandomBasicBlock(TreeNode *node,Function *f,Value *var,std::vector<BasicBlock*>::iterator &iter,std::map<BasicBlock*,TreeNode*> *bb_info)
{
// 若是叶子节点 则对该节点建立到基本块的映射
if(node->left==NULL && node->right==NULL)
{
BasicBlock *bb=*iter;
bb_info->insert(std::pair<BasicBlock*,TreeNode*>(bb,node));
iter++;
return bb;
}

// 若非叶子节点 建立一个分发块
BasicBlock *left_bb=createRandomBasicBlock(node->left,f,var,iter,bb_info);
BasicBlock *right_bb=createRandomBasicBlock(node->right,f,var,iter,bb_info);
BasicBlock *node_bb=BasicBlock::Create(f->getContext(),"knot",f);
if(node->left->r!=node->right->l) // 左闭右开区间断言
errs()<<"Error!\n";

// 在分发块末尾插入if语句 若switch变量小于 node->left->r 则跳转到node->left 否则跳转到node->right
LoadInst *load=new LoadInst(TYPE_I32, var, "", node_bb);
ICmpInst *condition=new ICmpInst(*node_bb, ICmpInst::ICMP_ULT, load, ConstantInt::get(TYPE_I32, node->left->r));
BranchInst::Create(left_bb,right_bb,(Value*)condition,node_bb);

return node_bb;
}

建完树后,对每个真实块设置一个val val属于[l, r)

1
2
3
4
5
6
7
8
9
10
11
12
13
unsigned int startNum=0;
for(BasicBlock *BB : origBB_list) {
unsigned int l=bb_map[BB]->l,r=bb_map[BB]->r;
unsigned int val=rand()%(r-l)+l; // [l, r)
nums_map[BB]=val;

// 对序言块之后的第一个真实块 将初始值赋到switchVar
if(BB==startBB) {
// 在序言末、loopEntry之前将switchVar设置为startNum
ConstantInt *startVal=cast<ConstantInt>(ConstantInt::get(TYPE_I32, startNum)); //Set the entry value
new StoreInst( startVal, switchVar, firstBB->getTerminator() );
}
}

处理真实块间跳转,将switchVar设置为目标块的val值(这部分和标准平坦化基本一样了

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
// Handle successors
for(std::vector<BasicBlock *>::iterator b = origBB_list.begin(); b !=origBB_list.end(); ++b) {
BasicBlock *block=*b;

// Ret BB
if(block->getTerminator()->getNumSuccessors() == 0)
continue;

// Non-Conditional Jump
if(block->getTerminator()->getNumSuccessors()==1)
{
BasicBlock *succ=block->getTerminator()->getSuccessor(0);
ConstantInt *caseNum=cast<ConstantInt>(ConstantInt::get(TYPE_I32,nums_map[succ]));
block->getTerminator()->eraseFromParent();
new StoreInst(caseNum,switchVar,block);
BranchInst::Create(loopEndList[rand()%loopEndNum], block);
}

// Conditional Jump
if(block->getTerminator()->getNumSuccessors()==2)
{
BasicBlock *succTrue=block->getTerminator()->getSuccessor(0);
BasicBlock *succFalse=block->getTerminator()->getSuccessor(1);
ConstantInt *numTrue=cast<ConstantInt>(ConstantInt::get(TYPE_I32, nums_map[succTrue]));
ConstantInt *numFalse=cast<ConstantInt>(ConstantInt::get(TYPE_I32, nums_map[succFalse]));
BranchInst *oldBr=cast<BranchInst>(block->getTerminator());
SelectInst *select=SelectInst::Create(oldBr->getCondition(),numTrue,numFalse,Twine("choice"),block->getTerminator());
block->getTerminator()->eraseFromParent();
new StoreInst(select,switchVar,block);
BranchInst::Create(loopEndList[rand()%loopEndNum], block);
}
} // Handle successors

最后,对函数内所有PHI节点和外部寄存器降级到变量

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
std::vector<PHINode *> tmpPhi;
std::vector<Instruction *> tmpReg;
BasicBlock *bbEntry = &*F.begin();
do
{
tmpPhi.clear();
tmpReg.clear();
for(Function::iterator i = F.begin();i!=F.end();i++)
{
for( BasicBlock::iterator j=i->begin();j!=i->end();j++)
{
if(isa<PHINode>(j))
{
PHINode *phi=cast<PHINode>(j);
tmpPhi.push_back(phi);
continue;
}
if (!(isa<AllocaInst>(j) && j->getParent()==bbEntry) && (Utils::valueEscapes(&*j) || j->isUsedOutsideOfBlock(&*i)))
{
tmpReg.push_back(&*j);
continue;
}
}
}
for(unsigned int i=0;i<tmpReg.size();i++)
DemoteRegToStack(*tmpReg.at(i), F.begin()->getTerminator());
for(unsigned int i=0;i<tmpPhi.size();i++)
DemotePHIToStack(tmpPhi.at(i), F.begin()->getTerminator());
}
while(tmpReg.size()!= 0 || tmpPhi.size()!= 0);

最后,总结一下,这个算法核心思想就是建二叉树取代标准平坦化的Switch跳转,但是从效果上来说,其实看不出太多区别,因为非叶子块仅仅也就是一个判断[l, r)范围的操作,和标准平坦化switch生成的子分发块没多少区别(就稍微复杂了一点点)。从平坦化的弱点来看,针对平坦化的处理都是适配识别出所有真实块并符号执行再修复块间关系,这种平坦化并没有提升真实块的识别难度(和标准平坦化一样都是loopEnd的前继就是真实块),所以总体来说,感觉这个混淆的性价比较低(比起这么复杂的Pass流程,并没有起到很好的效果)




浙ICP备19044916号-1