using namespace llvm;
using std::vector;
namespace{
struct Obfu : public FunctionPass{
static char
ID
;
Obfu() : FunctionPass(
ID
){}
bool
flatten(Function
*
f);
bool
runOnFunction(Function &F);
};
}
bool
Obfu::runOnFunction(Function &F){
return
flatten(&F);
}
bool
Obfu::flatten(Function
*
f){
IntegerType
*
int32Type
=
Type
::getInt32Ty(f
-
>getContext());
/
/
遍历函数所有基本块,将其存到vector中
vector<BasicBlock
*
> origBB;
for
(BasicBlock &BB:
*
f){
origBB.push_back(&BB);
}
/
/
基本块数量不超过
1
则无需平坦化
if
(origBB.size() <
=
1
){
return
false;
}
/
/
从vector中去除第一个基本块
origBB.erase(origBB.begin());
BasicBlock
*
firstBB
=
&f
-
>front();
/
/
如果第一个基本块的末尾是条件跳转
if
(isa<BranchInst>(firstBB
-
>getTerminator())){
BranchInst
*
br
=
cast<BranchInst>(firstBB
-
>getTerminator());
if
(br
-
>isConditional()){
CmpInst
*
cmpInst
=
cast<CmpInst>(firstBB
-
>getTerminator()
-
>getPrevNode());
BasicBlock
*
newBB
=
firstBB
-
>splitBasicBlock(cmpInst,
"newBB"
);
origBB.insert(origBB.begin(), newBB);
}
}
/
/
创建循环
BasicBlock
*
loopEntry
=
BasicBlock::Create(f
-
>getContext(),
"loopEntry"
, f, firstBB);
BasicBlock
*
loopEnd
=
BasicBlock::Create(f
-
>getContext(),
"loopEnd"
, f, firstBB);
firstBB
-
>moveBefore(loopEntry);
/
/
去除第一个基本块末尾的跳转
firstBB
-
>getTerminator()
-
>eraseFromParent();
/
/
初始化switch on变量
srand(time(
0
));
int
randNumCase
=
rand();
AllocaInst
*
swVarPtr
=
new AllocaInst(int32Type,
0
,
"swVar.ptr"
, firstBB);
new StoreInst(ConstantInt::get(int32Type, randNumCase), swVarPtr, firstBB);
/
/
使第一个基本块跳转到loopEntry
BranchInst::Create(loopEntry, firstBB);
/
/
在进入loopEntry读取switch on变量
LoadInst
*
swVar
=
new LoadInst(int32Type, swVarPtr,
"swVar"
, false, loopEntry);
BranchInst::Create(loopEntry, loopEnd);
/
/
初始化switch的default case
/
/
default case实际上不会被执行
BasicBlock
*
swDefault
=
BasicBlock::Create(f
-
>getContext(),
"swDefault"
, f, loopEnd);
BranchInst::Create(loopEnd, swDefault);
SwitchInst
*
swInst
=
SwitchInst::Create(swVar, swDefault,
0
, loopEntry);
/
/
插入原基本块到switch中,仅是位置意义上的插入,而不是逻辑意义上的
for
(BasicBlock
*
BB : origBB){
ConstantInt
*
numCase
=
cast<ConstantInt>(ConstantInt::get(int32Type, randNumCase));
BB
-
>moveBefore(loopEnd);
swInst
-
>addCase(numCase,BB);
randNumCase
=
rand();
}
/
/
添加case
for
(BasicBlock
*
BB : origBB){
/
/
retn BB
if
(BB
-
>getTerminator()
-
>getNumSuccessors()
=
=
0
){
continue
;
}
/
/
非条件跳转
if
(BB
-
>getTerminator()
-
>getNumSuccessors()
=
=
1
){
BasicBlock
*
sucBB
=
BB
-
>getTerminator()
-
>getSuccessor(
0
);
BB
-
>getTerminator()
-
>eraseFromParent();
ConstantInt
*
numCase
=
swInst
-
>findCaseDest(sucBB);
new StoreInst(numCase, swVarPtr, BB);
BranchInst::Create(loopEnd, BB);
continue
;
}
/
/
条件跳转
if
(BB
-
>getTerminator()
-
>getNumSuccessors()
=
=
2
){
ConstantInt
*
numCaseTrue
=
swInst
-
>findCaseDest(BB
-
>getTerminator()
-
>getSuccessor(
0
));
ConstantInt
*
numCaseFalse
=
swInst
-
>findCaseDest(BB
-
>getTerminator()
-
>getSuccessor(
1
));
BranchInst
*
br
=
cast<BranchInst>(BB
-
>getTerminator());
SelectInst
*
sel
=
SelectInst::Create(br
-
>getCondition(), numCaseTrue, numCaseFalse, "", BB
-
>getTerminator());
BB
-
>getTerminator()
-
>eraseFromParent();
new StoreInst(sel, swVarPtr, BB);
BranchInst::Create(loopEnd, BB);
}
}
return
true;
}
char Obfu::
ID
=
0
;
static RegisterPass<Obfu> X(
"obfu"
,
"My obfuscating pass"
);