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");
   |