ollvm源码分析之控制流扁平化（3）

参考网站：ollvm源码分析
参考网站：Obufuscator-llvm源码分析

Pass之flattening

实现功能是：增加switch-case语句，使其扁平化。

1 命令

命令	解析
-mllvm -fla	激活控制流扁平化

2 代码分析

2.1 入口函数runOnFunction

Flattening继承了FunctionPass，因此它的入口函数即为runOnFunction。

bool Flattening::runOnFunction(Function &F) {Function *tmp = &F;// Do we obfuscateif (toObfuscate(flag, tmp, "fla")) {if (flatten(tmp)) {++Flattened;}}return false;
}

toObfuscate函数判断是否启动了fla命令，判断完毕，再调用flatten函数。

2.1 flatten函数

vector<BasicBlock *> origBB;
for (Function::iterator i = f->begin(); i != f->end(); ++i) {BasicBlock *tmp = &*i;origBB.push_back(tmp);BasicBlock *bb = &*i;if (isa<InvokeInst>(bb->getTerminator())) {return false;}}
// Nothing to flatten
if (origBB.size() <= 1) {return false;
}// Remove first BB
origBB.erase(origBB.begin());

除了first BB，将所有基本块存储到vector容器origBB中，若origBB中不止一个基本块才进行后续处理。

// Get a pointer on the first BB
Function::iterator tmp = f->begin();  //++tmp;
BasicBlock *insert = &*tmp;
// If main begin with an if
BranchInst *br = NULL;
if (isa<BranchInst>(insert->getTerminator())) {br = cast<BranchInst>(insert->getTerminator());
}if ((br != NULL && br->isConditional()) ||insert->getTerminator()->getNumSuccessors() > 1) {BasicBlock::iterator i = insert->end();--i;if (insert->size() > 1) {--i;}BasicBlock *tmpBB = insert->splitBasicBlock(i, "first");origBB.insert(origBB.begin(), tmpBB);
}
// Remove jump
insert->getTerminator()->eraseFromParent();

若first BB未尾有条件跳转指令，获取该指令地址。通过splitBasicBlock将该块分为first BB与tmpBB，将tmpBB存储到origBB中。去除第一个基本块与下一个基本块间的跳转关系。

switchVar = new AllocaInst(Type::getInt32Ty(f->getContext()), 0,"switchVar", insert);
new StoreInst(ConstantInt::get(Type::getInt32Ty(f->getContext()),llvm::cryptoutils->scramble32(0, scrambling_key)),switchVar, insert);

创建switch变量并且初始化。

// Create main looploopEntry = BasicBlock::Create(f->getContext(), "loopEntry", f, insert);loopEnd = BasicBlock::Create(f->getContext(), "loopEnd", f, insert);load = new LoadInst(switchVar, "switchVar", loopEntry);// Move first BB on topinsert->moveBefore(loopEntry);BranchInst::Create(loopEntry, insert);// loopEnd jump to loopEntryBranchInst::Create(loopEntry, loopEnd);BasicBlock *swDefault =BasicBlock::Create(f->getContext(), "switchDefault", f, loopEnd);BranchInst::Create(loopEnd, swDefault);// Create switch instruction itself and set conditionswitchI = SwitchInst::Create(&*f->begin(), swDefault, 0, loopEntry);switchI->setCondition(load);

创建两个基本块loopEntry 、loopEnd。将first BB移至loopEntry 前面，创建loopEnd到时loopEntry 的跳转。创建swDefault 块，swDefault跳转到loopEnd，loopEntry 跳转到swDefault。

 f->begin()->getTerminator()->eraseFromParent();BranchInst::Create(loopEntry, &*f->begin());

删除first BB的跳转，将first BB跳转至loopEntry。

 // Put all BB in the switchfor (vector<BasicBlock *>::iterator b = origBB.begin(); b != origBB.end();++b) {BasicBlock *i = *b;ConstantInt *numCase = NULL;// Move the BB inside the switch (only visual, no code logic)i->moveBefore(loopEnd);// Add case to switchnumCase = cast<ConstantInt>(ConstantInt::get(switchI->getCondition()->getType(),llvm::cryptoutils->scramble32(switchI->getNumCases(), scrambling_key)));switchI->addCase(numCase, i);}

所有基本块加入switch中。

// Recalculate switchVarfor (vector<BasicBlock *>::iterator b = origBB.begin(); b != origBB.end();++b) {BasicBlock *i = *b;ConstantInt *numCase = NULL;// Ret BBif (i->getTerminator()->getNumSuccessors() == 0) {continue;}// If it's a non-conditional jumpif (i->getTerminator()->getNumSuccessors() == 1) {// Get successor and delete terminatorBasicBlock *succ = i->getTerminator()->getSuccessor(0);i->getTerminator()->eraseFromParent();// Get next casenumCase = switchI->findCaseDest(succ);// If next case == default case (switchDefault)if (numCase == NULL) {numCase = cast<ConstantInt>(ConstantInt::get(switchI->getCondition()->getType(),llvm::cryptoutils->scramble32(switchI->getNumCases() - 1, scrambling_key)));}// Update switchVar and jump to the end of loopnew StoreInst(numCase, load->getPointerOperand(), i);BranchInst::Create(loopEnd, i);continue;}// If it's a conditional jumpif (i->getTerminator()->getNumSuccessors() == 2) {// Get next casesConstantInt *numCaseTrue =switchI->findCaseDest(i->getTerminator()->getSuccessor(0));ConstantInt *numCaseFalse =switchI->findCaseDest(i->getTerminator()->getSuccessor(1));// Check if next case == default case (switchDefault)if (numCaseTrue == NULL) {numCaseTrue = cast<ConstantInt>(ConstantInt::get(switchI->getCondition()->getType(),llvm::cryptoutils->scramble32(switchI->getNumCases() - 1, scrambling_key)));}if (numCaseFalse == NULL) {numCaseFalse = cast<ConstantInt>(ConstantInt::get(switchI->getCondition()->getType(),llvm::cryptoutils->scramble32(switchI->getNumCases() - 1, scrambling_key)));}// Create a SelectInstBranchInst *br = cast<BranchInst>(i->getTerminator());SelectInst *sel =SelectInst::Create(br->getCondition(), numCaseTrue, numCaseFalse, "",i->getTerminator());// Erase terminatori->getTerminator()->eraseFromParent();// Update switchVar and jump to the end of loopnew StoreInst(sel, load->getPointerOperand(), i);BranchInst::Create(loopEnd, i);continue;}}

修改每个基本块之间的跳转关系，使得每个基本块执行完毕会重新设置switchVar 的值。从而回到switch并顺利跳转到下一个case，直到程序结束。
处理后的基本块间的关系图如下图所示。