From 6032fb15b49d28c2147bea997c47fab22d7f6cb8 Mon Sep 17 00:00:00 2001
From: "yangzhao.zy" <yangzhao.zy@alibaba-inc.com>
Date: Thu, 24 Dec 2015 10:09:48 +0800
Subject: [PATCH] Record each function's branch destination, add it into switch
 block.

---
 lib/Transforms/Obfuscation/Flattening.cpp | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/lib/Transforms/Obfuscation/Flattening.cpp b/lib/Transforms/Obfuscation/Flattening.cpp
index 2158a47449cd..0c5ab209e766 100644
--- a/lib/Transforms/Obfuscation/Flattening.cpp
+++ b/lib/Transforms/Obfuscation/Flattening.cpp
@@ -107,6 +107,7 @@ bool Flattening::flatten(Function *f) {
     br = cast<BranchInst>(insert->getTerminator());
   }
 
+  BasicBlock *zeroBB = NULL;
   if ((br != NULL && br->isConditional()) ||
       insert->getTerminator()->getNumSuccessors() > 1) {
     BasicBlock::iterator i = insert->back();
@@ -117,6 +118,9 @@ bool Flattening::flatten(Function *f) {
 
     BasicBlock *tmpBB = insert->splitBasicBlock(i, "first");
     origBB.insert(origBB.begin(), tmpBB);
+  } else if (br != NULL && insert->getTerminator()->getNumSuccessors() == 1) {
+    // Record the first BB branch to
+    zeroBB = insert->getTerminator()->getSuccessor(0);
   }
 
   // Remove jump
@@ -169,7 +173,12 @@ bool Flattening::flatten(Function *f) {
     numCase = cast<ConstantInt>(ConstantInt::get(
         switchI->getCondition()->getType(),
         llvm::cryptoutils->scramble32(switchI->getNumCases(), scrambling_key)));
-    switchI->addCase(numCase, i);
+    if (zeroBB && zeroBB != i) {
+      switchI->addCase(numCase, zeroBB);
+      zeroBB = NULL;
+    } else {
+      switchI->addCase(numCase, i);
+    }
   }
 
   // Recalculate switchVar