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 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99
| #include "llvm/IR/Function.h" #include "llvm/Pass.h" #include "llvm/IR/LegacyPassManager.h" #include "llvm/Transforms/IPO/PassManagerBuilder.h" #include "llvm/Support/raw_ostream.h" #include "llvm/IR/Instructions.h" #include "llvm/Transforms/Utils.h" #include "llvm/Support/CommandLine.h" #include "llvm/Transforms/Utils/Local.h" #include "SplitBasicBlock.h" #include "Utils.h" #include <vector> #include <cstdlib> #include <ctime> using namespace llvm; using std::vector; namespace{ class Flattening :public FunctionPass{ public : static char ID; Flattening():FunctionPass(ID){ srand(time(0)); } void flatten(Function &F); bool runOnFunction(Function &F);
};
} bool Flattening::runOnFunction(Function &F){ INIT_CONTEXT(F); FunctionPass * pass =createSplitBasicBlockPass(); pass->runOnFunction(F); flatten(F); return true; } void Flattening::flatten(Function &F){ vector<BasicBlock*> origBB; for(BasicBlock& BB:F){ origBB.push_back(&BB); } origBB.erase(origBB.begin()); BasicBlock &entryBlock=F.getEntryBlock(); if(BranchInst *br=dyn_cast<BranchInst>(entryBlock.getTerminator())){ if(br->isConditional()){ origBB.insert(origBB.begin(),entryBlock.splitBasicBlock(br)); } } BasicBlock *dispatchBB=BasicBlock::Create(*CONTEXT,"dispatchBB",&F,&entryBlock); BasicBlock *retBB=BasicBlock::Create(*CONTEXT,"retBB",&F,&entryBlock); entryBlock.moveBefore(dispatchBB); entryBlock.getTerminator()->eraseFromParent(); BranchInst::Create(dispatchBB,&entryBlock); BranchInst::Create(dispatchBB,retBB);
int randNumcase=rand(); AllocaInst * swVarPtr =new AllocaInst(TYPE_I32,0,"swVar.ptr",entryBlock.getTerminator()); new StoreInst(CONST_I32(randNumcase),swVarPtr,entryBlock.getTerminator()); LoadInst *swVar=new LoadInst(TYPE_I32,swVarPtr,"swVAR",dispatchBB); BasicBlock *defaultBB = BasicBlock::Create(*CONTEXT,"defaultBB",&F,retBB); BranchInst::Create(retBB,defaultBB); SwitchInst *swInst=SwitchInst::Create(swVar,defaultBB,0,dispatchBB);
for(BasicBlock *BB:origBB){ BB->moveBefore(retBB); swInst->addCase(CONST_I32(randNumcase),BB); randNumcase=rand();
} for(BasicBlock *BB:origBB){ if(BB->getTerminator()->getNumSuccessors()==0){ continue; } else if(BB->getTerminator()->getNumSuccessors()==1){ ConstantInt *numcase = swInst->findCaseDest(BB->getTerminator()->getSuccessor(0)); new StoreInst(numcase,swVarPtr,BB->getTerminator()); BB->getTerminator()->eraseFromParent(); BranchInst::Create(retBB,BB); }else if(BB->getTerminator()->getNumSuccessors()==2){ ConstantInt *numcase1 = swInst->findCaseDest(BB->getTerminator()->getSuccessor(0)); ConstantInt *numcase2 = swInst->findCaseDest(BB->getTerminator()->getSuccessor(1)); BranchInst *br1=cast<BranchInst>(BB->getTerminator()); SelectInst *sel=SelectInst::Create(br1->getCondition(),numcase1,numcase2,"",BB->getTerminator()); new StoreInst(sel,swVarPtr,BB->getTerminator()); BB->getTerminator()->eraseFromParent(); BranchInst::Create(retBB,BB); } } fixStack(F); }
char Flattening::ID=0; static RegisterPass<Flattening>X("fla","My control flow flattening obfuscation");
|