Fix two GlobalsEncryption pass bugs

llvm/lib/Transforms/Obfuscation/Flattening.cpp

@@ -23,6 +23,17 @@ bool Flattening::runOnFunction(Function &F) {
     return false;
 }
 
+bool has_unsupported_ir(Function &F){
+    for (BasicBlock &BB : F) {
+        for (Instruction &I : BB) {
+            if (isa<LandingPadInst>(&I) || isa<PHINode>(&I))
+                return true;
+        }
+    }
+    return false;
+}
+
+
 void Flattening::flatten(Function &F) {
     // 基本块数量不超过1则无需平坦化
     if (F.size() <= 1) {
@@ -34,9 +45,27 @@ void Flattening::flatten(Function &F) {
     // 将除入口块（第一个基本块）以外的基本块保存到一个 vector
     // 容器中，便于后续处理 首先保存所有基本块
     vector<BasicBlock *> origBB;
+
+    int use_flat = 1;
+
     for (BasicBlock &BB : F) {
         origBB.push_back(&BB);
+
+        unsigned int op = BB.getTerminator()->getOpcode();
+        if(op != Instruction::Br && op != Instruction::Ret){
+            use_flat = 0;
+            break;
+        }
+    }
+
+    if(has_unsupported_ir(F)){
+        use_flat = 0;
     }
+
+    if(!use_flat){
+        return;
+    }
+
     // 从vector中去除第一个基本块
     origBB.erase(origBB.begin());
     BasicBlock &entryBB = F.getEntryBlock();

llvm/lib/Transforms/Obfuscation/GlobalsEncryption.cpp

@@ -28,19 +28,25 @@ bool GlobalsEncryption::runOnModule(Module &M) {
     INIT_CONTEXT(M);
     vector<GlobalVariable *> GVs;
     for (GlobalVariable &GV : M.getGlobalList()) {
+        // only process non llvm-generated IRs
+        if(GV.getName().contains("llvm"))
+            continue;
         GVs.push_back(&GV);
     }
     for (int i = 0; i < ObfuTimes; i++) {
         for (GlobalVariable *GV : GVs) {
-        // Only encrypt globals of integer and array
-        if (!GV->getValueType()->isIntegerTy() &&
-            !GV->getValueType()->isArrayTy()) {
-            continue;
-        }
+            if(GV->getValueType()->isArrayTy()){ // the value can be array
+                ArrayType *ArrTy = dyn_cast<ArrayType>(GV->getValueType());
+                if(!ArrTy->getElementType()->isIntegerTy()){  // but the array must be integerty
+                    continue;
+                }
+            }
+            else if (!GV->getValueType()->isIntegerTy()){  // or, the value must be integerty
+                continue;
+            }
+
         if (GV->hasInitializer() && GV->getInitializer() &&
-            (GV->getName().contains(".str") || !OnlyStr)
-            // Do not encrypt globals having a section named "llvm.metadata"
-            && !GV->getSection().equals("llvm.metadata")) {
+            (GV->getName().contains(".str") || !OnlyStr)) {
             Constant *initializer = GV->getInitializer();
             ConstantInt *intData = dyn_cast<ConstantInt>(initializer);
             ConstantDataArray *arrData = dyn_cast<ConstantDataArray>(initializer);
@@ -61,6 +67,7 @@ bool GlobalsEncryption::runOnModule(Module &M) {
             GV->setConstant(false);
             insertArrayDecryption(M, {GV, key, eleNum});
             } else if (intData) {
+                continue;  // TODO: fix the bug
             uint64_t key = cryptoutils->get_uint64_t();
             ConstantInt *enc =
                 CONST(intData->getType(), key ^ intData->getZExtValue());