Skip to content

Commit

Permalink
Refine Flattening implementation
Browse files Browse the repository at this point in the history
  • Loading branch information
@bluesadi
bluesadi committed Jan 1, 2024
1 parent 6c9060e commit dfd1979
Showing 1 changed file with 29 additions and 48 deletions.
77 changes: 29 additions & 48 deletions llvm/lib/Transforms/Obfuscation/Flattening.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,13 +12,13 @@ using std::find;
using std::queue;
using std::vector;

void fixVariables(Function &F, vector<BasicBlock *> &normalBlocks) {
void fixVariables(Function &F) {
vector<PHINode *> origPHI;
vector<Instruction *> origReg;
BasicBlock &entryBB = F.getEntryBlock();

for (BasicBlock *BB : normalBlocks) {
for (Instruction &I : *BB) {
for (BasicBlock &BB : F) {
for (Instruction &I : BB) {
if (PHINode *PN = dyn_cast<PHINode>(&I)) {
origPHI.push_back(PN);
}
Expand All @@ -28,9 +28,9 @@ void fixVariables(Function &F, vector<BasicBlock *> &normalBlocks) {
DemotePHIToStack(PN, entryBB.getTerminator());
}

for (BasicBlock *BB : normalBlocks) {
for (Instruction &I : *BB) {
if (!(isa<AllocaInst>(&I) && I.getParent() == &entryBB) && I.isUsedOutsideOfBlock(BB)) {
for (BasicBlock &BB : F) {
for (Instruction &I : BB) {
if (!(isa<AllocaInst>(&I) && I.getParent() == &entryBB) && I.isUsedOutsideOfBlock(&BB)) {
origReg.push_back(&I);
}
}
Expand All @@ -44,43 +44,29 @@ PreservedAnalyses Pluto::Flattening::run(Function &F, FunctionAnalysisManager &A
auto &context = F.getContext();
IRBuilder<> builder(context);

// 如果只有一个基本块则无需平坦化
// No need to do flattening if only there is only one block
if (F.size() <= 1) {
return PreservedAnalyses::all();
}

vector<BasicBlock *> normalBlocks;
queue<BasicBlock *> queue;
queue.push(&F.getEntryBlock());
while (queue.size()) {
BasicBlock *BB = queue.front();
queue.pop();
if (find(normalBlocks.begin(), normalBlocks.end(), BB) != normalBlocks.end()) {
continue;
}
normalBlocks.push_back(BB);
if (InvokeInst *invoke = dyn_cast_or_null<InvokeInst>(BB->getTerminator())) {
queue.push(invoke->getNormalDest());
} else {
for (size_t i = 0; i < BB->getTerminator()->getNumSuccessors(); i++) {
BasicBlock *successor = BB->getTerminator()->getSuccessor(i);
queue.push(successor);
}
}
vector<BasicBlock *> origBB;

for (BasicBlock &BB : F) {
origBB.push_back(&BB);
}

BasicBlock &entryBB = F.getEntryBlock();
normalBlocks.erase(normalBlocks.begin());
origBB.erase(origBB.begin());

// 如果入口块的末尾是条件跳转，单独分离
// If the entry block ends with a conditional branch, seperate it as a new block
if (entryBB.getTerminator()->getNumSuccessors() > 1) {
BasicBlock *newBB = entryBB.splitBasicBlock(entryBB.getTerminator(), "newBB");
normalBlocks.insert(normalBlocks.begin(), newBB);
origBB.insert(origBB.begin(), newBB);
}

for (BasicBlock *BB : normalBlocks) {
// This is for register demotion
// The return value of the invoke instruction will be store into stack in the bridge block
for (BasicBlock *BB : origBB) {
if (InvokeInst *invoke = dyn_cast_or_null<InvokeInst>(BB->getTerminator())) {
BasicBlock *bridgeBB = BasicBlock::Create(context, "bridgeBB", &F, invoke->getNormalDest());
invoke->getNormalDest()->replacePhiUsesWith(BB, bridgeBB);
Expand All @@ -90,44 +76,44 @@ PreservedAnalyses Pluto::Flattening::run(Function &F, FunctionAnalysisManager &A
}
}

// 创建分发块和返回块
// Create the dispatch block and return block
BasicBlock *dispatchBB = BasicBlock::Create(context, "dispatchBB", &F, &entryBB);
BasicBlock *returnBB = BasicBlock::Create(context, "returnBB", &F, &entryBB);
BranchInst::Create(dispatchBB, returnBB);
entryBB.moveBefore(dispatchBB);

// 去除入口块末尾的跳转，并使入口块跳转到dispatchBB
// Make the entry block go to the dispatchBB directly
entryBB.getTerminator()->eraseFromParent();
BranchInst *brDispatchBB = BranchInst::Create(dispatchBB, &entryBB);

builder.SetInsertPoint(brDispatchBB);

// 在入口块插入alloca和store指令创建并初始化switch变量，初始值为随机值
// Now insert an alloca and a store instruction at the end of entry block, and initialize the
// switch variable with a random value.
// Now insert an alloca and a store instruction at the end of entry block, and initialize the switch variable with
// a random value.
uint32_t randNum = cryptoutils->get_uint32_t();
AllocaInst *swVarPtr = builder.CreateAlloca(Type::getInt32Ty(context), 0, "swVar.ptr");
builder.CreateStore(ConstantInt::get(Type::getInt32Ty(context), randNum), swVarPtr);

// 在分发块插入load指令读取switch变量
// Insert a load instruction at the end of dispatch block
builder.SetInsertPoint(dispatchBB);
LoadInst *swVar = builder.CreateLoad(Type::getInt32Ty(context), swVarPtr, false, "swVar");

// 在分发块插入switch指令实现基本块的调度
// Insert a switch instruction to implement block dispatching
// Insert a switch instruction to dispatch blocks
BasicBlock *swDefault = BasicBlock::Create(context, "swDefault", &F, returnBB);
builder.SetInsertPoint(swDefault);
builder.CreateBr(returnBB);
builder.SetInsertPoint(dispatchBB);
SwitchInst *swInst = builder.CreateSwitch(swVar, swDefault, normalBlocks.size());
SwitchInst *swInst = builder.CreateSwitch(swVar, swDefault, origBB.size());

// Record used random numbers to avoid depulicate case values in switch
std::set<uint32_t> usedNum;

// 将原基本块插入到返回块之前，并分配case值
// Insert original basic blocks before return block and assign a random case value for each one.
for (BasicBlock *BB : normalBlocks) {
for (BasicBlock *BB : origBB) {
// Do not add error handling blocks into switch
if (BB->isEHPad()) {
continue;
}
usedNum.insert(randNum);
BB->moveBefore(returnBB);
swInst->addCase(ConstantInt::get(Type::getInt32Ty(context), randNum), BB);
Expand All @@ -136,16 +122,12 @@ PreservedAnalyses Pluto::Flattening::run(Function &F, FunctionAnalysisManager &A
} while (find(usedNum.begin(), usedNum.end(), randNum) != usedNum.end());
}

// 在每个基本块最后添加修改switch变量的指令和跳转到返回块的指令
// Insert a store instruction at the end of each block to modify swVar, and make them jump back
// to dispatch block then.
for (BasicBlock *BB : normalBlocks) {
// 跳过没有后继的基本块
// Insert a store instruction at the end of each block to modify swVar, and make them jump back to dispatch block.
for (BasicBlock *BB : origBB) {
// Skip blocks with no successor
if (BB->getTerminator()->getNumSuccessors() == 0) {
continue;
}
// 非条件跳转
// Branch
else if (BB->getTerminator()->getNumSuccessors() == 1) {
BasicBlock *sucBB = BB->getTerminator()->getSuccessor(0);
Expand All @@ -159,7 +141,6 @@ PreservedAnalyses Pluto::Flattening::run(Function &F, FunctionAnalysisManager &A
}
}
}
// 条件跳转
// Conditional branch
else if (BB->getTerminator()->getNumSuccessors() == 2) {
ConstantInt *numIfTrue = swInst->findCaseDest(BB->getTerminator()->getSuccessor(0));
Expand All @@ -175,6 +156,6 @@ PreservedAnalyses Pluto::Flattening::run(Function &F, FunctionAnalysisManager &A
}
}
}
fixVariables(F, normalBlocks);
fixVariables(F);
return PreservedAnalyses::none();
}

0 comments on commit dfd1979

Please sign in to comment.